Showing
1 changed file
with
53 additions
and
0 deletions
code/tfrecord_writer.py
0 → 100644
| 1 | +import tensorflow as tf | ||
| 2 | + | ||
| 3 | +def bytes_feature(values): | ||
| 4 | + return tf.train.Feature(bytes_list=tf.train.BytesList(value=[values])) | ||
| 5 | + | ||
| 6 | +def int64_feature(values): | ||
| 7 | + if not isinstance(values, (tuple, list)): | ||
| 8 | + values = [values] | ||
| 9 | + | ||
| 10 | + return tf.train.Feature(int64_list=tf.train.Int64List(value=values)) | ||
| 11 | + | ||
| 12 | +def readImage(path): | ||
| 13 | + file = open(path, 'rb') | ||
| 14 | + byte = file.read() | ||
| 15 | + return byte | ||
| 16 | + | ||
| 17 | +def main(): | ||
| 18 | + ANNOTATION_PATH = 'train.txt' #annotation set (train/val/test) text file | ||
| 19 | + IMAGE_DIRECTORY = 'image_data/' #image directory | ||
| 20 | + SAVE_PATH = 'train.tfrecord' #save path for tfrecord | ||
| 21 | + | ||
| 22 | + print('Tensorflow version:', tf.__version__) #tensorflow version should be 1.x | ||
| 23 | + | ||
| 24 | + file = open(ANNOTATION_PATH, 'r') | ||
| 25 | + lines = file.readlines() | ||
| 26 | + file.close() | ||
| 27 | + | ||
| 28 | + options = tf.python_io.TFRecordOptions(tf.python_io.TFRecordCompressionType.GZIP) #compress option | ||
| 29 | + writer = tf.python_io.TFRecordWriter(path=SAVE_PATH, options=options) | ||
| 30 | + | ||
| 31 | + for line in lines: | ||
| 32 | + parsed = line.split(' ') | ||
| 33 | + print('Current Doing...', parsed[1]) #debug messages | ||
| 34 | + image = readImage(IMAGE_DIRECTORY + '/' + parsed[1]) | ||
| 35 | + boxes = [] | ||
| 36 | + | ||
| 37 | + for i in range(4, len(parsed)): | ||
| 38 | + boxes.append(int(parsed[i])) | ||
| 39 | + | ||
| 40 | + data = tf.train.Example(features=tf.train.Features(feature={ | ||
| 41 | + 'index': int64_feature(int(parsed[0])), | ||
| 42 | + 'image': bytes_feature(image), | ||
| 43 | + 'width': int64_feature(int(parsed[2])), | ||
| 44 | + 'height': int64_feature(int(parsed[3])), | ||
| 45 | + 'boxes': int64_feature(boxes) # boxes = [label1, xmin1, ymin1, xmax1, ymax1, label2, xmax2, ...] | ||
| 46 | + })) | ||
| 47 | + | ||
| 48 | + writer.write(data.SerializeToString()) | ||
| 49 | + | ||
| 50 | + writer.close() | ||
| 51 | + | ||
| 52 | +if __name__ == '__main__': | ||
| 53 | + main() | ||
| ... | \ No newline at end of file | ... | \ No newline at end of file |
-
Please register or login to post a comment