Update detect_deepsort.py

detect_deepsort.py

@@ -1,4 +1,3 @@
-#import spaces
 import argparse
 import os
 import platform
@@ -23,7 +22,6 @@ from utils.general import (LOGGER, Profile, check_file, check_img_size, check_im
 from utils.plots import Annotator, colors, save_one_box
 from utils.torch_utils import select_device, smart_inference_mode
 
-#@spaces.GPU(duration=120)
 # def initialize_deepsort():
 #     # Create the Deep SORT configuration object and load settings from the YAML file
 #     cfg_deep = get_config()
@@ -42,12 +40,12 @@ from utils.torch_utils import select_device, smart_inference_mode
 #                         max_age=cfg_deep.DEEPSORT.MAX_AGE, n_init=cfg_deep.DEEPSORT.N_INIT,
 #                         #nn_budget: It sets the budget for the nearest-neighbor search.
 #                         nn_budget=cfg_deep.DEEPSORT.NN_BUDGET,
-#                         use_cuda=True
+#                         use_cuda=False
 #         )
 
 #     return deepsort
 
-# deepsort = initialize_deepsort()
+#deepsort = initialize_deepsort()
 data_deque = {}
 def classNames():
     cocoClassNames = ["Bus", "Bike", "Car", "Pedestrian", "Truck"
@@ -56,16 +54,16 @@ def classNames():
 className = classNames()
 
 def colorLabels(classid):
-    if classid == 0: #person
-        color = (85, 45, 255)
-    elif classid == 1: #car
-        color = (222, 82, 175)
-    elif classid == 2: #Motorbike
-        color = (0, 204, 255)
-    elif classid == 3: #Bus
-        color = (0,149,255)
-    else:
-        color = (200, 100,0)
+    if classid == 0: #Bus
+        color = (0, 0, 255)
+    elif classid == 1: #Bike  250, 247, 0
+        color = (0,148,255)
+    elif classid == 2: #Car 
+        color = (0, 255, 10)
+    elif classid == 3: #Pedestrian
+        color = (250,247,0) 
+    else: #Truck
+        color = (235,0,255)   
     return tuple(color)
 
 def draw_boxes(frame, bbox_xyxy, draw_trails, identities=None, categories=None, offset=(0,0)):
@@ -110,7 +108,6 @@ def draw_boxes(frame, bbox_xyxy, draw_trails, identities=None, categories=None,
                   cv2.line(frame, data_deque[id][i - 1], data_deque[id][i], color, thickness)    
     return frame
 
-#@spaces.GPU(duration=120)
 @smart_inference_mode()
 def run_deepsort(
         weights=ROOT / 'yolo.pt',  # model path or triton URL
@@ -149,17 +146,11 @@ def run_deepsort(
     save_dir = increment_path(Path(project) / name, exist_ok=exist_ok)  # increment run
     save_dir.mkdir(parents=True, exist_ok=True)  # make dir
 
-    # Load model
-    device = select_device(device)
-    model = DetectMultiBackend(weights, device=device, dnn=dnn, data=data, fp16=half)
-    stride, names, pt = model.stride, model.names, model.pt
-    imgsz = check_img_size(imgsz, s=stride)  # check image size
-
-    #Initalize deep sort
+    #Initalize deepsort
     # Create the Deep SORT configuration object and load settings from the YAML file
     cfg_deep = get_config()
     cfg_deep.merge_from_file("deep_sort_pytorch/configs/deep_sort.yaml")
-    
+
     # Initialize the DeepSort tracker
     deepsort = DeepSort(cfg_deep.DEEPSORT.REID_CKPT,
                         max_dist=cfg_deep.DEEPSORT.MAX_DIST,
@@ -173,8 +164,15 @@ def run_deepsort(
                         max_age=cfg_deep.DEEPSORT.MAX_AGE, n_init=cfg_deep.DEEPSORT.N_INIT,
                         #nn_budget: It sets the budget for the nearest-neighbor search.
                         nn_budget=cfg_deep.DEEPSORT.NN_BUDGET,
-                        use_cuda=True
+                        use_cuda=False
         )
+    
+    # Load model
+    device = select_device(device)
+    model = DetectMultiBackend(weights, device=device, dnn=dnn, data=data, fp16=half)
+    stride, names, pt = model.stride, model.names, model.pt
+    imgsz = check_img_size(imgsz, s=stride)  # check image size
+
     # Dataloader
     bs = 1  # batch_size
     if webcam:
@@ -202,7 +200,7 @@ def run_deepsort(
         with dt[1]:
             visualize = increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False
             pred = model(im, augment=augment, visualize=visualize)
-            pred = pred[0][0]
+            pred = pred[0][1]
 
         # NMS
         with dt[2]:
@@ -216,14 +214,14 @@ def run_deepsort(
             seen += 1
             if webcam:  # batch_size >= 1
                 p, im0, frame = path[i], im0s[i].copy(), dataset.count
-                #s += f'{i}: '
+                s += f'{i}: '
             else:
                 p, im0, frame = path, im0s.copy(), getattr(dataset, 'frame', 0)
 
             p = Path(p)  # to Path
             save_path = str(save_dir / p.name)  # im.jpg
             txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}')  # im.txt
-            #s += '%gx%g ' % im.shape[2:]  # print string
+            s += '%gx%g ' % im.shape[2:]  # print string
             gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh
             ims = im0.copy()
             if len(det):
@@ -233,7 +231,7 @@ def run_deepsort(
                 # Print results
                 for c in det[:, 5].unique():
                     n = (det[:, 5] == c).sum()  # detections per class
-                    #s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string
+                    s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string
                 xywh_bboxs = []
                 confs = []
                 oids = []
@@ -271,23 +269,23 @@ def run_deepsort(
                 cv2.imshow(str(p), ims)
                 cv2.waitKey(1)  # 1 millisecond
             # Save results (image with detections)
-            #if save_img:
-            #if vid_path[i] != save_path:  # new video
-            #vid_path[i] = save_path
-            if isinstance(vid_writer[i], cv2.VideoWriter):
-                vid_writer[i].release()  # release previous video writer
-            if vid_cap:  # video
-                fps = vid_cap.get(cv2.CAP_PROP_FPS)
-                w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-                h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-            else:  # stream
-                fps, w, h = 30, ims.shape[1], ims.shape[0]
-            save_path = str(Path(save_path).with_suffix('.mp4'))  # force *.mp4 suffix on results videos
-            vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
-            vid_writer[i].write(ims)
+            if save_img:
+                if vid_path[i] != save_path:  # new video
+                    vid_path[i] = save_path
+                    if isinstance(vid_writer[i], cv2.VideoWriter):
+                        vid_writer[i].release()  # release previous video writer
+                    if vid_cap:  # video
+                        fps = vid_cap.get(cv2.CAP_PROP_FPS)
+                        w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+                        h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+                    else:  # stream
+                        fps, w, h = 30, ims.shape[1], ims.shape[0]
+                    save_path = str(Path(save_path).with_suffix('.mp4'))  # force *.mp4 suffix on results videos
+                    vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc('m','p','4','v'), fps, (w, h))
+                vid_writer[i].write(ims)
 
         # Print time (inference-only)
-        #LOGGER.info(f"{s}{'' if len(det) else '(no detections), '}{dt[1].dt * 1E3:.1f}ms")
+        LOGGER.info(f"{s}{'' if len(det) else '(no detections), '}{dt[1].dt * 1E3:.1f}ms")
     if update:
         strip_optimizer(weights[0])  # update model (to fix SourceChangeWarning)
     return save_path