Only set backend CUDA if prefer GPU

vietanhdev · May 4, 2023 · 431ad85 · 431ad85
1 parent 28b7f53
commit 431ad85
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git a/anylabeling/services/auto_labeling/yolov5.py b/anylabeling/services/auto_labeling/yolov5.py
@@ -6,6 +6,7 @@
 from PyQt5 import QtCore
 from PyQt5.QtCore import QCoreApplication
 
+from anylabeling.app_info import __preferred_device__
 from anylabeling.views.labeling.shape import Shape
 from anylabeling.views.labeling.utils.opencv import qt_img_to_rgb_cv_img
 from .model import Model
@@ -45,8 +46,9 @@ def __init__(self, model_config, on_message) -> None:
             raise Exception(f"Model not found: {model_abs_path}")
 
         self.net = cv2.dnn.readNet(model_abs_path)
-        self.net.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
-        self.net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA)
+        if __preferred_device__ == "GPU":
+            self.net.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
+            self.net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA)
         self.classes = self.config["classes"]
 
     def pre_process(self, input_image, net):

diff --git a/anylabeling/services/auto_labeling/yolov8.py b/anylabeling/services/auto_labeling/yolov8.py
@@ -6,6 +6,7 @@
 from PyQt5 import QtCore
 from PyQt5.QtCore import QCoreApplication
 
+from anylabeling.app_info import __preferred_device__
 from anylabeling.views.labeling.shape import Shape
 from anylabeling.views.labeling.utils.opencv import qt_img_to_rgb_cv_img
 from .model import Model
@@ -45,8 +46,9 @@ def __init__(self, model_config, on_message) -> None:
             raise Exception(f"Model not found: {model_abs_path}")
 
         self.net = cv2.dnn.readNet(model_abs_path)
-        self.net.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
-        self.net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA)
+        if __preferred_device__ == "GPU":
+            self.net.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
+            self.net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA)
         self.classes = self.config["classes"]
 
     def pre_process(self, input_image, net):