Add support for yolo generic

This commit is contained in:
Nicolas Mowen 2025-03-23 13:36:36 -06:00
parent b25de0bd90
commit a4950da1fe
2 changed files with 9 additions and 2 deletions

View File

@ -97,7 +97,10 @@ class ONNXDetector(DetectionApi):
x_max / self.w,
]
return detections
elif self.onnx_model_type == ModelTypeEnum.yolov9:
elif (
self.onnx_model_type == ModelTypeEnum.yolov9
or self.onnx_model_type == ModelTypeEnum.yologeneric
):
predictions: np.ndarray = tensor_output[0]
return post_process_yolov9(predictions, self.w, self.h)
else:

View File

@ -34,6 +34,7 @@ class OvDetector(DetectionApi):
ModelTypeEnum.ssd,
ModelTypeEnum.yolonas,
ModelTypeEnum.yolov9,
ModelTypeEnum.yologeneric,
ModelTypeEnum.yolox,
]
@ -231,7 +232,10 @@ class OvDetector(DetectionApi):
x_max / self.w,
]
return detections
elif self.ov_model_type == ModelTypeEnum.yolov9:
elif (
self.ov_model_type == ModelTypeEnum.yolov9
or self.ov_model_type == ModelTypeEnum.yologeneric
):
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: