Feature/base footprint detection

config/yoeo.cfg

@@ -4,7 +4,7 @@
 #subdivisions=1
 # Training
 batch=64
-subdivisions=8
+subdivisions=4
 width=416
 height=416
 channels=3
@@ -19,8 +19,8 @@ learning_rate=0.001
 burn_in=100
 max_batches = 4000
 policy=steps
-steps=50000,60000
-scales=.1,.1
+steps=80000,250000
+scales=.1,.2
 
 ####
 # Like YOEO rev 2 but with deeper skip connections
@@ -213,15 +213,15 @@ activation=leaky
 size=1
 stride=1
 pad=1
-filters=24
+filters=12
 activation=linear
 
 
 
 [yolo]
-mask = 3,4,5
-anchors = 17, 32,  13,171,  37, 67,  30,224,  69,112, 116,212
-classes=3
+mask = 0
+anchors = 100, 100
+classes=4
 num=6
 jitter=.3
 ignore_thresh = .7
@@ -257,13 +257,13 @@ activation=leaky
 size=1
 stride=1
 pad=1
-filters=24
+filters=12
 activation=linear
 
 [yolo]
-mask = 0,1,2
-anchors = 17, 32,  13,171,  37, 67,  30,224,  69,112, 116,212
-classes=3
+mask = 0
+anchors = 100, 100
+classes=4
 num=6
 jitter=.3
 ignore_thresh = .7

yoeo/detect.py

@@ -71,10 +71,10 @@ def detect_directory(model_path, weights_path, img_path, class_config: ClassConf
     print(f"---- Detections were saved to: '{output_path}' ----")
 
 
-def detect_image(model, 
-                 image: np.ndarray,  
-                 img_size: int = 416, 
-                 conf_thres: float = 0.5, 
+def detect_image(model,
+                 image: np.ndarray,
+                 img_size: int = 416,
+                 conf_thres: float = 0.5,
                  nms_thres: float = 0.5,
                  group_config: Optional[GroupConfig] = None
                  ):
@@ -97,7 +97,7 @@ def detect_image(model,
     :rtype: nd.array, nd.array
     """
     model.eval()  # Set model to evaluation mode
-    
+
     # Configure input
     input_img = transforms.Compose([
         DEFAULT_TRANSFORMS,
@@ -113,9 +113,9 @@ def detect_image(model,
     with torch.no_grad():
         detections, segmentations = model(input_img)
         detections = non_max_suppression(
-            prediction=detections, 
-            conf_thres=conf_thres, 
-            iou_thres=nms_thres, 
+            prediction=detections,
+            conf_thres=conf_thres,
+            iou_thres=nms_thres,
             group_config=group_config
         )
         detections = rescale_boxes(detections[0], img_size, image.shape[0:2])
@@ -124,9 +124,9 @@ def detect_image(model,
 
 
 def detect(model,
-           dataloader: DataLoader, 
-           output_path: str, 
-           conf_thres: float = 0.5, 
+           dataloader: DataLoader,
+           output_path: str,
+           conf_thres: float = 0.5,
            nms_thres: float = 0.5,
            group_config: Optional[GroupConfig] = None
             ):
@@ -169,9 +169,9 @@ def detect(model,
         with torch.no_grad():
             detections, segmentations = model(input_imgs)
             detections = non_max_suppression(
-                prediction=detections, 
-                conf_thres=conf_thres, 
-                iou_thres=nms_thres, 
+                prediction=detections,
+                conf_thres=conf_thres,
+                iou_thres=nms_thres,
                 group_config=group_config
             )
 
@@ -254,7 +254,7 @@ def _draw_and_save_output_image(image_path, detections, seg, img_size, output_pa
     # Bounding-box colors
     cmap = plt.get_cmap("tab20b")
     colors = [cmap(i) for i in np.linspace(0, 1, len(classes))]
-    for x1, y1, x2, y2, conf, cls_pred in detections:
+    for x1, y1, x2, y2, conf, cls_pred, b, bx, by in detections:
 
         print(f"\t+ Label: {classes[int(cls_pred)]} | Confidence: {conf.item():0.4f}")
 
@@ -265,16 +265,19 @@ def _draw_and_save_output_image(image_path, detections, seg, img_size, output_pa
         bbox = patches.Rectangle((x1, y1), box_w, box_h, linewidth=2, edgecolor=colors[int(cls_pred)], facecolor="none")
         # Add the bbox to the plot
         ax.add_patch(bbox)
+
+        # Add red dot for base_footprint if it exists
+        if b > 0.5:
+            ax.plot(bx, by, 'ro')
+
         # Add label
-        """
         plt.text(
             x1,
             y1,
             s=classes[int(cls_pred)],
             color="white",
             verticalalignment="top",
             bbox={"color": colors[int(cls_pred)], "pad": 0})
-        """
 
     # Save generated image with detections
     plt.axis("off")

yoeo/models.py

@@ -8,7 +8,7 @@
 import numpy as np
 
 from yoeo.utils.parse_config import parse_model_config
 from yoeo.utils.utils import weights_init_normal, to_cpu, seg_iou
 
 
 def create_modules(module_defs):
@@ -117,7 +117,7 @@
         x = F.interpolate(x, scale_factor=self.scale_factor, mode=self.mode)
         return x
 
 class Mish(nn.Module):
     """ The MISH activation function (https://github.com/digantamisra98/Mish) """
 
     def __init__(self):
@@ -135,7 +135,7 @@
         self.num_classes = num_classes
         self.mse_loss = nn.MSELoss()
         self.bce_loss = nn.BCELoss()
-        self.no = num_classes + 5  # number of outputs per anchor
+        self.no = num_classes + 5 + 3   # TODO make basefootprint param    # number of outputs per anchor
         self.grid = torch.zeros(1)  # TODO
 
         anchors = torch.tensor(list(chain(*anchors))).float().view(-1, 2)
@@ -157,7 +157,9 @@
             x = torch.cat([
                 (x[..., 0:2].sigmoid() + self.grid) * stride,  # xy
                 torch.exp(x[..., 2:4]) * self.anchor_grid, # wh
-                x[..., 4:].sigmoid(),
+                x[..., 4:-3].sigmoid(),
+                x[..., -3, None].sigmoid(), # basefootprint confidence
+                (x[..., -2:].tan() + self.grid) * stride # basefoot print xy
             ], axis=4).view(bs, -1, self.no)
 
         return x

yoeo/scripts/createYOEOLabelsFromTORSO-21.py

@@ -30,6 +30,7 @@ def range_limited_float_type_0_to_1(arg):
 parser.add_argument("--skip-concealed", action="store_true", help="Skip concealed labels")
 parser.add_argument("--skip-classes", nargs="+", default=[], help="These bounding box classes will be skipped")
 parser.add_argument("--robots-with-team-colors", action="store_true", help="The robot class will be subdivided into subclasses, one for each team color (currently either 'blue', 'red' or 'unknown').")
+parser.add_argument("--add-base-footprint", action="store_true", help="Add base_footprint to bounding box annotations")
 args = parser.parse_args()
 
 # Available classes for YOEO
@@ -54,6 +55,9 @@ def range_limited_float_type_0_to_1(arg):
 dataset_collection_dir = args.dataset_collection_dir
 destination_dir = args.dataset_collection_dir
 
+# Keep track of all classes with base_footprint annotations
+classes_with_base_footprint = set()
+
 # Overwrite defaults, if destination path is given
 if args.destination_dir:
     create_symlinks = True
@@ -138,10 +142,10 @@ def range_limited_float_type_0_to_1(arg):
                 (args.skip_concealed and annotation.get('concealed', False))):
                 continue
             elif class_name in CLASSES['bb_classes']:  # Handle bounding boxes
-                min_x = min(map(lambda x: x[0], annotation['vector']))
-                max_x = max(map(lambda x: x[0], annotation['vector']))
-                min_y = min(map(lambda x: x[1], annotation['vector']))
-                max_y = max(map(lambda x: x[1], annotation['vector']))
+                min_x = min(pt[0] for pt in annotation['vector'])
+                max_x = max(pt[0] for pt in annotation['vector'])
+                min_y = min(pt[1] for pt in annotation['vector'])
+                max_y = max(pt[1] for pt in annotation['vector'])
 
                 annotation_width = max_x - min_x
                 annotation_height = max_y - min_y
@@ -154,8 +158,27 @@ def range_limited_float_type_0_to_1(arg):
                 relative_center_y = center_y / img_height
 
                 # Derive classID from index in predefined classes
-                classID = CLASSES['bb_classes'].index(class_name)                
-                annotations.append(f"{classID} {relative_center_x} {relative_center_y} {relative_annotation_width} {relative_annotation_height}")
+                classID = CLASSES['bb_classes'].index(class_name)
+
+                # Serialize bounding box
+                serialized_bounding_box = f"{classID} {relative_center_x} {relative_center_y} {relative_annotation_width} {relative_annotation_height}"
+
+                # Extract other properties
+                if args.add_base_footprint:
+                    # Default value if key point is not present
+                    base_footprint_x, base_footprint_y = -1, -1
+                    # Check if the annotation exists (some classes do not have base_footprint)
+                    if 'base_footprint' in annotation:
+                        classes_with_base_footprint.add(class_name)
+                        # Check if the base_footprint is not None (negative label)
+                        if annotation['base_footprint']:
+                            base_footprint_x = annotation['base_footprint'][0] / img_width
+                            base_footprint_y = annotation['base_footprint'][1] / img_height
+                    # Serialize base_footprint key point
+                    serialized_bounding_box += f" {base_footprint_x} {base_footprint_y}"
+
+                # Add row to file
+                annotations.append(serialized_bounding_box)
             else:
                 print(f"The annotation type '{class_name}' is not supported. Image: '{img_name_with_extension}'")
 
@@ -180,8 +203,14 @@ def range_limited_float_type_0_to_1(arg):
 names_path = os.path.join(destination_dir, "yoeo_names.yaml")
 names = {
     'detection': CLASSES['bb_classes'],
-    'segmentation': CLASSES["segmentation_classes"],
+    'segmentation': CLASSES["segmentation_classes"]
 }
+
+# Add base_footprint classes to names file
+if args.add_base_footprint:
+    names['base_footprint'] = list(classes_with_base_footprint)
+
+# Save names to yaml file
 with open(names_path, "w") as names_file:
     yaml.dump(names, names_file)
 

yoeo/test.py

@@ -50,12 +50,11 @@
     :type nms_thres: float, optional
     :param verbose: If True, prints stats of model, defaults to True
     :type verbose: bool, optional
-    :return: Returns precision, recall, AP, f1, ap_class
     """
     dataloader = _create_validation_data_loader(
         img_path, batch_size, img_size, n_cpu)
     model = load_model(model_path, weights_path)
-    metrics_output, seg_class_ious, secondary_metric = _evaluate(
+    metrics_output, seg_class_ious, secondary_metric, base_footprint_metric  = _evaluate(
         model,
         dataloader,
         class_config,
@@ -64,13 +63,14 @@
         conf_thres,
         nms_thres,
         verbose)
-    return metrics_output, seg_class_ious, secondary_metric
+    return metrics_output, seg_class_ious, secondary_metric, base_footprint_metric
 
 
-def print_eval_stats(metrics_output: Optional[Tuple[np.ndarray]], 
-                     seg_class_ious: List[np.float64], 
-                     secondary_metric: Optional[Metric], 
-                     class_config: ClassConfig, 
+def print_eval_stats(metrics_output: Optional[Tuple[np.ndarray]],
+                     seg_class_ious: List[np.float64],
+                     secondary_metric: Optional[Metric],
+                     base_footprint_metrics: dict[str, float],
+                     class_config: ClassConfig,
                      verbose: bool
                      ):
     # Print detection statistics
@@ -80,7 +80,7 @@
         if verbose:
             # Prints class AP and mean AP
             ap_table = [["Index", "Class", "AP"]]
-            class_names = class_config.get_squeezed_det_class_names()
+            class_names = class_config.get_ungrouped_det_class_names()
             for i, c in enumerate(ap_class):
                 ap_table += [[c, class_names[c], "%.5f" % AP[i]]]
             print(AsciiTable(ap_table).table)
@@ -95,7 +95,7 @@
         if verbose:
             classes = class_config.get_group_class_names()
             mbACC_per_class = [secondary_metric.bACC(i) for i in range(len(classes))]
-                        
+
             sec_table = [["Index", "Class", "bACC"]]
             for i, c in enumerate(classes):
                 sec_table += [[i, c, "%.5f" % mbACC_per_class[i]]]
@@ -116,6 +116,13 @@
     mean_seg_class_ious = np.array(seg_class_ious).mean()
     print(f"----Average IoU {mean_seg_class_ious:.5f} ----")
 
+    print("#### Base Footprint ####")
+    # Print base footprint statistics
+    base_footprint_table = [["Metric", "Value"]]
+    for metric, value in base_footprint_metrics.items():
+        base_footprint_table += [[metric, "%.5f" % value]]
+    print(AsciiTable(base_footprint_table).table)
+
 
 def _evaluate(model, dataloader, class_config, img_size, iou_thres, conf_thres, nms_thres, verbose):
     """Evaluate model on validation dataset.
@@ -136,7 +143,7 @@
     :type nms_thres: float
     :param verbose: If True, prints stats of model
     :type verbose: bool
-    :return: Returns precision, recall, AP, f1, ap_class
+    :return: Returns
     """
     model.eval()  # Set model to evaluation mode
 
@@ -157,13 +164,13 @@
         # Extract labels
         labels += bb_targets[:, 1].tolist()
 
-        # If a subset of the detection classes should be grouped into one class for non-maximum suppression and the 
+        # If a subset of the detection classes should be grouped into one class for non-maximum suppression and the
         # subsequent AP-computation, we need to group those class labels here.
         if class_config.classes_should_be_grouped():
             labels = class_config.group(labels)
 
         # Rescale target
-        bb_targets[:, 2:] = xywh2xyxy(bb_targets[:, 2:])
+        bb_targets[:, 2:-2] = xywh2xyxy(bb_targets[:, 2:-2])
         bb_targets[:, 2:] *= img_size
 
         imgs = Variable(imgs.type(Tensor), requires_grad=False)
@@ -180,9 +187,9 @@
             )
 
         sample_stat, secondary_stat = get_batch_statistics(
-            yolo_outputs, 
-            bb_targets, 
-            iou_threshold=iou_thres, 
+            yolo_outputs,
+            bb_targets,
+            iou_threshold=iou_thres,
             group_config=class_config.get_group_config()
         )
 
@@ -200,12 +207,28 @@
     print(f"Times: Mean {1 / np.array(times).mean()}fps | Std: {np.array(times).std()} ms")
 
     # Concatenate sample statistics
-    true_positives, pred_scores, pred_labels = [
+    true_positives, pred_scores, pred_labels, base_bootprint_visibility_tp, base_bootprint_visibility_fn, base_bootprint_visibility_fp, base_bootprint_visibility_tn, base_footprint_point_distances = [
         np.concatenate(x, 0) for x in list(zip(*sample_metrics))]
-
+
+    # Calculate base footprint metrics
+    base_bootprint_visibility_fn = base_bootprint_visibility_fn.sum()
+    base_bootprint_visibility_fp = base_bootprint_visibility_fp.sum()
+    base_bootprint_visibility_tp = base_bootprint_visibility_tp.sum()
+    base_bootprint_visibility_tn = base_bootprint_visibility_tn.sum()
+    base_footprint_visibility_precision = base_bootprint_visibility_tp / (base_bootprint_visibility_tp + base_bootprint_visibility_fp + np.finfo(float).eps)
+    base_footprint_visibility_recall = base_bootprint_visibility_tp / (base_bootprint_visibility_tp + base_bootprint_visibility_fn + np.finfo(float).eps)
+    base_footprint_metrics = {
+        "precision": base_footprint_visibility_precision,
+        "recall": base_footprint_visibility_recall,
+        "f1": 2 * base_footprint_visibility_precision * base_footprint_visibility_recall / (base_footprint_visibility_precision + base_footprint_visibility_recall + np.finfo(float).eps),
+        'keypoint_distance': base_footprint_point_distances.mean()
+    }
+
+    # Calculate yolo metrics
     yolo_metrics_output = ap_per_class(
         true_positives, pred_scores, pred_labels, labels)
 
+    # Calculate segmentation metrics
     def seg_iou_mean_without_nan(seg_iou: List[float]) -> np.ndarray:
         """This helper function is needed to remove cases, where the segmentation IOU is NaN.
         This is the case, if a whole batch does not contain any pixels of a segmentation class.
@@ -218,9 +241,9 @@
 
     seg_class_ious = [seg_iou_mean_without_nan(class_ious) for class_ious in list(zip(*seg_ious))]
 
-    print_eval_stats(yolo_metrics_output, seg_class_ious, secondary_metric, class_config, verbose)
+    print_eval_stats(yolo_metrics_output, seg_class_ious, secondary_metric, base_footprint_metrics, class_config, verbose)
 
-    return yolo_metrics_output, seg_class_ious, secondary_metric
+    return yolo_metrics_output, seg_class_ious, secondary_metric, base_footprint_metrics
 
 
 def _create_validation_data_loader(img_path, batch_size, img_size, n_cpu):