convert_fc_to_tfrecords.py 文件源码

def convert_dataset(indices, name):
    # Open a TFRRecordWriter
    filename = os.path.join(FLAGS.out, name + '.tfrecords')
    writeOpts = tf.python_io.TFRecordOptions(tf.python_io.TFRecordCompressionType.ZLIB)
    writer = tf.python_io.TFRecordWriter(filename, options=writeOpts)

    # Load each data sample (image_a, image_b, flow) and write it to the TFRecord
    count = 0
    pbar = ProgressBar(widgets=[Percentage(), Bar()], maxval=len(indices)).start()
    for i in indices:
        image_a_path = os.path.join(FLAGS.data_dir, '%05d_img1.ppm' % (i + 1))
        image_b_path = os.path.join(FLAGS.data_dir, '%05d_img2.ppm' % (i + 1))
        flow_path = os.path.join(FLAGS.data_dir, '%05d_flow.flo' % (i + 1))

        image_a = imread(image_a_path)
        image_b = imread(image_b_path)

        # Convert from RGB -> BGR
        image_a = image_a[..., [2, 1, 0]]
        image_b = image_b[..., [2, 1, 0]]

        # Scale from [0, 255] -> [0.0, 1.0]
        image_a = image_a / 255.0
        image_b = image_b / 255.0

        image_a_raw = image_a.tostring()
        image_b_raw = image_b.tostring()
        flow_raw = open_flo_file(flow_path).tostring()

        example = tf.train.Example(features=tf.train.Features(feature={
            'image_a': _bytes_feature(image_a_raw),
            'image_b': _bytes_feature(image_b_raw),
            'flow': _bytes_feature(flow_raw)}))
        writer.write(example.SerializeToString())
        pbar.update(count + 1)
        count += 1
    writer.close()