diff --git a/docs/docs/configuration/object_detectors.md b/docs/docs/configuration/object_detectors.md index 8510e0602..c04b92474 100644 --- a/docs/docs/configuration/object_detectors.md +++ b/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) diff --git a/frigate/detectors/detector_config.py b/frigate/detectors/detector_config.py index 1d096a648..c8aea0a1d 100644 --- a/frigate/detectors/detector_config.py +++ b/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" diff --git a/frigate/detectors/plugins/onnx.py b/frigate/detectors/plugins/onnx.py index 7004f28fa..c8589145a 100644 --- a/frigate/detectors/plugins/onnx.py +++ b/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." diff --git a/frigate/util/model.py b/frigate/util/model.py index ce2c9538c..75b545cfb 100644 --- a/frigate/util/model.py +++ b/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 diff --git a/frigate/video.py b/frigate/video.py index 3632a87e9..cb922500d 100755 --- a/frigate/video.py +++ b/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,