From e0da7eca04552e5c130b2176ce715b13a64c727b Mon Sep 17 00:00:00 2001 From: Jim Martens Date: Mon, 15 Jul 2019 11:07:20 +0200 Subject: [PATCH] Save images only from one batch Signed-off-by: Jim Martens --- src/twomartens/masterthesis/ssd.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/twomartens/masterthesis/ssd.py b/src/twomartens/masterthesis/ssd.py index 82314c3..afdcc67 100644 --- a/src/twomartens/masterthesis/ssd.py +++ b/src/twomartens/masterthesis/ssd.py @@ -268,13 +268,16 @@ def _predict_loop(generator: Generator, use_dropout: bool, steps_per_epoch: int, transform_func: callable, save_func: callable) -> None: batch_counter = 0 + saved_images = False for inputs, filenames, inverse_transforms, original_labels in generator: if use_dropout: predictions = dropout_step(inputs) else: predictions = vanilla_step(inputs) - save_images(inputs, predictions) + if not saved_images: + save_images(inputs, predictions) + saved_images = True transformed_predictions = transform_func(predictions, inverse_transforms) save_func(transformed_predictions, original_labels, filenames, batch_nr=batch_counter)