setup cuda

classification/inference_onnx.py

@@ -33,8 +33,14 @@ def load_model(model_path: str):
         raise FileNotFoundError(f"Model file not found: {model_path}")
 
     try:
-        # Try CPU provider first
-        providers = ["CPUExecutionProvider"]
+        available_providers = ort.get_available_providers()
+
+        if "CUDAExecutionProvider" in available_providers:
+            providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
+            print("Using CUDA provider")
+        else:
+            providers = ["CPUExecutionProvider"]
+            print("Using CPU provider")
         session = ort.InferenceSession(model_path, providers=providers)
         return session
 

ultrafast/ultrafastLaneDetector.py

@@ -130,7 +130,14 @@ class UltrafastLaneDetector:
         self.initialize_model(model_path)
 
     def initialize_model(self, model_path):
-        providers = ["CPUExecutionProvider"]
+        available_providers = onnxruntime.get_available_providers()
+
+        if 'CUDAExecutionProvider' in available_providers:
+            providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
+            print("Using CUDA provider")
+        else:
+            providers = ['CPUExecutionProvider']
+            print("Using CPU provider")
 
         self.session = onnxruntime.InferenceSession(model_path, providers=providers)