Update inference_onnx.py

inference_onnx.py

@@ -1,9 +1,8 @@
 """
-ONNX Inference implementation for LSNet models
+ONNX Inference implementation for Kaloscope LSNet model
 """
 
 import numpy as np
-import onnxruntime as ort
 from timm.data import resolve_data_config
 from timm.data.transforms_factory import create_transform
 from timm.models import create_model
@@ -20,6 +19,7 @@ class ONNXInference:
     def __init__(self, model_path, model_arch="lsnet_xl_artist", device="cpu"):
         """
         Initialize ONNX inference session
+        Tries CUDA GPU execution when selected and available, and OpenVINO for CPU with CPUExecutionProvider as last fallback
 
         Args:
             model_path: Path to ONNX model file
@@ -29,18 +29,68 @@ class ONNXInference:
         self.model_path = model_path
         self.model_arch = model_arch
         self.device = device
+        self.use_openvino = False
 
-        # Set providers based on device | barebones, theres a lot more https://onnxruntime.ai/docs/execution-providers/
         if device == "cuda":
-            providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
-        else:
-            providers = ["CPUExecutionProvider"]
+            # Try CUDA first for GPU
+            try:
+                import onnxruntime as ort
+
+                # Set session options to suppress warnings
+                sess_options = ort.SessionOptions()
+                sess_options.log_severity_level = 3  # 0:Verbose, 1:Info, 2:Warning, 3:Error, 4:Fatal
+
+                providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
+                self.session = ort.InferenceSession(model_path, sess_options=sess_options, providers=providers)
+                self.execution_provider = self.session.get_providers()[0]
+
+                # Check if CUDA is actually being used
+                if self.execution_provider == "CUDAExecutionProvider":
+                    print(f"Using ONNX Runtime with {self.execution_provider}")
+                    # Get transform from timm model
+                    self.transform = self._get_transform()
+                    return
+                else:
+                    # CUDA failed, fall through to CPU logic
+                    print("CUDA not available in ONNX Runtime, falling back to CPU options")
+            except Exception as e:
+                print(f"ONNX Runtime CUDA initialization failed: {e}, falling back to CPU options")
+
+        # For CPU or if CUDA failed, prefer OpenVINO
+        try:
+            import openvino as ov
+
+            # error here on purpose
+            # raise ImportError("aaa")
+
+            core = ov.Core()
+            self.model = core.read_model(model_path)
+            self.session = core.compile_model(self.model, "CPU")
+            self.execution_provider = "CPU – OpenVINO™"
+            self.use_openvino = True
+            print("Using OpenVINO runtime for inference on CPU")
+        except ImportError:
+            print("OpenVINO not available, falling back to ONNX Runtime CPU")
+            self._init_onnx_runtime_cpu(model_path)
+        except Exception as e:
+            print(f"OpenVINO initialization failed: {e}, falling back to ONNX Runtime CPU")
+            self._init_onnx_runtime_cpu(model_path)
+
+        # Get transform from timm model
+        self.transform = self._get_transform()
 
-        # Load ONNX session
-        self.session = ort.InferenceSession(model_path, providers=providers)
+    def _init_onnx_runtime_cpu(self, model_path):
+        """Initialize ONNX Runtime with CPU as fallback"""
+        import onnxruntime as ort
 
-        # Store the actual provider being used
+        # Set session options to suppress warnings
+        sess_options = ort.SessionOptions()
+        sess_options.log_severity_level = 3  # Only show errors and fatal messages
+
+        providers = ["CPUExecutionProvider"]
+        self.session = ort.InferenceSession(model_path, sess_options=sess_options, providers=providers)
         self.execution_provider = self.session.get_providers()[0]
+        print(f"Using ONNX Runtime with {self.execution_provider}")
 
         # Get transform from timm model
         self.transform = self._get_transform()
@@ -81,11 +131,16 @@ class ONNXInference:
         """
         input_tensor = self.preprocess(image)
 
-        input_name = self.session.get_inputs()[0].name
-        output_name = self.session.get_outputs()[0].name
-
-        results = self.session.run([output_name], {input_name: input_tensor})
-        logits = results[0][0]
+        if self.use_openvino:
+            # OpenVINO inference
+            results = self.session(input_tensor)
+            logits = list(results.values())[0][0]
+        else:
+            # ONNX Runtime inference
+            input_name = self.session.get_inputs()[0].name
+            output_name = self.session.get_outputs()[0].name
+            results = self.session.run([output_name], {input_name: input_tensor})
+            logits = results[0][0]
 
         return logits