Extracted debug functionality into separate module

Signed-off-by: Jim Martens <github@2martens.de>

src/twomartens/masterthesis/cli.py

@@ -173,6 +173,7 @@ def _ssd_train(args: argparse.Namespace) -> None:
     image_size = conf.get_property("Parameters.ssd_image_size")
     use_dropout = False if args.network == "ssd" else True
 
+    summary_path = conf.get_property("Paths.summaries")
     weights_path = conf.get_property("Paths.weights")
     pre_trained_weights_file = f"{weights_path}/VGG_coco_SSD_300x300_iter_400000.h5"
     weights_path = f"{weights_path}/train/{args.network}/"
@@ -213,57 +214,20 @@ def _ssd_train(args: argparse.Namespace) -> None:
     del file_names_train, instances_train, file_names_val, instances_val
     
     if args.debug and conf.get_property("Debug.train_images"):
-        from matplotlib import pyplot
-        import numpy as np
-        from PIL import Image
-        
         from twomartens.masterthesis.ssd_keras.eval_utils import coco_utils
+        from twomartens.masterthesis import debug
         
         train_data = next(train_generator)
         train_length -= batch_size
         train_images = train_data[0]
         train_labels = train_data[1]
+        output_path = f"{summary_path}/train/{args.network}/{args.iteration}"
 
-        annotation_file_train = f"{args.coco_path}/annotations/instances_train2014.json"
+        debug.save_ssd_train_images(train_images, train_labels, output_path)
-        _, _, _, classes_to_names = coco_utils.get_coco_category_maps(annotation_file_train)
-        colors = pyplot.cm.hsv(np.linspace(0, 1, 81)).tolist()
-        summary_path = f"{args.summary_path}/train/{args.network}/{args.iteration}/"
-        os.makedirs(summary_path, exist_ok=True)
-
-        nr_images = len(train_images)
-        nr_digits = math.ceil(math.log10(nr_images))
-        
-        for i, train_image in enumerate(train_images):
-            instances = train_labels[i]
-            image = Image.fromarray(train_image)
-            image.save(f"{summary_path}"
-                       f"train_image{str(i).zfill(nr_digits)}.png")
-            
-            figure = pyplot.figure(figsize=(20, 12))
-            pyplot.imshow(image)
-
-            current_axis = pyplot.gca()
-
-            for instance in instances:
-                xmin = instance[-12] * image_size
-                ymin = instance[-11] * image_size
-                xmax = instance[-10] * image_size
-                ymax = instance[-9] * image_size
-                class_id = np.argmax(instance[:-12], axis=0)
-                color = colors[class_id]
-                label = f"{classes_to_names[class_id]}"
-                current_axis.add_patch(
-                    pyplot.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin, color=color, fill=False,
-                                     linewidth=2))
-                current_axis.text(xmin, ymin, label, size='x-large', color='white',
-                                  bbox={'facecolor': color, 'alpha': 1.0})
-            pyplot.savefig(f"{args.summary_path}/train/{args.network}/{args.iteration}/bboxes{str(i).zfill(nr_digits)}.png")
-            pyplot.close(figure)
     
     nr_batches_train = int(math.floor(train_length / batch_size))
     nr_batches_val = int(math.floor(val_length / batch_size))
     
-    summary_path = conf.get_property("Paths.summaries")
     tensorboard_callback = tf.keras.callbacks.TensorBoard(
         log_dir=f"{summary_path}/train/{args.network}/{args.iteration}"
     )

src/twomartens/masterthesis/debug.py

@@ -0,0 +1,71 @@
+# -*- coding: utf-8 -*-
+
+#   Copyright 2018 Timon Brüning, Inga Kempfert, Anne Kunstmann, Jim Martens,
+#                  Marius Pierenkemper, Yanneck Reiss
+#
+#   Licensed under the Apache License, Version 2.0 (the "License");
+#   you may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#   See the License for the specific language governing permissions and
+#   limitations under the License.
+
+"""
+Handle debug functionality.
+
+Functions:
+    save_ssd_train_images(images, labels, output_path):
+        saves the first batch of SSD train images with overlaid ground truth bounding boxes
+"""
+import os
+
+import math
+import numpy as np
+from matplotlib import pyplot
+from PIL import Image
+
+from twomartens.masterthesis import config
+from twomartens.masterthesis.ssd_keras.eval_utils import coco_utils
+
+
+def save_ssd_train_images(images: np.ndarray, labels: np.ndarray, output_path: str) -> None:
+    annotation_file_train = f"{config.get_property('Paths.coco')}/annotations/instances_train2014.json"
+    _, _, _, classes_to_names = coco_utils.get_coco_category_maps(annotation_file_train)
+    colors = pyplot.cm.hsv(np.linspace(0, 1, 81)).tolist()
+    os.makedirs(output_path, exist_ok=True)
+    
+    nr_images = len(images)
+    nr_digits = math.ceil(math.log10(nr_images))
+    image_size = config.get_property("Parameters.ssd_image_size")
+    
+    for i, train_image in enumerate(images):
+        instances = labels[i]
+        image = Image.fromarray(train_image)
+        image.save(f"{output_path}"
+                   f"train_image{str(i).zfill(nr_digits)}.png")
+        
+        figure = pyplot.figure(figsize=(20, 12))
+        pyplot.imshow(image)
+        
+        current_axis = pyplot.gca()
+        
+        for instance in instances:
+            xmin = instance[-12] * image_size
+            ymin = instance[-11] * image_size
+            xmax = instance[-10] * image_size
+            ymax = instance[-9] * image_size
+            class_id = np.argmax(instance[:-12], axis=0)
+            color = colors[class_id]
+            label = f"{classes_to_names[class_id]}"
+            current_axis.add_patch(
+                pyplot.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin, color=color, fill=False,
+                                 linewidth=2))
+            current_axis.text(xmin, ymin, label, size='x-large', color='white',
+                              bbox={'facecolor': color, 'alpha': 1.0})
+        pyplot.savefig(f"{output_path}/bboxes{str(i).zfill(nr_digits)}.png")
+        pyplot.close(figure)