def _convert_dataset(self, dataset_split):
"""Converts the specified dataset split to TFRecord format.
Args:
dataset_split: The dataset split (e.g., train, test).
Raises:
RuntimeError: If loaded image and label have different shape.
"""
dataset = os.path.basename(dataset_split)[:-4]
sys.stdout.write('Processing ' + dataset)
filenames = [x.strip('\n') for x in open(dataset_split, 'r')]
num_images = len(filenames)
num_per_shard = int(math.ceil(num_images / _NUM_SHARDS))
image_reader = build_data.ImageReader('jpeg', channels=3)
label_reader = build_data.ImageReader('png', channels=1)
for shard_id in range(_NUM_SHARDS):
output_filename = os.path.join(
self.output_dir,
'%s-%05d-of-%05d.tfrecord' % (dataset, shard_id, _NUM_SHARDS))
with tf.python_io.TFRecordWriter(
output_filename) as tfrecord_writer:
start_idx = shard_id * num_per_shard
end_idx = min((shard_id + 1) * num_per_shard, num_images)
for i in range(start_idx, end_idx):
sys.stdout.write('\r>> Converting image %d/%d shard %d' %
(i + 1, len(filenames), shard_id))
sys.stdout.flush()
# Read the image.
image_filename = os.path.join(
self.image_folder,
filenames[i] + '.' + self.image_format)
image_data = tf.gfile.GFile(image_filename, 'rb').read()
height, width = image_reader.read_image_dims(image_data)
# Read the semantic segmentation annotation.
seg_filename = os.path.join(
self.semantic_segmentation_folder,
filenames[i] + '.' + self.label_format)
seg_data = tf.gfile.GFile(seg_filename, 'rb').read()
seg_height, seg_width = label_reader.read_image_dims(
seg_data)
if height != seg_height or width != seg_width:
raise RuntimeError(
'Shape mismatched between image and label.')
# Convert to tf example.
example = build_data.image_seg_to_tfexample(
image_data, filenames[i], height, width, seg_data)
tfrecord_writer.write(example.SerializeToString())
sys.stdout.write('\n')
sys.stdout.flush()
def _convert_dataset(tfrec_name, dataset_dir, dataset_label_dir):
img_names = tf.gfile.Glob(os.path.join(dataset_dir,
f'*.{PARAM.image_ext}'))
random.shuffle(img_names)
seg_names = []
for f in img_names:
basename = os.path.basename(f).split('.')[0]
seg = os.path.join(dataset_label_dir, f'{basename}.{PARAM.label_ext}')
seg_names.append(seg)
num_images = len(img_names)
num_per_shard = int(math.ceil(num_images / float(PARAM.num_shards)))
image_reader = build_data.ImageReader(
'png' if PARAM.image_ext == 'png' else 'jpeg',
channels=PARAM.image_nchannels)
label_reader = build_data.ImageReader(
'png' if PARAM.image_ext == 'png' else 'jpeg',
channels=PARAM.label_nchannels)
for shard_id in range(PARAM.num_shards):
output_filename = os.path.join(
PARAM.output_dir, '%s-%05d-of-%05d.tfrecord' %
(tfrec_name, shard_id, PARAM.num_shards))
with tf.python_io.TFRecordWriter(output_filename) as tfrecord_writer:
start_idx = shard_id * num_per_shard
end_idx = min((shard_id + 1) * num_per_shard, num_images)
for i in range(start_idx, end_idx):
print(
f'\r>> Converting image {i+1}/{num_images} shard {shard_id}'
)
# Read the image.
image_filename = img_names[i]
image_data = tf.gfile.FastGFile(image_filename, 'rb').read()
height, width = image_reader.read_image_dims(image_data)
# Read the semantic segmentation annotation.
seg_filename = seg_names[i]
seg_data = tf.gfile.FastGFile(seg_filename, 'rb').read()
seg_height, seg_width = label_reader.read_image_dims(seg_data)
if height != seg_height or width != seg_width:
raise RuntimeError(
'Shape mismatched between image and label.')
# Convert to tf example.
example = build_data.image_seg_to_tfexample(
image_data, img_names[i], height, width, seg_data)
tfrecord_writer.write(example.SerializeToString())
def _convert_dataset(dataset_split):
"""Converts the specified dataset split to TFRecord format.
Args:
dataset_split: The dataset split (e.g., train, val).
Raises:
RuntimeError: If loaded image and label have different shape, or if the
image file with specified postfix could not be found.
"""
image_files = _get_files('image', dataset_split)
label_files = _get_files('label', dataset_split)
image_files = image_files[:10]
label_files = image_files[:10]
num_images = len(image_files)
num_per_shard = int(math.ceil(num_images / float(_NUM_SHARDS)))
image_reader = build_data.ImageReader('png', channels=3)
label_reader = build_data.ImageReader('png', channels=1)
for shard_id in range(_NUM_SHARDS):
shard_filename = '%s-%05d-of-%05d.tfrecord' % (dataset_split, shard_id,
_NUM_SHARDS)
output_filename = os.path.join(FLAGS.output_dir, shard_filename)
with tf.python_io.TFRecordWriter(output_filename) as tfrecord_writer:
start_idx = shard_id * num_per_shard
end_idx = min((shard_id + 1) * num_per_shard, num_images)
for i in range(start_idx, end_idx):
sys.stdout.write('\r>> Converting image %d/%d shard %d' %
(i + 1, num_images, shard_id))
sys.stdout.flush()
# Read the image.
image_data = tf.gfile.FastGFile(image_files[i], 'rb').read()
height, width = image_reader.read_image_dims(image_data)
# Read the semantic segmentation annotation.
seg_data = tf.gfile.FastGFile(label_files[i], 'rb').read()
# print('seg_data type',type(seg_data))
label = cv2.imread(label_files[i], cv2.IMREAD_GRAYSCALE)
edge = get_edge(label)
# edge_bytes=tf.train.BytesList(value=edge)
# edge_feature=tf.train.Feature(int64_list=tf.train.Int64List(value=edge.reshape(-1)))
edge_feature = tf.train.Feature(float_list=tf.train.FloatList(
value=edge))
seg_height, seg_width = label_reader.read_image_dims(seg_data)
if height != seg_height or width != seg_width:
raise RuntimeError(
'Shape mismatched between image and label.')
# Convert to tf example.
re_match = _IMAGE_FILENAME_RE.search(image_files[i])
if re_match is None:
raise RuntimeError('Invalid image filename: ' +
image_files[i])
filename = os.path.basename(re_match.group(1))
example = build_data.image_seg_edge_to_tfexample(
image_data, filename, height, width, seg_data,
edge_feature)
tfrecord_writer.write(example.SerializeToString())
sys.stdout.write('\n')
sys.stdout.flush()
def main(unused_argv):
tf.logging.set_verbosity(tf.logging.INFO)
tf.gfile.MakeDirs(FLAGS.inference_dir)
g = tf.Graph()
with g.as_default():
image_name = FLAGS.image_path.split('/')[-1]
image_name, image_extension = image_name.split('.')
supported_extensions = ['png', 'jpeg', 'jpg']
if not any(image_extension == extension
for extension in supported_extensions):
raise ValueError('Image extension "{}" not supported...'.format(
image_extension))
reader = build_data.ImageReader(image_extension)
image = reader.decode_image(
tf.gfile.FastGFile(FLAGS.image_path, 'r').read())
image = tf.identity(image)
original_image_dimensions = image.get_shape().as_list()[0:2]
original_image_dimensions = reversed(original_image_dimensions)
image = tf.image.resize_images(image, [480, 640],
method=tf.image.ResizeMethod.BILINEAR,
align_corners=True)
image.set_shape([None, None, 3])
image = tf.expand_dims(image, 0)
model_options = common.ModelOptions(
outputs_to_num_classes={common.OUTPUT_TYPE: FLAGS.num_classes},
crop_size=FLAGS.inference_crop_size,
atrous_rates=FLAGS.atrous_rates,
output_stride=FLAGS.output_stride)
predictions = model.predict_labels(image,
model_options=model_options,
image_pyramid=None)
predictions = predictions[common.OUTPUT_TYPE]
# predictions = tf.image.resize_images(
# predictions, original_image_dimensions,
# method=tf.image.ResizeMethod.BILINEAR,
# align_corners=True)
param_stats = tf.profiler.profile(
tf.get_default_graph(),
options=tf.profiler.ProfileOptionBuilder.
trainable_variables_parameter())
print('Total parameters: ', param_stats.total_parameters)
total_parameters = 0
for variable in tf.trainable_variables():
shape = variable.get_shape()
variable_parameters = 1
for dim in shape:
variable_parameters *= dim.value
total_parameters += variable_parameters
print('Total parameters: ', total_parameters)
tf.train.get_or_create_global_step()
saver = tf.train.Saver(slim.get_variables_to_restore())
sv = tf.train.Supervisor(graph=g,
logdir=FLAGS.inference_dir,
init_op=tf.global_variables_initializer(),
summary_op=None,
summary_writer=None,
global_step=None,
saver=saver)
with sv.managed_session(start_standard_services=False) as sess:
sv.start_queue_runners(sess)
sv.saver.restore(sess, FLAGS.checkpoint_path)
semantic_predictions = sess.run(predictions)
result = np.array(semantic_predictions, dtype=np.uint8)
result = np.squeeze(result)
result = cv2.resize(result, tuple(original_image_dimensions))
# save raw result...
save_annotation.save_annotation(result,
FLAGS.inference_dir,
_RAW_FORMAT % image_name,
add_colormap=False)
# save result as color image...
save_annotation.save_annotation(result,
FLAGS.inference_dir,
_PREDICTION_FORMAT % image_name,
add_colormap=True,
colormap_type=FLAGS.dataset)