diff --git a/src/pipeline.py b/src/pipeline.py index 2bbb75e..497321f 100644 --- a/src/pipeline.py +++ b/src/pipeline.py @@ -92,6 +92,17 @@ def run_pipeline( effective_filter = class_filter or ( model_config.known_classes if model_config.known_classes else None ) + + # Check onnxruntime for .onnx models + if model_config.path.endswith(".onnx"): + try: + import onnxruntime # noqa: F401 + except ImportError: + raise RuntimeError( + f"onnxruntime is not installed. Cannot load .onnx model '{model_config.filename}'. " + f"Either install it (pip install onnxruntime) or use the .engine version of this model." + ) + shared_model = YOLO(model_config.path) detector = BaseDetector( shared_model, conf=sack_conf, class_filter=effective_filter