def test_write_read_tfrecord(self):
"""Test it can write and read the tfrecord file correctly.
"""
# Constants.
DATA_CATEGORY = ['train', 'val', 'test']
# Make a dummy tfrecord file.
input_dir, gt_data_list = _create_sample_cityscapes_structure(
self.tmpdir)
output_dir = input_dir
# Convert from py.path.local to str.
data_list = get_file_path(input_dir)
write_tfrecord(data_list, output_dir)
# Read the created tfrecord file.
init_op = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
for category in DATA_CATEGORY:
dataset = read_tfrecord(
os.path.join(output_dir, category + '_0000.tfrecord'))
next_element = dataset.make_one_shot_iterator().get_next()
with self.test_session() as sess:
# The op for initializing the variables.
sess.run(init_op)
i = 0
while True:
try:
sample = sess.run(next_element)
gt_image = np.array(
Image.open(open(sample['filename'].decode(),
'rb')).convert('RGB'))
assert np.array_equal(sample['image'], gt_image)
self.assertEqual(sample['height'], IMAGE_HEIGHT)
self.assertEqual(sample['width'], IMAGE_WIDTH)
i += 1
except tf.errors.OutOfRangeError:
assert i == 4
break
raise AttributeError('--save_dir should be existing directory path.')
save_dir_fullpath = os.path.join(save_dir_fullpath, str(datetime.now()))
os.makedirs(save_dir_fullpath)
logging.info('Created save directory to {}'.format(save_dir_fullpath))
resume_fullpath = None
if options['resume_path']:
resume_fullpath = os.path.expanduser(options['resume_path'])
finetune_from = None
if options['finetune_from']:
finetune_from = options['finetune_from']
# Data part should live in cpu.
with tf.device('/cpu:0'):
# Read from tfrecord format data made by src.data.tfrecord.TFRecordWriter.
train_dataset = read_tfrecord(
os.path.join(options['tfdata_dir'], 'train_*.tfrecord'))
train_data_processor = DataPreprocessor(
dataset=train_dataset,
num_parallel_calls=options['num_parallel_calls'],
batch_size=options['batch_size'],
shuffle_buffer_size=options['shuffle_buffer_size'],
prefetch_buffer_size=options['prefetch_buffer_size'])
val_dataset = read_tfrecord(
os.path.join(options['tfdata_dir'], 'val_*.tfrecord'))
val_data_processor = DataPreprocessor(
dataset=val_dataset,
num_parallel_calls=options['num_parallel_calls'],
batch_size=1,
shuffle_buffer_size=1,
prefetch_buffer_size=1)
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
save_dir_fullpath = os.path.expanduser(options['save_dir'])
if not os.path.isdir(save_dir_fullpath):
raise AttributeError('--save_dir should be existing directory path.')
save_dir_fullpath = os.path.join(save_dir_fullpath, str(datetime.now()))
os.makedirs(save_dir_fullpath)
logging.info('Created save directory to {}'.format(save_dir_fullpath))
ckpt_fullpath = os.path.expanduser(options['ckpt_path'])
# Data part should live in cpu.
with tf.device('/cpu:0'):
# Read from tfrecord format data made by src.data.tfrecord.TFRecordWriter.
test_dataset = read_tfrecord(
os.path.join(options['tfdata_dir'], 'test_*.tfrecord'))
test_data_processor = DataPreprocessor(
dataset=test_dataset,
num_parallel_calls=options['num_parallel_calls'],
batch_size=options['batch_size'],
shuffle_buffer_size=None,
prefetch_buffer_size=1)
# Pre-process test data.
# Add more pre-procesing blocks.
logging.info('Preprocess test data')
if options['test_resized_height'] is not None and options[
'test_resized_width'] is not None:
logging.info('Resize image to ({}, {})'.format(
options['test_resized_height'], options['test_resized_width']))
test_data_processor.process_image_and_label(
import project_root
from src.data.cityscapes import id2label
from src.data.tfrecord import read_tfrecord
if __name__ == '__main__':
parser = argparse.ArgumentParser(
description="The script to compute cityscapes class distribution.")
parser.add_argument(
'data_path', type=str, help='Regex to the tfrecord data paths.')
parser.add_argument('num_classes', type=int, help='Number of classes.')
args = parser.parse_args()
# Read from tfrecord format data made by src.data.tfrecord.TFRecordWriter.
train_dataset = read_tfrecord(glob(os.path.expanduser(args.data_path)))
next_element = train_dataset.make_one_shot_iterator().get_next()
# Count the number of classes.
class_counts = np.zeros((args.num_classes, 1), dtype=np.int)
with tf.Session() as sess:
loop = 0
while True:
if loop % 100 == 0:
print(loop)
try:
label = sess.run(next_element['label'])
label = np.array(
[id2label[i].trainId for i in label.flatten()]).reshape(
label.shape)
# ignore id 255.