diff --git a/src/twomartens/masterthesis/debug.py b/src/twomartens/masterthesis/debug.py index a1d3a81..db88b25 100644 --- a/src/twomartens/masterthesis/debug.py +++ b/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)) - pyplot.imshow(image) + _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=[ + (colors, instances) + ]) + + +def _draw_bbox_image(image: Image, + filename: str, + draw_func: callable, + drawables: Sequence[Tuple[Sequence, Sequence[np.ndarray]]]): + figure = pyplot.figure(figsize=(6.4, 4.8)) + pyplot.imshow(image) + + current_axis = pyplot.gca() + for colors, instances in drawables: + draw_func(instances=instances, + axis=current_axis, + colors=colors) + + pyplot.savefig(filename) + pyplot.close(figure) + + +def _draw_bboxes(instances: Sequence[np.ndarray], axis: pyplot.Axes, + image_size: int, + colors: Sequence, + classes_to_names: Dict[int, str]) -> None: + for instance in instances: + if not len(instance): + continue + else: + class_id, xmin, ymin, xmax, ymax = _get_bbox_info(instance, image_size) - current_axis = pyplot.gca() + if class_id == 0: + continue - for instance in instances: - 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] - 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] - ymin = instance[-11] - xmax = instance[-10] - ymax = instance[-9] - - if class_id == 0: - continue - - 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}/{custom_string}bboxes{str(i).zfill(nr_digits)}.png") - pyplot.close(figure) + 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