Update my_model/object_detection.py

my_model/object_detection.py

@@ -30,6 +30,7 @@ class ObjectDetector:
         self.model = None
         self.processor = None
         self.model_name = None
+        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
         
 
     def load_model(self, model_name='detic', pretrained=True, model_version='yolov5s'):
@@ -64,8 +65,8 @@ class ObjectDetector:
 
         try:
             model_path = get_model_path('deformable-detr-detic')
-            self.processor = AutoImageProcessor.from_pretrained(model_path)
-            self.model = AutoModelForObjectDetection.from_pretrained(model_path)
+            self.processor = AutoImageProcessor.from_pretrained(model_path, device_map = self.device)
+            self.model = AutoModelForObjectDetection.from_pretrained(model_path, device_map = self.device)
         except Exception as e:
             print(f"Error loading Detic model: {e}")
             raise
@@ -83,9 +84,9 @@ class ObjectDetector:
         try:
             model_path = get_model_path ('yolov5')
             if model_path and os.path.exists(model_path):
-                self.model = torch.hub.load(model_path, model_version, pretrained=pretrained, source='local')
+                self.model = torch.hub.load(model_path, model_version, pretrained=pretrained, source='local', device_map = self.device)
             else:
-                self.model = torch.hub.load('ultralytics/yolov5', model_version, pretrained=pretrained)
+                self.model = torch.hub.load('ultralytics/yolov5', model_version, pretrained=pretrained, device_map = self.device)
         except Exception as e:
             print(f"Error loading YOLOv5 model: {e}")
             raise