diff --git a/frigate/detectors/plugins/edgetpu_tfl.py b/frigate/detectors/plugins/edgetpu_tfl.py index 07dfc127d9..4cd87e7102 100644 --- a/frigate/detectors/plugins/edgetpu_tfl.py +++ b/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] diff --git a/frigate/detectors/plugins/onnx.py b/frigate/detectors/plugins/onnx.py index 428b680782..d49263bdd3 100644 --- a/frigate/detectors/plugins/onnx.py +++ b/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) diff --git a/frigate/detectors/plugins/rocm.py b/frigate/detectors/plugins/rocm.py index d5d0ba585a..6638a5b674 100644 --- a/frigate/detectors/plugins/rocm.py +++ b/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) diff --git a/frigate/detectors/yolo_utils.py b/frigate/detectors/yolo_utils.py index 02442ab659..51bd3c839b 100644 --- a/frigate/detectors/yolo_utils.py +++ b/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