Added functionality to use SSD with COCO data set

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

src/twomartens/masterthesis/cli.py

@@ -227,6 +227,7 @@ def _ssd_test(args: argparse.Namespace) -> None:
     batch_size, image_size, learning_rate, \
         forward_passes_per_image, nr_classes, iou_threshold, dropout_rate, \
         use_entropy_threshold, entropy_threshold_min, entropy_threshold_max, \
+        use_coco, \
         top_k, nr_trajectories, test_pretrained, \
         coco_path, output_path, weights_path, ground_truth_path = _ssd_test_get_config_values(args, conf.get_property)
     
@@ -252,6 +253,8 @@ def _ssd_test(args: argparse.Namespace) -> None:
     ssd.compile_model(ssd_model, learning_rate, loss_func)
 
     test_generator, length_dataset, test_debug_generator = _ssd_test_get_generators(args,
+                                                                                    use_coco,
+                                                                                    data.load_coco_val_ssd,
                                                                                     data.load_scenenet_data,
                                                                                     file_names,
                                                                                     instances,
@@ -531,6 +534,7 @@ def _ssd_test_get_config_values(args: argparse.Namespace,
                                 config_get: Callable[[str], Union[str, float, int, bool]]
                                 ) -> Tuple[int, int, float, int, int, float, float,
                                            bool, float, float,
+                                           bool,
                                            int, int, bool,
                                            str, str, str, str]:
     
@@ -544,6 +548,7 @@ def _ssd_test_get_config_values(args: argparse.Namespace,
     use_entropy_threshold = config_get("Parameters.ssd_use_entropy_threshold")
     entropy_threshold_min = config_get("Parameters.ssd_entropy_threshold_min")
     entropy_threshold_max = config_get("Parameters.ssd_entropy_threshold_max")
+    use_coco = config_get("Parameters.ssd_use_coco")
     top_k = config_get("Parameters.ssd_top_k")
     nr_trajectories = config_get("Parameters.nr_trajectories")
     test_pretrained = config_get("Parameters.ssd_test_pretrained")
@@ -569,6 +574,8 @@ def _ssd_test_get_config_values(args: argparse.Namespace,
         entropy_threshold_min,
         entropy_threshold_max,
         #
+        use_coco,
+        #
         top_k,
         nr_trajectories,
         test_pretrained,
@@ -798,7 +805,9 @@ def _ssd_train_get_generators(args: argparse.Namespace,
 
 
 def _ssd_test_get_generators(args: argparse.Namespace,
-                             load_data: callable,
+                             use_coco: bool,
+                             load_data_coco: callable,
+                             load_data_scenenet: callable,
                              file_names: Sequence[Sequence[str]],
                              instances: Sequence[Sequence[Sequence[dict]]],
                              coco_path: str,
@@ -807,10 +816,21 @@ def _ssd_test_get_generators(args: argparse.Namespace,
                              nr_trajectories: int,
                              predictor_sizes: Sequence[Sequence[int]]) -> Tuple[Generator, int, Generator]:
     
+    from twomartens.masterthesis import data
+    
     if nr_trajectories == -1:
         nr_trajectories = None
     
-    generator, length, debug_generator = load_data(file_names, instances, coco_path,
+    if use_coco:
+        generator, length, debug_generator = load_data_coco(data.clean_dataset,
+                                                            data.group_bboxes_to_images,
+                                                            coco_path,
+                                                            batch_size,
+                                                            image_size,
+                                                            training=False, evaluation=True, augment=False,
+                                                            debug=args.debug)
+    else:
+        generator, length, debug_generator = load_data_scenenet(file_names, instances, coco_path,
                                                                 predictor_sizes=predictor_sizes,
                                                                 batch_size=batch_size,
                                                                 image_size=image_size,

src/twomartens/masterthesis/config.py

@@ -63,6 +63,7 @@ _CONFIG_PROPS = {
         "ssd_entropy_threshold_min": (float, "0.1"),
         "ssd_entropy_threshold_max": (float, "2.5"),
         "ssd_test_pretrained": (bool, "False"),
+        "ssd_use_coco": (bool, "False"),
         "nr_trajectories": (int, "-1")
     }
 }

src/twomartens/masterthesis/data.py

@@ -78,7 +78,7 @@ def load_coco_train(data_path: str, category: int,
     file_names.update(file_names_val)
     ids_to_images = {image['id']: image for image in images}
     
-    checked_file_names, checked_bboxes = _clean_dataset(annotations, file_names, ids_to_images)
+    checked_file_names, checked_bboxes = clean_dataset(annotations, file_names, ids_to_images)
     length_dataset = len(checked_file_names)
     
     # build image data set
@@ -119,7 +119,7 @@ def load_coco_val(data_path: str, category: int,
     file_names = {image['id']: f"{data_path}/val2014/{image['file_name']}" for image in images}
     ids_to_images = {image['id']: image for image in images}
     
-    checked_file_names, checked_bboxes = _clean_dataset(annotations, file_names, ids_to_images)
+    checked_file_names, checked_bboxes = clean_dataset(annotations, file_names, ids_to_images)
     length_dataset = len(checked_file_names)
     
     # build image data set
@@ -133,7 +133,7 @@ def load_coco_val(data_path: str, category: int,
     return dataset
 
 
-def _clean_dataset(annotations: Sequence[dict], file_names: Mapping[str, str],
+def clean_dataset(annotations: Sequence[dict], file_names: Mapping[str, str],
                   ids_to_images: Mapping[str, dict]) -> Tuple[List[str], List[List[float]]]:
     """
     Cleans a given data set from problematic cases and returns cleaned version.
@@ -152,7 +152,7 @@ def _clean_dataset(annotations: Sequence[dict], file_names: Mapping[str, str],
         img_id = annotation['image_id']
         image = ids_to_images[img_id]
         file_name = file_names[img_id]
-        bbox = annotation['bbox']
+        bbox = annotation['bbox']  # type: List[float]
         target_height = round(bbox[3])
         target_width = round(bbox[2])
         image_width, image_height = image['width'], image['height']
@@ -178,9 +178,16 @@ def _clean_dataset(annotations: Sequence[dict], file_names: Mapping[str, str],
             continue
         bbox[2] = target_width
         bbox[3] = target_height
+        new_bbox = [
+            annotation['category_id'],
+            x1,
+            y1,
+            round(bbox[0] + bbox[2]),
+            round(bbox[1] + bbox[3])
+        ]
         
         checked_file_names.append(file_name)
-        checked_bboxes.append(bbox)
+        checked_bboxes.append(new_bbox)
     
     return checked_file_names, checked_bboxes
 
@@ -232,6 +239,141 @@ def _load_images_callback(resized_shape: Sequence[int]) -> Callable[
     return _load_images
 
 
+def group_bboxes_to_images(file_names: Sequence[str], bboxes: Sequence[Sequence[int]]) -> Tuple[List[str],
+                                                                                                List[List[List[int]]]]:
+    return_labels = {}
+    for file_name, bbox in zip(file_names, bboxes):
+        if file_name not in return_labels:
+            return_labels[file_name] = []
+        return_labels[file_name].append(bbox)
+    
+    return list(return_labels.keys()), list(return_labels.values())
+
+
+def load_coco_val_ssd(clean_dataset: callable,
+                      group_bboxes_to_images: callable,
+                      coco_path: str,
+                      batch_size: int,
+                      image_size: int,
+                      training: bool,
+                      evaluation: bool,
+                      augment: bool,
+                      debug: bool,
+                      predictor_sizes: Optional[np.ndarray]) -> Tuple[Generator, int, Optional[Generator]]:
+    """
+    Loads the COCO minival2014/val2017 data and returns a data set.
+
+    Args:
+        clean_dataset: function that cleans the data set
+        group_bboxes_to_images: function that groups bounding boxes to corresponding file name
+        coco_path: path to the COCO data set
+        batch_size: batch size
+        image_size: size of images after resizing them
+        training: True if training data is desired
+        evaluation: True if evaluation-ready data is desired
+        augment: True if training data should be augmented
+        debug: True if a more extensive generator should be added to output
+        predictor_sizes: sizes of the predictor layers, can be None for evaluation
+        
+    Returns:
+        coco data set generator
+        length of dataset
+        generator which offers processed_labels as well (only if debug is True)
+    """
+    from pycocotools import coco
+    
+    from twomartens.masterthesis.ssd_keras.eval_utils import coco_utils
+    
+    annotation_file_minival = f"{coco_path}/annotations/instances_minival2014.json"
+    resized_shape = (image_size, image_size)
+    
+    coco_val = coco.COCO(annotation_file_minival)
+    img_ids = coco_val.getImgIds()  # return all image IDs belonging to given category
+    images = coco_val.loadImgs(img_ids)  # load all images
+    annotation_ids = coco_val.getAnnIds(img_ids)
+    annotations = coco_val.loadAnns(annotation_ids)  # load all image annotations
+    file_names = {image['id']: f"{coco_path}/val2014/{image['file_name']}" for image in images}
+    ids_to_images = {image['id']: image for image in images}
+
+    annotation_file_train = f"{coco_path}/annotations/instances_train2014.json"
+    cats_to_classes, _, _, _ = coco_utils.get_coco_category_maps(annotation_file_train)
+    
+    checked_image_paths, checked_bboxes = clean_dataset(annotations, file_names, ids_to_images)
+    final_image_paths, final_labels = group_bboxes_to_images(checked_image_paths, checked_bboxes)
+
+    data_generator = object_detection_2d_data_generator.DataGenerator(
+        filenames=final_image_paths,
+        labels=final_labels
+    )
+
+    shuffle = True if training else False
+
+    if training and augment:
+        transformations = [data_augmentation_chain_original_ssd.SSDDataAugmentation(
+            img_width=resized_shape[0],
+            img_height=resized_shape[1]
+        )]
+    else:
+        transformations = [
+            object_detection_2d_photometric_ops.ConvertTo3Channels(),
+            object_detection_2d_geometric_ops.Resize(height=resized_shape[0],
+                                                     width=resized_shape[1])
+        ]
+
+    returns = {'processed_images', 'encoded_labels'}
+    returns_debug = {'processed_images', 'encoded_labels', 'processed_labels'}
+
+    if not training and evaluation:
+        returns = {
+            'processed_images',
+            'filenames',
+            'inverse_transform',
+            'original_labels'}
+        label_encoder = None
+    else:
+        if predictor_sizes is None:
+            raise ValueError("predictor_sizes cannot be None for training/validation")
+        label_encoder = ssd_input_encoder.SSDInputEncoder(
+            img_height=resized_shape[0],
+            img_width=resized_shape[1],
+            n_classes=len(cats_to_classes),  # 80
+            predictor_sizes=predictor_sizes,
+            steps=[8, 16, 32, 64, 100, 300],
+            coords="corners",
+            aspect_ratios_per_layer=[[1.0, 2.0, 0.5],
+                                     [1.0, 2.0, 0.5, 3.0, 1.0 / 3.0],
+                                     [1.0, 2.0, 0.5, 3.0, 1.0 / 3.0],
+                                     [1.0, 2.0, 0.5, 3.0, 1.0 / 3.0],
+                                     [1.0, 2.0, 0.5],
+                                     [1.0, 2.0, 0.5]]
+        )
+
+    if debug:
+        debug_generator = data_generator.generate(
+            batch_size=batch_size,
+            shuffle=shuffle,
+            transformations=transformations,
+            label_encoder=label_encoder,
+            returns=returns_debug,
+            keep_images_without_gt=False
+        )
+    else:
+        debug_generator = None
+
+    length_dataset = data_generator.dataset_size
+
+    generator = data_generator.generate(
+        batch_size=batch_size,
+        shuffle=shuffle,
+        transformations=transformations,
+        label_encoder=label_encoder,
+        returns=returns,
+        keep_images_without_gt=False
+    )
+    
+    return generator, length_dataset, debug_generator
+
+
 def load_scenenet_data(photo_paths: Sequence[Sequence[str]],
                        instances: Sequence[Sequence[Sequence[dict]]],
                        coco_path: str,