Added function to draw single bbox figure

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

src/twomartens/masterthesis/debug.py

@@ -82,10 +82,48 @@ def save_ssd_train_images(images: Union[np.ndarray, Sequence[str]], labels: np.n
                          ])
 
 
+def draw_bbox_figure(image_filename: str, labels: Sequence[np.ndarray],
+                     instances: Sequence[Sequence[np.ndarray]],
+                     image_size: int,
+                     output_path: str, coco_path: str,
+                     get_coco_cat_maps_func: callable) -> None:
+    """
+    Draws a bounding box figure and saves it under the image file name under the output path.
+    
+    Args:
+        image_filename: complete path to image file
+        labels: ground truth labels for image
+        instances: list of predictions to be compared against each other
+        image_size: size of the resized images
+        output_path: path to save the images in
+        coco_path: path to the COCO data set
+        get_coco_cat_maps_func: callable that returns the COCO category maps for a given annotation file
+    """
+    annotation_file_train = f"{coco_path}/annotations/instances_minival2014.json"
+    _, _, _, classes_to_names = get_coco_cat_maps_func(annotation_file_train)
+
+    colors = pyplot.cm.hsv(np.linspace(0, 1, len(instances) + 1)).tolist()
+    os.makedirs(output_path, exist_ok=True)
+    with Image.open(image_filename) as _image:
+        np_image = np.array(_image, dtype=np.uint8)
+    image = Image.fromarray(np_image)
+    figure_filename = f"{output_path}/{os.path.basename(image_filename)}_bboxes.png"
+    drawables = [(colors[i], _instances) for i, _instances in enumerate(instances)]
+    drawables.append((colors[-1], labels))
+    _draw_bbox_image(image=image,
+                     filename=figure_filename,
+                     draw_func=functools.partial(
+                         _draw_bboxes,
+                         image_size=image_size,
+                         classes_to_names=classes_to_names
+                     ),
+                     drawables=drawables)
+    
+
 def _draw_bbox_image(image: Image,
                      filename: str,
                      draw_func: callable,
-                     drawables: Sequence[Tuple[Sequence, Sequence[np.ndarray]]]):
+                     drawables: Sequence[Tuple[Union[Sequence, float], Sequence[np.ndarray]]]):
     figure = pyplot.figure(figsize=(6.4, 4.8))
     pyplot.imshow(image)
     
@@ -107,13 +145,18 @@ def _draw_bboxes(instances: Sequence[np.ndarray], axis: pyplot.Axes,
         if not len(instance):
             continue
         else:
-            class_id, xmin, ymin, xmax, ymax = _get_bbox_info(instance, image_size)
+            class_id, confidence, xmin, ymin, xmax, ymax = _get_bbox_info(instance, image_size)
         
         if class_id == 0:
             continue
         
-        color = colors[class_id]
+        if type(colors) is float:
-        label = f"{classes_to_names[class_id]}"
+            color = colors
+        else:
+            color = colors[class_id]
+        label = f"{classes_to_names[class_id]} - {confidence}" \
+            if confidence is not None \
+            else f"{classes_to_names[class_id]}"
         axis.add_patch(
             pyplot.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin, color=color, fill=False,
                              linewidth=2))
@@ -121,21 +164,24 @@ def _draw_bboxes(instances: Sequence[np.ndarray], axis: pyplot.Axes,
                   bbox={'facecolor': color, 'alpha': 1.0})
 
 
-def _get_bbox_info(instance: np.ndarray, image_size: int) -> Tuple[int, float, float, float, float]:
+def _get_bbox_info(instance: np.ndarray, image_size: int) -> Tuple[int, Union[float, None], float, float, float, float]:
     if len(instance) == 5:  # ground truth
         class_id = int(instance[0])
+        confidence = None
         xmin = instance[1]
         ymin = instance[2]
         xmax = instance[3]
         ymax = instance[4]
     elif len(instance) == 7:  # predictions
         class_id = int(instance[0])
+        confidence = instance[1]
         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])
+        confidence = instance[1]
         xmin = instance[2]
         ymin = instance[3]
         xmax = instance[4]
@@ -143,6 +189,7 @@ def _get_bbox_info(instance: np.ndarray, image_size: int) -> Tuple[int, float, f
     else:
         instance = np.copy(instance)
         class_id = np.argmax(instance[:-12], axis=0)
+        confidence = np.amax(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)
@@ -154,4 +201,4 @@ def _get_bbox_info(instance: np.ndarray, image_size: int) -> Tuple[int, float, f
         xmax = instance[-10]
         ymax = instance[-9]
     
-    return class_id, xmin, ymin, xmax, ymax
+    return class_id, confidence, xmin, ymin, xmax, ymax