detectors: implement class aggregation

frigate/detectors/plugins/edgetpu_tfl.py

@@ -57,6 +57,8 @@ def __init__(self, detector_config: EdgeTpuDetectorConfig):
         self.tensor_output_details = self.interpreter.get_output_details()
         self.model_type = detector_config.model.model_type
 
+        self.class_aggregation = yolo_utils.generate_class_aggregation_from_config(detector_config)
+
     def detect_raw(self, tensor_input):
         if self.model_type == 'yolov8':
             scale, zero_point = self.tensor_input_details[0]['quantization']
@@ -72,7 +74,7 @@ def detect_raw(self, tensor_input):
             model_input_shape = self.tensor_input_details[0]['shape']
             tensor_output[:, [0, 2]] *= model_input_shape[2]
             tensor_output[:, [1, 3]] *= model_input_shape[1]
-            return yolo_utils.yolov8_postprocess(model_input_shape, tensor_output)
+            return yolo_utils.yolov8_postprocess(model_input_shape, tensor_output, class_aggregation = self.class_aggregation)
 
         boxes = self.interpreter.tensor(self.tensor_output_details[0]["index"])()[0]
         class_ids = self.interpreter.tensor(self.tensor_output_details[1]["index"])()[0]

frigate/detectors/plugins/onnx.py

@@ -47,6 +47,8 @@ def __init__(self, detector_config: ONNXDetectorConfig):
         self.model = onnxruntime.InferenceSession(path)
         logger.info(f"ONNX: {path} loaded")
 
+        self.class_aggregation = yolo_utils.generate_class_aggregation_from_config(detector_config)
+
     def detect_raw(self, tensor_input):
         model_input_name = self.model.get_inputs()[0].name
         model_input_shape = self.model.get_inputs()[0].shape
@@ -55,5 +57,5 @@ def detect_raw(self, tensor_input):
 
         tensor_output = self.model.run(None, {model_input_name: tensor_input})[0]
 
-        return yolo_utils.yolov8_postprocess(model_input_shape, tensor_output)
+        return yolo_utils.yolov8_postprocess(model_input_shape, tensor_output, class_aggregation = self.class_aggregation)
 

frigate/detectors/plugins/rocm.py

@@ -98,6 +98,8 @@ def __init__(self, detector_config: ROCmDetectorConfig):
             migraphx.save(self.model, mxr_path)
         logger.info(f"AMD/ROCm: model loaded")
 
+        self.class_aggregation = yolo_utils.generate_class_aggregation_from_config(detector_config)
+
     def detect_raw(self, tensor_input):
         model_input_name = self.model.get_parameter_names()[0];
         model_input_shape = tuple(self.model.get_parameter_shapes()[model_input_name].lens());
@@ -109,5 +111,5 @@ def detect_raw(self, tensor_input):
         addr = ctypes.cast(detector_result.data_ptr(), ctypes.POINTER(ctypes.c_float))
         tensor_output = np.ctypeslib.as_array(addr, shape=detector_result.get_shape().lens())
 
-        return yolo_utils.yolov8_postprocess(model_input_shape, tensor_output)
+        return yolo_utils.yolov8_postprocess(model_input_shape, tensor_output, class_aggregation = self.class_aggregation)
 

frigate/detectors/yolo_utils.py

@@ -3,8 +3,34 @@
 import numpy as np
 import cv2
 
+from frigate.util.builtin import load_labels
+
 logger = logging.getLogger(__name__)
 
+def generate_class_aggregation(labels):
+    if isinstance(labels, dict):
+        labels = [labels.get(i, 'unknown') for i in range(0, max(labels.keys()) + 1)]
+    while len(labels) > 0 and labels[-1] in ('unknown', 'other'):
+        labels = labels[:-1]
+    labels = np.array(labels)
+    unique_labels = np.unique(labels)
+    if len(unique_labels) == len(labels):
+        # nothing to aggregate, so there is no mapping
+        return None
+    ret = []
+    for label in unique_labels:
+        if label == 'other' or label == 'unknown':
+            continue
+        index = np.where(labels == label)[0]
+        ret.append(((label, index[0]), index))
+    return ret
+
+def generate_class_aggregation_from_config(config):
+    labelmap_path = config.model.labelmap_path
+    if labelmap_path is None:
+        return None
+    return generate_class_aggregation(load_labels(labelmap_path))
+
 def preprocess(tensor_input, model_input_shape, model_input_element_type):
     model_input_shape = tuple(model_input_shape)
     assert tensor_input.dtype == np.uint8, f'tensor_input.dtype: {tensor_input.dtype}'
@@ -22,14 +48,21 @@ def preprocess(tensor_input, model_input_shape, model_input_element_type):
     # cv2.dnn.blobFromImage is faster than numpying it
     return cv2.dnn.blobFromImage(tensor_input[0], 1.0 / 255, (model_input_shape[3], model_input_shape[2]), None, swapRB=False)
 
-def yolov8_postprocess(model_input_shape, tensor_output, box_count = 20, score_threshold = 0.3, nms_threshold = 0.5):
+def yolov8_postprocess(model_input_shape, tensor_output, box_count = 20, score_threshold = 0.5, nms_threshold = 0.5, class_aggregation = None):
     model_box_count = tensor_output.shape[2]
-    probs = tensor_output[0, 4:, :]
-    all_ids = np.argmax(probs, axis=0)
-    all_confidences = probs.T[np.arange(model_box_count), all_ids]
+    probs = tensor_output[0, 4:, :].T
+    if class_aggregation is not None:
+        new_probs = np.zeros((probs.shape[0], len(class_aggregation)), dtype=probs.dtype)
+        for index, ((label, class_id), selector) in enumerate(class_aggregation):
+            new_probs[:, index] = np.sum(probs[:, selector], axis=1)
+        probs = new_probs
+    all_ids = np.argmax(probs, axis=1)
+    all_confidences = probs[np.arange(model_box_count), all_ids]
     all_boxes = tensor_output[0, 0:4, :].T
     mask = (all_confidences > score_threshold)
     class_ids = all_ids[mask]
+    if class_aggregation is not None:
+        class_ids = np.array([class_aggregation[index][0][1] for index in class_ids])
     confidences = all_confidences[mask]
     cx, cy, w, h = all_boxes[mask].T