Initial attempt at adding Yolo26 support for OpenVino detectors

This commit is contained in:
quantumrand 2026-01-26 03:32:51 -08:00
parent 0a8f499640
commit f41c55c1f9
2 changed files with 36 additions and 0 deletions

View File

@ -42,6 +42,7 @@ class ModelTypeEnum(str, Enum):
yolox = "yolox"
yolonas = "yolonas"
yologeneric = "yolo-generic"
yolo26 = "yolo26"
class ModelConfig(BaseModel):

View File

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