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