-
Notifications
You must be signed in to change notification settings - Fork 8
/
Copy pathconvert_to_tfrecord.py
93 lines (77 loc) · 3.71 KB
/
convert_to_tfrecord.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
# ref: https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/using_your_own_dataset.md
import os
import tensorflow as tf
import data_utils as util
from PIL import Image
from tqdm import tqdm
from object_detection.utils import dataset_util
flags = tf.app.flags
flags.DEFINE_string('output_path', '', 'Path to output TFRecord')
FLAGS = flags.FLAGS
TYPE_OF_DATASET = 'train'
fileformat = 'png' # 'jpg'
if TYPE_OF_DATASET == 'train':
if fileformat == 'jpg':
source_img_dir = util.d_train.img_dir_jpg
else:
source_img_dir = util.d_train.img_dir
source_label_dir = util.d_train.label_dir
elif TYPE_OF_DATASET == 'valid':
if fileformat == 'jpg':
source_img_dir = util.d_valid.img_dir_jpg
else:
source_img_dir = util.d_valid.img_dir
source_label_dir = util.d_valid.label_dir
else:
print('[warning] Type of dataset is not defined')
def create_tf_example(name):
# TODO(user): Populate the following variables from your example.
b_image = util.encode_image_png(os.path.join(source_img_dir, name) + '.' + fileformat)
label_objects = util.parse_dota_poly(os.path.join(source_label_dir, name) + '.txt')
width, height = Image.open(os.path.join(source_img_dir, name) + '.' + fileformat).size # Image width, height
filename = name.encode() # Filename of the image. Empty if image is not from file
encoded_image_data = b_image # Encoded image bytes
if fileformat == 'jpg':
image_format = b'jpeg'
else:
image_format = b'png' # b'jpeg' or b'png'
xmins = [] # List of normalized left x coordinates in bounding box (1 per box)
xmaxs = [] # List of normalized right x coordinates in bounding box (1 per box)
ymins = [] # List of normalized top y coordinates in bounding box (1 per box)
ymaxs = [] # List of normalized bottom y coordinates in bounding box (1 per box)
classes_text = [] # List of string class name of bounding box (1 per box)
classes = [] # List of integer class id of bounding box (1 per box)
for obj in label_objects:
poly = obj['poly']
xmin, xmax, ymin, ymax = util.dots4ToRec4(poly)
xmins.append(xmin)
xmaxs.append(xmax)
ymins.append(ymin)
ymaxs.append(ymax)
classes_text.append(obj['name'].encode())
classes.append(util.dota_15.index(obj['name']))
tf_example = tf.train.Example(features=tf.train.Features(feature={
'image/height': dataset_util.int64_feature(height),
'image/width': dataset_util.int64_feature(width),
'image/filename': dataset_util.bytes_feature(filename),
'image/source_id': dataset_util.bytes_feature(filename),
'image/encoded': dataset_util.bytes_feature(encoded_image_data),
'image/format': dataset_util.bytes_feature(image_format),
'image/object/bbox/xmin': dataset_util.float_list_feature(xmins),
'image/object/bbox/xmax': dataset_util.float_list_feature(xmaxs),
'image/object/bbox/ymin': dataset_util.float_list_feature(ymins),
'image/object/bbox/ymax': dataset_util.float_list_feature(ymaxs),
'image/object/class/text': dataset_util.bytes_list_feature(classes_text),
'image/object/class/label': dataset_util.int64_list_feature(classes),
}))
return tf_example
def main(_):
writer = tf.python_io.TFRecordWriter('./datasets/{}_dataset.record'.format(TYPE_OF_DATASET))
img_paths = util.get_file_from_this_rootdir(source_img_dir, fileformat)
for idx, path in enumerate(tqdm(img_paths)):
name = os.path.splitext(os.path.basename(path))[0]
tf_example = create_tf_example(name)
writer.write(tf_example.SerializeToString())
writer.close()
if __name__ == '__main__':
tf.app.run()