mirror of
https://github.com/blakeblackshear/frigate.git
synced 2026-04-11 01:27:36 +03:00
Initial attempt at adding Yolo26 support for OpenVino detectors
This commit is contained in:
parent
0a8f499640
commit
f41c55c1f9
@ -42,6 +42,7 @@ class ModelTypeEnum(str, Enum):
|
||||
yolox = "yolox"
|
||||
yolonas = "yolonas"
|
||||
yologeneric = "yolo-generic"
|
||||
yolo26 = "yolo26"
|
||||
|
||||
|
||||
class ModelConfig(BaseModel):
|
||||
|
||||
@ -33,6 +33,7 @@ class OvDetector(DetectionApi):
|
||||
ModelTypeEnum.yolonas,
|
||||
ModelTypeEnum.yologeneric,
|
||||
ModelTypeEnum.yolox,
|
||||
ModelTypeEnum.yolo26,
|
||||
]
|
||||
|
||||
def __init__(self, detector_config: OvDetectorConfig):
|
||||
@ -82,6 +83,27 @@ class OvDetector(DetectionApi):
|
||||
logger.error(f"SSD model output doesn't match. Found {output_shape}.")
|
||||
self.model_invalid = True
|
||||
|
||||
if self.ov_model_type == ModelTypeEnum.yolo26:
|
||||
model_inputs = self.runner.compiled_model.inputs
|
||||
model_outputs = self.runner.compiled_model.outputs
|
||||
|
||||
if len(model_inputs) != 1:
|
||||
logger.error(
|
||||
f"Yolo26 models must only have 1 input. Found {len(model_inputs)}."
|
||||
)
|
||||
self.model_invalid = True
|
||||
if len(model_outputs) != 1:
|
||||
logger.error(
|
||||
f"Yolo26 models must be exported in flat format and only have 1 output. Found {len(model_outputs)}."
|
||||
)
|
||||
self.model_invalid = True
|
||||
output_shape = model_outputs[0].partial_shape
|
||||
if output_shape[-1] != 6:
|
||||
logger.error(
|
||||
f"Yolo26 models must be exported in flat format. Model output doesn't match (1, N, 6). Found {output_shape}."
|
||||
)
|
||||
self.model_invalid = True
|
||||
|
||||
if self.ov_model_type == ModelTypeEnum.yolonas:
|
||||
model_inputs = self.runner.compiled_model.inputs
|
||||
model_outputs = self.runner.compiled_model.outputs
|
||||
@ -193,6 +215,19 @@ class OvDetector(DetectionApi):
|
||||
x_max / self.w,
|
||||
]
|
||||
return detections
|
||||
elif self.ov_model_type == ModelTypeEnum.yolo26:
|
||||
# Output shape (batch, predictions, [x_center, y_center, width, height, confidence_score, class_id])
|
||||
predictions = outputs[0][0]
|
||||
|
||||
for i, prediction in enumerate(predictions):
|
||||
if i == 20:
|
||||
break
|
||||
(x, y, w, h, confidence, class_id) = prediction
|
||||
|
||||
if class_id < 0:
|
||||
continue
|
||||
detections[i] = self.process_yolo(class_id, confidence, [x, y, w, h])
|
||||
return detections
|
||||
elif self.ov_model_type == ModelTypeEnum.yologeneric:
|
||||
return post_process_yolo(outputs, self.w, self.h)
|
||||
elif self.ov_model_type == ModelTypeEnum.yolox:
|
||||
|
||||
Loading…
Reference in New Issue
Block a user