Implement support for YOLOv9 via ONNX (#16459)

* WIP yolov9

* Implement post processing for yolov9

* Cleanup detection

* Update docs to make note of supported yolov9

* Move post processing to separate utility

* Add note about other models

docs/docs/configuration/object_detectors.md

@@ -450,7 +450,7 @@ Note that the labelmap uses a subset of the complete COCO label set that has onl
 
 ## ONNX
 
-ONNX is an open format for building machine learning models, Frigate supports running ONNX models on CPU, OpenVINO, and TensorRT. On startup Frigate will automatically try to use a GPU if one is available.
+ONNX is an open format for building machine learning models, Frigate supports running ONNX models on CPU, OpenVINO, ROCm, and TensorRT. On startup Frigate will automatically try to use a GPU if one is available.
 
 :::info
 
@@ -517,6 +517,33 @@ model:
   labelmap_path: /labelmap/coco-80.txt
 ```
 
+#### YOLOv9
+
+[YOLOv9](https://github.com/MultimediaTechLab/YOLO) models are supported, but not included by default.
+
+:::tip
+
+The YOLOv9 detector has been designed to support YOLOv9 models, but may support other YOLO model architectures as well.
+
+:::
+
+After placing the downloaded onnx model in your config folder, you can use the following configuration:
+
+```yaml
+detectors:
+  onnx:
+    type: onnx
+
+model:
+  model_type: yolov9
+  width: 640 # <--- should match the imgsize set during model export
+  height: 640 # <--- should match the imgsize set during model export
+  input_tensor: nchw
+  input_dtype: float
+  path: /config/model_cache/yolov9-t.onnx
+  labelmap_path: /labelmap/coco-80.txt
+```
+
 Note that the labelmap uses a subset of the complete COCO label set that has only 80 objects.
 
 ## CPU Detector (not recommended)

frigate/detectors/detector_config.py

@@ -35,6 +35,7 @@ class InputDTypeEnum(str, Enum):
 class ModelTypeEnum(str, Enum):
     ssd = "ssd"
     yolox = "yolox"
+    yolov9 = "yolov9"
     yolonas = "yolonas"
 
 

frigate/detectors/plugins/onnx.py

@@ -9,7 +9,7 @@ from frigate.detectors.detector_config import (
     BaseDetectorConfig,
     ModelTypeEnum,
 )
-from frigate.util.model import get_ort_providers
+from frigate.util.model import get_ort_providers, post_process_yolov9
 
 logger = logging.getLogger(__name__)
 
@@ -79,6 +79,9 @@ class ONNXDetector(DetectionApi):
                     x_max / self.w,
                 ]
             return detections
+        elif self.onnx_model_type == ModelTypeEnum.yolov9:
+            predictions: np.ndarray = tensor_output[0]
+            return post_process_yolov9(predictions, self.w, self.h)
         else:
             raise Exception(
                 f"{self.onnx_model_type} is currently not supported for rocm. See the docs for more info on supported models."

frigate/util/model.py

@@ -4,6 +4,8 @@ import logging
 import os
 from typing import Any
 
+import cv2
+import numpy as np
 import onnxruntime as ort
 
 try:
@@ -14,6 +16,43 @@ except ImportError:
 
 logger = logging.getLogger(__name__)
 
+### Post Processing
+
+
+def post_process_yolov9(predictions: np.ndarray, width, height) -> np.ndarray:
+    predictions = np.squeeze(predictions).T
+    scores = np.max(predictions[:, 4:], axis=1)
+    predictions = predictions[scores > 0.4, :]
+    scores = scores[scores > 0.4]
+    class_ids = np.argmax(predictions[:, 4:], axis=1)
+
+    # Rescale box
+    boxes = predictions[:, :4]
+
+    input_shape = np.array([width, height, width, height])
+    boxes = np.divide(boxes, input_shape, dtype=np.float32)
+    indices = cv2.dnn.NMSBoxes(boxes, scores, score_threshold=0.4, nms_threshold=0.4)
+    detections = np.zeros((20, 6), np.float32)
+    for i, (bbox, confidence, class_id) in enumerate(
+        zip(boxes[indices], scores[indices], class_ids[indices])
+    ):
+        if i == 20:
+            break
+
+        detections[i] = [
+            class_id,
+            confidence,
+            bbox[1] - bbox[3] / 2,
+            bbox[0] - bbox[2] / 2,
+            bbox[1] + bbox[3] / 2,
+            bbox[0] + bbox[2] / 2,
+        ]
+
+    return detections
+
+
+### ONNX Utilities
+
 
 def get_ort_providers(
     force_cpu: bool = False, device: str = "AUTO", requires_fp16: bool = False

frigate/video.py

@@ -481,7 +481,7 @@ def detect(
     detect_config: DetectConfig,
     object_detector,
     frame,
-    model_config,
+    model_config: ModelConfig,
     region,
     objects_to_track,
     object_filters,