Improve readability of the predict script
@@ -31,7 +31,6 @@ DUMMY_FILENAME = "DUMMY_FILE"
def _get_images(filenames, pre_processor, data_format):
- """ """
images = []
raw_images = []
@@ -52,17 +51,16 @@ def _get_images(filenames, pre_processor, data_format):
def _all_image_files(directory):
- all_image_files = []
- for file_path in glob(os.path.join(directory, "*")):
- if os.path.isfile(file_path) and imghdr.what(file_path) in {"jpeg", "png"}:
- all_image_files.append(os.path.abspath(file_path))
-
- return all_image_files
+ return [
+ os.path.abspath(file_path)
+ for file_path in glob(os.path.join(directory, "*"))
+ if os.path.isfile(file_path) and imghdr.what(file_path) in {"jpeg", "png"}
+ ]
def _run(input_dir, output_dir, config, restore_path, save_images):
ModelClass = config.NETWORK_CLASS
- network_kwargs = dict((key.lower(), val) for key, val in config.NETWORK.items())
+ network_kwargs = {key.lower(): val for key, val in config.NETWORK.items()}
graph = tf.Graph()
with graph.as_default():
@@ -99,34 +97,26 @@ def _run(input_dir, output_dir, config, restore_path, save_images):
results = []
for step in range(step_size):
- start_index = (step) * config.BATCH_SIZE
+ start_index = step * config.BATCH_SIZE
end_index = (step + 1) * config.BATCH_SIZE
image_files = all_image_files[start_index:end_index]
- while len(image_files) != config.BATCH_SIZE:
+ if len(image_files) < config.BATCH_SIZE:
# add dummy image.
- image_files.append(DUMMY_FILENAME)
+ image_files += [DUMMY_FILENAME] * (config.BATCH_SIZE - len(image_files))
images, raw_images = _get_images(
image_files, config.DATASET.PRE_PROCESSOR, config.DATA_FORMAT)
- feed_dict = {images_placeholder: images}
- outputs = sess.run(output_op, feed_dict=feed_dict)
+ outputs = sess.run(output_op, feed_dict={images_placeholder: images})
if config.POST_PROCESSOR:
outputs = config.POST_PROCESSOR(outputs=outputs)["outputs"]
results.append(outputs)
- writer.write(
- output_dir,
- outputs,
- raw_images,
- image_files,
- step,
- save_material=save_images
- )
+ writer.write(output_dir, outputs, raw_images, image_files, step, save_material=save_images)
return results
@@ -138,7 +128,7 @@ def run(input_dir, output_dir, experiment_id, config_file, restore_path, save_im
config = config_util.merge(config, config_util.load(config_file))
if not os.path.isdir(input_dir):
- raise Exception("Input directory {} does not exist.".format(input_dir))
+ raise FileNotFoundError("Input directory not found: '{}'".format(input_dir))
if restore_path is None:
restore_file = search_restore_filename(environment.CHECKPOINTS_DIR)
@@ -147,7 +137,7 @@ def run(input_dir, output_dir, experiment_id, config_file, restore_path, save_im
print("Restore from {}".format(restore_path))
if not os.path.exists("{}.index".format(restore_path)):
- raise Exception("restore file {} dont exists.".format(restore_path))
+ raise FileNotFoundError("Checkpoint file not found: '{}'".format(restore_path))
print("---- start predict ----")
