activate cuda

detector/utils.py

@@ -8,14 +8,18 @@ import requests
 
 BASE_DIR = os.path.abspath(os.getcwd())
 
+device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+print('Loading models...', device)
+
 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_plates.to(device).eval()
 
 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)
+model_chars.to(device).eval()
 
 
 def pad_img_to_fit_bbox(img, x1, x2, y1, y2):
@@ -36,12 +40,12 @@ def imcrop(img, bbox):
 
 
 def detect_plates(img):
-    return model_plates(img, device=device)
+    return model_plates(img)
 
 
 def detect_chars(img):
     img = cv2.resize(img, (640, 320))
-    detect = model_chars(img, device=device)
+    detect = model_chars(img)
     records = detect.pandas().xyxy[0].to_dict(orient='records')
     text = ''
     if records: