mirror of
https://github.com/blakeblackshear/frigate.git
synced 2025-07-30 13:48:07 +02:00
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:
parent
72209986b6
commit
198d067e25
@ -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)
|
||||
|
@ -35,6 +35,7 @@ class InputDTypeEnum(str, Enum):
|
||||
class ModelTypeEnum(str, Enum):
|
||||
ssd = "ssd"
|
||||
yolox = "yolox"
|
||||
yolov9 = "yolov9"
|
||||
yolonas = "yolonas"
|
||||
|
||||
|
||||
|
@ -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."
|
||||
|
@ -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
|
||||
|
@ -481,7 +481,7 @@ def detect(
|
||||
detect_config: DetectConfig,
|
||||
object_detector,
|
||||
frame,
|
||||
model_config,
|
||||
model_config: ModelConfig,
|
||||
region,
|
||||
objects_to_track,
|
||||
object_filters,
|
||||
|
Loading…
Reference in New Issue
Block a user