Moved code into separate functions

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

src/twomartens/masterthesis/debug.py

@@ -22,10 +22,12 @@ Functions:
     save_ssd_train_images(...):
         saves the first batch of SSD train images with overlaid ground truth bounding boxes
 """
+import functools
 import os
+from typing import Dict
+from typing import Tuple
 from typing import Union, Sequence
 
-import cv2
 import math
 import numpy as np
 from matplotlib import pyplot
@@ -68,56 +70,88 @@ def save_ssd_train_images(images: Union[np.ndarray, Sequence[str]], labels: np.n
         image = Image.fromarray(train_image)
         image.save(f"{output_path}/"
                    f"{custom_string}train_image{str(i).zfill(nr_digits)}.png")
+        figure_filename = f"{output_path}/{custom_string}bboxes{str(i).zfill(nr_digits)}.png"
         
-        figure = pyplot.figure(figsize=(6.4, 4.8))
+        _draw_bbox_image(image=image,
-        pyplot.imshow(image)
+                         filename=figure_filename,
+                         draw_func=functools.partial(_draw_bboxes,
+                                                     image_size=image_size,
+                                                     classes_to_names=classes_to_names),
+                         drawables=[
+                             (colors, instances)
+                         ])
 
-        current_axis = pyplot.gca()
 
-        for instance in instances:
+def _draw_bbox_image(image: Image,
-            if len(instance) == 5:  # ground truth
+                     filename: str,
-                class_id = int(instance[0])
+                     draw_func: callable,
-                xmin = instance[1]
+                     drawables: Sequence[Tuple[Sequence, Sequence[np.ndarray]]]):
-                ymin = instance[2]
+    figure = pyplot.figure(figsize=(6.4, 4.8))
-                xmax = instance[3]
+    pyplot.imshow(image)
-                ymax = instance[4]
-            elif len(instance) == 7:  # predictions
-                class_id = int(instance[0])
-                xmin = instance[3]
-                ymin = instance[4]
-                xmax = instance[5]
-                ymax = instance[6]
-            elif len(instance) == 6:  # predictions using Caffe method
-                class_id = int(instance[0])
-                xmin = instance[2]
-                ymin = instance[3]
-                xmax = instance[4]
-                ymax = instance[5]
-            elif not len(instance):
-                continue
-            else:
-                instance = np.copy(instance)
-                class_id = np.argmax(instance[:-12], axis=0)
-                instance[-12:-8] *= instance[-4:]  # multiply with variances
-                instance[[-11, -9]] *= np.expand_dims(instance[-5] - instance[-7], axis=-1)
-                instance[[-12, -10]] *= np.expand_dims(instance[-6] - instance[-8], axis=-1)
-                instance[-12:-8] += instance[-8:-4]
-                instance[-12:-8] *= image_size
     
-                xmin = instance[-12]
+    current_axis = pyplot.gca()
-                ymin = instance[-11]
+    for colors, instances in drawables:
-                xmax = instance[-10]
+        draw_func(instances=instances,
-                ymax = instance[-9]
+                  axis=current_axis,
+                  colors=colors)
     
-            if class_id == 0:
+    pyplot.savefig(filename)
-                continue
+    pyplot.close(figure)
 
-            color = colors[class_id]
+
-            label = f"{classes_to_names[class_id]}"
+def _draw_bboxes(instances: Sequence[np.ndarray], axis: pyplot.Axes,
-            current_axis.add_patch(
+                 image_size: int,
-                pyplot.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin, color=color, fill=False,
+                 colors: Sequence,
-                                 linewidth=2))
+                 classes_to_names: Dict[int, str]) -> None:
-            current_axis.text(xmin, ymin, label, size='x-large', color='white',
+    for instance in instances:
-                              bbox={'facecolor': color, 'alpha': 1.0})
+        if not len(instance):
-        pyplot.savefig(f"{output_path}/{custom_string}bboxes{str(i).zfill(nr_digits)}.png")
+            continue
-        pyplot.close(figure)
+        else:
+            class_id, xmin, ymin, xmax, ymax = _get_bbox_info(instance, image_size)
+        
+        if class_id == 0:
+            continue
+        
+        color = colors[class_id]
+        label = f"{classes_to_names[class_id]}"
+        axis.add_patch(
+            pyplot.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin, color=color, fill=False,
+                             linewidth=2))
+        axis.text(xmin, ymin, label, size='x-large', color='white',
+                  bbox={'facecolor': color, 'alpha': 1.0})
+
+
+def _get_bbox_info(instance: np.ndarray, image_size: int) -> Tuple[int, float, float, float, float]:
+    if len(instance) == 5:  # ground truth
+        class_id = int(instance[0])
+        xmin = instance[1]
+        ymin = instance[2]
+        xmax = instance[3]
+        ymax = instance[4]
+    elif len(instance) == 7:  # predictions
+        class_id = int(instance[0])
+        xmin = instance[3]
+        ymin = instance[4]
+        xmax = instance[5]
+        ymax = instance[6]
+    elif len(instance) == 6:  # predictions using Caffe method
+        class_id = int(instance[0])
+        xmin = instance[2]
+        ymin = instance[3]
+        xmax = instance[4]
+        ymax = instance[5]
+    else:
+        instance = np.copy(instance)
+        class_id = np.argmax(instance[:-12], axis=0)
+        instance[-12:-8] *= instance[-4:]  # multiply with variances
+        instance[[-11, -9]] *= np.expand_dims(instance[-5] - instance[-7], axis=-1)
+        instance[[-12, -10]] *= np.expand_dims(instance[-6] - instance[-8], axis=-1)
+        instance[-12:-8] += instance[-8:-4]
+        instance[-12:-8] *= image_size
+        
+        xmin = instance[-12]
+        ymin = instance[-11]
+        xmax = instance[-10]
+        ymax = instance[-9]
+    
+    return class_id, xmin, ymin, xmax, ymax