activate gpu

app.py

@@ -63,6 +63,8 @@ def fn_video(video, initial_time, duration):
     while cap.isOpened():
         try:
             ret, frame = cap.read()
+            gpu_frame = cv2.cuda_GpuMat()
+            gpu_frame.upload(frame)
             if not ret:
                 break
             frame_copy = frame.copy()
@@ -72,7 +74,7 @@ def fn_video(video, initial_time, duration):
         if num_frames < min_frame:
             num_frames += 1
             continue
-        yolo_detections = detect_plates(frame)
+        yolo_detections = detect_plates(gpu_frame)
         detections = yolo_to_norfair(yolo_detections)
         tracked_objects = tracker.update(detections=detections)
         for obj in tracked_objects:
@@ -80,22 +82,22 @@ def fn_video(video, initial_time, duration):
                 bbox = obj.last_detection.points
                 bbox = int(bbox[0][0]), int(bbox[0][1]), int(bbox[1][0]), int(bbox[1][1])
                 if obj.id not in plates.keys():
-                    crop = imcrop(frame, bbox)
+                    crop = imcrop(gpu_frame, bbox)
                     text = detect_chars(crop)
                     plates[obj.id] = text
                     thread = Thread(target=send_request, args=(frame_copy, text, bbox))
                     thread.start()
 
                 cv2.rectangle(
-                    frame,
+                    gpu_frame,
                     (bbox[0], bbox[1]),
                     (bbox[2], bbox[3]),
                     (0, 255, 0),
                     2,
                 )
-                draw_text(frame, plates[obj.id], (bbox[0], bbox[1]))
+                draw_text(gpu_frame, plates[obj.id], (bbox[0], bbox[1]))
                 cv2.putText(
-                    frame,
+                    gpu_frame,
                     plates[obj.id],
                     (bbox[0], bbox[1]),
                     cv2.FONT_HERSHEY_SIMPLEX,
@@ -103,7 +105,7 @@ def fn_video(video, initial_time, duration):
                     (0, 255, 0),
                     2,
                 )
-        final_video.write(frame)
+        final_video.write(gpu_frame)
         num_frames += 1
         if num_frames == max_frame:
             break

detector/utils.py

@@ -9,14 +9,13 @@ import requests
 BASE_DIR = os.path.abspath(os.getcwd())
 
 device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
-cpu = torch.device('cpu')
-print('Loading models...', cpu)
+print('Loading models...', device)
 
 model_plates = torch.hub.load('ultralytics/yolov5', 'custom',
-                              path=os.path.join(BASE_DIR, 'detector', 'static', 'plates.pt'), device=cpu)
+                              path=os.path.join(BASE_DIR, 'detector', 'static', 'plates.pt'), device=device)
 
 model_chars = torch.hub.load('ultralytics/yolov5', 'custom',
-                             path=os.path.join(BASE_DIR, 'detector', 'static', 'chars.pt'), device=cpu)
+                             path=os.path.join(BASE_DIR, 'detector', 'static', 'chars.pt'), device=device)
 
 
 def pad_img_to_fit_bbox(img, x1, x2, y1, y2):