activate cuda

detector/utils.py

@@ -9,11 +9,14 @@ import requests
 BASE_DIR = os.path.abspath(os.getcwd())
 
 model_plates = torch.hub.load('ultralytics/yolov5', 'custom',
-                              path=os.path.join(BASE_DIR, 'detector', 'static', 'plates.pt'))
+                              path=os.path.join(BASE_DIR, 'detector', 'static', 'plates.pt'), )
 
 model_chars = torch.hub.load('ultralytics/yolov5', 'custom',
                              path=os.path.join(BASE_DIR, 'detector', 'static', 'chars.pt'))
 
+device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+print('Loading models...', device)
+
 
 def pad_img_to_fit_bbox(img, x1, x2, y1, y2):
     img = np.pad(img, ((np.abs(np.minimum(0, y1)), np.maximum(y2 - img.shape[0], 0)),
@@ -33,12 +36,12 @@ def imcrop(img, bbox):
 
 
 def detect_plates(img):
-    return model_plates(img)
+    return model_plates(img, device=device)
 
 
 def detect_chars(img):
     img = cv2.resize(img, (640, 320))
-    detect = model_chars(img)
+    detect = model_chars(img, device=device)
     records = detect.pandas().xyxy[0].to_dict(orient='records')
     text = ''
     if records: