def build_dataset(name, file_list):
writer = tf.python_io.TFRecordWriter(
'nyuv2.{}.tfrecords'.format(name))
image_dir = os.path.join(dataset_dir)
lines = open(os.path.join(image_dir, file_list)).readlines()
for line in tqdm(lines):
fin = os.path.join(image_dir, line.strip() + '_i.png')
fdn = os.path.join(image_dir, line.strip() + '_f.png')
image_data = tf.gfile.FastGFile(fin, 'rb').read()
depth_data = tf.gfile.FastGFile(fdn, 'rb').read()
w, h, c = 427, 561, 3
# write record
example = tf.train.Example(features=tf.train.Features(
feature={
'image':
hem.bytes_feature(image_data),
'depth':
hem.bytes_feature(depth_data),
'width':
hem.int64_feature(w),
'height':
hem.int64_feature(h),
'channels':
hem.int64_feature(c),
'filename':
hem.bytes_feature(tf.compat.as_bytes(fin)),
'depth_filename':
hem.bytes_feature(tf.compat.as_bytes(fdn))
}))
writer.write(example.SerializeToString())
def build_dataset(name, file_list):
writer = tf.python_io.TFRecordWriter(os.path.join(storage_dir, 'floorplans.{}.tfrecords'.format(name)))
image_dir = os.path.join(dataset_dir) # '/mnt/research/datasets/floorplans')
lines = open(os.path.join(image_dir, file_list)).readlines()
for line in tqdm(lines):
fn = os.path.join(image_dir, line.strip())
# read image
with tf.gfile.FastGFile(fn, 'rb') as f:
image_data = f.read()
# determine shape
f = open(fn, 'rb')
stuff = f.read()
f.close()
nparr = np.fromstring(stuff, np.uint8)
i = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
w, h, c = i.shape
# write record
example = tf.train.Example(features=tf.train.Features(
feature={
'image': hem.bytes_feature(image_data),
'width': hem.int64_feature(w),
'height': hem.int64_feature(h),
'channels': hem.int64_feature(c),
'filename': hem.bytes_feature(tf.compat.as_bytes(fn))}))
writer.write(example.SerializeToString())
def write_tfrecord(images, labels, fn):
writer = tf.python_io.TFRecordWriter(os.path.join(storage_dir, fn))
for i in range(images.shape[0]):
img = images[i].tostring()
label = int(labels[i])
record = tf.train.Example(features=tf.train.Features(
feature={
'image': hem.bytes_feature(img),
'label': hem.int64_feature(label)
}))
writer.write(record.SerializeToString())
writer.close()
def build_dataset(name):
image_dirs = {
'train': 'train2014',
'validate': 'val2014',
'test': 'test2014'
}
annotate_files = {
'train': 'instances_train2014.json',
'validate': 'instances_val2014.json',
'test': 'image_info_test2014.json'
}
image_dir = os.path.join(download_dir, image_dirs[name])
annotate_file = os.path.join(download_dir, 'annotations',
annotate_files[name])
coco = COCO(annotate_file)
categories = coco.loadCats(coco.getCatIds())
category_names = set([c['name'] for c in categories])
category_ids = coco.getCatIds(catNms=category_names)
image_ids = coco.getImgIds(catIds=category_ids)
file_list = []
for i in image_ids:
# keys = license, file_name, coco_url, height, width, date_captured, flickr_url, id
img = coco.loadImgs(i)[0]
file_list.append(img)
writer = tf.python_io.TFRecordWriter(
os.path.join(storage_dir, 'coco.{}.tfrecords'.format(name)))
for img in tqdm(file_list):
path = os.path.join(image_dir, img['file_name'])
with tf.gfile.FastGFile(path, 'rb') as f:
image_data = f.read()
# create mask
annotations = coco.loadAnns(coco.getAnnIds(imgIds=img['id']))
total_mask = np.zeros((img['height'], img['width'], 1),
dtype=np.uint8)
label_list = []
bbox_list = []
crowd_list = []
area_list = []
for a in annotations:
mask = coco.annToMask(a)
total_mask[mask == 1] = int(a['category_id'])
for b in a['bbox']:
bbox_list.append(b)
# bbox_list.append(a['bbox'])
crowd_list.append(a['iscrowd'])
area_list.append(a['area'])
label_list.append(a['category_id'])
example = tf.train.Example(features=tf.train.Features(
feature={
'image':
hem.bytes_feature(image_data),
'annotations':
hem.bytes_feature(total_mask.tostring()),
'filename':
hem.bytes_feature(tf.compat.as_bytes(
img['file_name'])),
'width':
hem.int64_feature(img['width']),
'height':
hem.int64_feature(img['height']),
'image_id':
hem.int64_feature(img['id']),
'bboxes':
hem.float_feature(bbox_list),
'iscrowds':
hem.int64_feature(crowd_list),
'areas':
hem.float_feature(area_list),
'labels':
hem.int64_feature(label_list)
}))
writer.write(example.SerializeToString())