From 265f089987fdb3979ace595b0162b0675528079d Mon Sep 17 00:00:00 2001 From: Nicolas Mowen Date: Mon, 14 Apr 2025 07:54:57 -0600 Subject: [PATCH] Implement blobbed yolov7 post processing and consolidate yolo implementation --- frigate/detectors/detector_config.py | 1 - frigate/detectors/plugins/onnx.py | 9 +-- frigate/detectors/plugins/openvino.py | 16 ++--- frigate/util/model.py | 96 ++++++++++++++++++++++++++- 4 files changed, 106 insertions(+), 16 deletions(-) diff --git a/frigate/detectors/detector_config.py b/frigate/detectors/detector_config.py index ce77384938..2c54d11e5a 100644 --- a/frigate/detectors/detector_config.py +++ b/frigate/detectors/detector_config.py @@ -37,7 +37,6 @@ class ModelTypeEnum(str, Enum): rfdetr = "rfdetr" ssd = "ssd" yolox = "yolox" - yolov9 = "yolov9" yolonas = "yolonas" yologeneric = "yolo-generic" diff --git a/frigate/detectors/plugins/onnx.py b/frigate/detectors/plugins/onnx.py index a10447b489..c3893ea25e 100644 --- a/frigate/detectors/plugins/onnx.py +++ b/frigate/detectors/plugins/onnx.py @@ -13,7 +13,7 @@ from frigate.util.model import ( get_ort_providers, post_process_dfine, post_process_rfdetr, - post_process_yolov9, + post_process_yolo, ) logger = logging.getLogger(__name__) @@ -97,12 +97,9 @@ class ONNXDetector(DetectionApi): x_max / self.w, ] return detections - elif ( - self.onnx_model_type == ModelTypeEnum.yolov9 - or self.onnx_model_type == ModelTypeEnum.yologeneric - ): + elif self.onnx_model_type == ModelTypeEnum.yologeneric: predictions: np.ndarray = tensor_output[0] - return post_process_yolov9(predictions, self.w, self.h) + return post_process_yolo(tensor_output, self.w, self.h) else: raise Exception( f"{self.onnx_model_type} is currently not supported for onnx. See the docs for more info on supported models." diff --git a/frigate/detectors/plugins/openvino.py b/frigate/detectors/plugins/openvino.py index d903527726..9c7ed5248d 100644 --- a/frigate/detectors/plugins/openvino.py +++ b/frigate/detectors/plugins/openvino.py @@ -13,7 +13,7 @@ from frigate.detectors.detector_config import BaseDetectorConfig, ModelTypeEnum from frigate.util.model import ( post_process_dfine, post_process_rfdetr, - post_process_yolov9, + post_process_yolo, ) logger = logging.getLogger(__name__) @@ -33,7 +33,6 @@ class OvDetector(DetectionApi): ModelTypeEnum.rfdetr, ModelTypeEnum.ssd, ModelTypeEnum.yolonas, - ModelTypeEnum.yolov9, ModelTypeEnum.yologeneric, ModelTypeEnum.yolox, ] @@ -232,12 +231,13 @@ class OvDetector(DetectionApi): x_max / self.w, ] return detections - elif ( - self.ov_model_type == ModelTypeEnum.yolov9 - or self.ov_model_type == ModelTypeEnum.yologeneric - ): - out_tensor = infer_request.get_output_tensor(0).data - return post_process_yolov9(out_tensor, self.w, self.h) + elif self.ov_model_type == ModelTypeEnum.yologeneric: + out_tensor = [] + + for item in infer_request.output_tensors: + out_tensor.append(item.data) + + return post_process_yolo(out_tensor, self.w, self.h) elif self.ov_model_type == ModelTypeEnum.yolox: out_tensor = infer_request.get_output_tensor() # [x, y, h, w, box_score, class_no_1, ..., class_no_80], diff --git a/frigate/util/model.py b/frigate/util/model.py index 19b3b1bf5b..a4ff9bd755 100644 --- a/frigate/util/model.py +++ b/frigate/util/model.py @@ -99,7 +99,94 @@ def post_process_rfdetr(tensor_output: list[np.ndarray, np.ndarray]) -> np.ndarr return detections -def post_process_yolov9(predictions: np.ndarray, width, height) -> np.ndarray: +def __post_process_multipart_yolo( + output_list, + width, + height, +): + anchors = [ + [(12, 16), (19, 36), (40, 28)], + [(36, 75), (76, 55), (72, 146)], + [(142, 110), (192, 243), (459, 401)], + ] + + stride_map = {0: 8, 1: 16, 2: 32} + + all_boxes = [] + all_scores = [] + all_class_ids = [] + + for i, output in enumerate(output_list): + bs, _, ny, nx = output.shape + stride = stride_map[i] + anchor_set = anchors[i] + + num_anchors = len(anchor_set) + output = output.reshape(bs, num_anchors, 85, ny, nx) + output = output.transpose(0, 1, 3, 4, 2) + output = output[0] + + for a_idx, (anchor_w, anchor_h) in enumerate(anchor_set): + for y in range(ny): + for x in range(nx): + pred = output[a_idx, y, x] + class_probs = pred[5:] + class_id = np.argmax(class_probs) + class_conf = class_probs[class_id] + conf = class_conf * pred[4] + + if conf < 0.4: + continue + + dx = pred[0] + dy = pred[1] + dw = pred[2] + dh = pred[3] + + bx = ((dx * 2.0 - 0.5) + x) * stride + by = ((dy * 2.0 - 0.5) + y) * stride + bw = ((dw * 2.0) ** 2) * anchor_w + bh = ((dh * 2.0) ** 2) * anchor_h + + x1 = max(0, bx - bw / 2) / width + y1 = max(0, by - bh / 2) / height + x2 = min(width, bx + bw / 2) / width + y2 = min(height, by + bh / 2) / height + + all_boxes.append([x1, y1, x2, y2]) + all_scores.append(conf) + all_class_ids.append(class_id) + + formatted_boxes = [ + [ + int(x1 * width), + int(y1 * height), + int((x2 - x1) * width), + int((y2 - y1) * height), + ] + for x1, y1, x2, y2 in all_boxes + ] + + indices = cv2.dnn.NMSBoxes( + bboxes=formatted_boxes, + scores=all_scores, + score_threshold=0.4, + nms_threshold=0.4, + ) + + results = np.zeros((20, 6), np.float32) + + if len(indices) > 0: + for i, idx in enumerate(indices.flatten()[:20]): + class_id = all_class_ids[idx] + conf = all_scores[idx] + x1, y1, x2, y2 = all_boxes[idx] + results[i] = [class_id, conf, y1, x1, y2, x2] + + return np.array(results, dtype=np.float32) + + +def __post_process_nms_yolo(predictions: np.ndarray, width, height) -> np.ndarray: predictions = np.squeeze(predictions).T scores = np.max(predictions[:, 4:], axis=1) predictions = predictions[scores > 0.4, :] @@ -131,6 +218,13 @@ def post_process_yolov9(predictions: np.ndarray, width, height) -> np.ndarray: return detections +def post_process_yolo(output: list[np.ndarray], width: int, height: int) -> np.ndarray: + if len(output) > 1: + return __post_process_multipart_yolo(output, width, height) + else: + return __post_process_nms_yolo(output[0], width, height) + + ### ONNX Utilities