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
This commit is contained in:
Nicolas Mowen 2025-02-10 14:00:12 -07:00 committed by GitHub
parent 72209986b6
commit 198d067e25
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 73 additions and 3 deletions

View File

@ -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)

View File

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

View File

@ -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."

View File

@ -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

View File

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