administrator/frigate
1
0
Fork 0
mirror of https://github.com/blakeblackshear/frigate.git synced 2026-02-19 01:17:06 +03:00
Code Issues Actions 2 Packages Projects Releases Wiki Activity

Add support for yolov9 to OpenVINO

Browse Source
This commit is contained in:
Nicolas Mowen 2025-02-11 10:53:47 -07:00
parent a3ede3cf8a
commit cf50f64fbf
1 changed files with 13 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
frigate/detectors/plugins/openvino.py
View File

@ -9,6 +9,7 @@ from typing_extensions import Literal
from frigate.detectors.detection_api import DetectionApi from frigate.detectors.detection_api import DetectionApi
from frigate.detectors.detector_config import BaseDetectorConfig, ModelTypeEnum from frigate.detectors.detector_config import BaseDetectorConfig, ModelTypeEnum
from frigate.util.model import post_process_yolov9
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@ -22,7 +23,12 @@ class OvDetectorConfig(BaseDetectorConfig):
class OvDetector(DetectionApi): class OvDetector(DetectionApi):
type_key = DETECTOR_KEY type_key = DETECTOR_KEY
supported_models = [ModelTypeEnum.ssd, ModelTypeEnum.yolonas, ModelTypeEnum.yolox] supported_models = [
ModelTypeEnum.ssd,
ModelTypeEnum.yolonas,
ModelTypeEnum.yolov9,
ModelTypeEnum.yolox,
]
def __init__(self, detector_config: OvDetectorConfig): def __init__(self, detector_config: OvDetectorConfig):
self.ov_core = ov.Core() self.ov_core = ov.Core()
@ -160,8 +166,7 @@ class OvDetector(DetectionApi):
if self.model_invalid: if self.model_invalid:
return detections return detections
elif self.ov_model_type == ModelTypeEnum.ssd:
if self.ov_model_type == ModelTypeEnum.ssd:
results = infer_request.get_output_tensor(0).data[0][0] results = infer_request.get_output_tensor(0).data[0][0]
for i, (_, class_id, score, xmin, ymin, xmax, ymax) in enumerate(results): for i, (_, class_id, score, xmin, ymin, xmax, ymax) in enumerate(results):
@ -176,8 +181,7 @@ class OvDetector(DetectionApi):
xmax, xmax,
] ]
return detections return detections
elif self.ov_model_type == ModelTypeEnum.yolonas:
if self.ov_model_type == ModelTypeEnum.yolonas:
predictions = infer_request.get_output_tensor(0).data predictions = infer_request.get_output_tensor(0).data
for i, prediction in enumerate(predictions): for i, prediction in enumerate(predictions):
@ -196,8 +200,10 @@ class OvDetector(DetectionApi):
x_max / self.w, x_max / self.w,
] ]
return detections return detections
elif self.ov_model_type == ModelTypeEnum.yolov9:
if self.ov_model_type == ModelTypeEnum.yolox: 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.yolox:
out_tensor = infer_request.get_output_tensor() out_tensor = infer_request.get_output_tensor()
# [x, y, h, w, box_score, class_no_1, ..., class_no_80], # [x, y, h, w, box_score, class_no_1, ..., class_no_80],
results = out_tensor.data results = out_tensor.data