Update detect_deepsort.py

detect_deepsort.py

@@ -6,8 +6,10 @@ from pathlib import Path
 import math
 import torch
 import numpy as np
+import re
 from deep_sort_pytorch.utils.parser import get_config
 from deep_sort_pytorch.deep_sort import DeepSort
+import pandas as pd
 from collections import deque
 FILE = Path(__file__).resolve()
 ROOT = FILE.parents[0]  # YOLO root directory
@@ -22,50 +24,59 @@ 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
 
-# def initialize_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,
-#                         # min_confidence  parameter sets the minimum tracking confidence required for an object detection to be considered in the tracking process
-#                         min_confidence=cfg_deep.DEEPSORT.MIN_CONFIDENCE,
-#                         #nms_max_overlap specifies the maximum allowed overlap between bounding boxes during non-maximum suppression (NMS)
-#                         nms_max_overlap=cfg_deep.DEEPSORT.NMS_MAX_OVERLAP,
-#                         #max_iou_distance parameter defines the maximum intersection-over-union (IoU) distance between object detections
-#                         max_iou_distance=cfg_deep.DEEPSORT.MAX_IOU_DISTANCE,
-#                         # Max_age: If an object's tracking ID is lost (i.e., the object is no longer detected), this parameter determines how many frames the tracker should wait before assigning a new id
-#                         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=False
-#         )
-
-#     return deepsort
-
-#deepsort = initialize_deepsort()
+def initialize_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,
+                        # min_confidence  parameter sets the minimum tracking confidence required for an object detection to be considered in the tracking process
+                        min_confidence=cfg_deep.DEEPSORT.MIN_CONFIDENCE,
+                        #nms_max_overlap specifies the maximum allowed overlap between bounding boxes during non-maximum suppression (NMS)
+                        nms_max_overlap=cfg_deep.DEEPSORT.NMS_MAX_OVERLAP,
+                        #max_iou_distance parameter defines the maximum intersection-over-union (IoU) distance between object detections
+                        max_iou_distance=cfg_deep.DEEPSORT.MAX_IOU_DISTANCE,
+                        # Max_age: If an object's tracking ID is lost (i.e., the object is no longer detected), this parameter determines how many frames the tracker should wait before assigning a new id
+                        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=False
+        )
+
+    return deepsort
+
+deepsort = initialize_deepsort()
 data_deque = {}
 def classNames():
     cocoClassNames = ["Bus", "Bike", "Car", "Pedestrian", "Truck"
                   ]
     return cocoClassNames
 className = classNames()
-
+# def convert_to_int(x):
+#     if isinstance(x, str):
+#         # Extract numeric value from tensor string using regular expressions
+#         match = re.match(r'tensor\((\d+)\)', x)
+#         if match:
+#             return int(match.group(1))
+#     return x
 def colorLabels(classid):
     if classid == 0: #Bus
         color = (0, 0, 255)
     elif classid == 1: #Bike  250, 247, 0
-        color = (250, 247, 0)
+        color = (0,148,255)
     elif classid == 2: #Car 
         color = (0, 255, 10)
     elif classid == 3: #Pedestrian
-        color = (0,148,255) 
+        color = (250,247,0) 
     else: #Truck
         color = (235,0,255)   
     return tuple(color)
 
+def convert_to_int(tensor):
+    return tensor.type(torch.int16).item()
+
 def draw_boxes(frame, bbox_xyxy, draw_trails, identities=None, categories=None, offset=(0,0)):
     height, width, _ = frame.shape
     for key in list(data_deque):
@@ -146,27 +157,6 @@ 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
 
-    #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,
-                        # min_confidence  parameter sets the minimum tracking confidence required for an object detection to be considered in the tracking process
-                        min_confidence=cfg_deep.DEEPSORT.MIN_CONFIDENCE,
-                        #nms_max_overlap specifies the maximum allowed overlap between bounding boxes during non-maximum suppression (NMS)
-                        nms_max_overlap=cfg_deep.DEEPSORT.NMS_MAX_OVERLAP,
-                        #max_iou_distance parameter defines the maximum intersection-over-union (IoU) distance between object detections
-                        max_iou_distance=cfg_deep.DEEPSORT.MAX_IOU_DISTANCE,
-                        # Max_age: If an object's tracking ID is lost (i.e., the object is no longer detected), this parameter determines how many frames the tracker should wait before assigning a new id
-                        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
-        )
-    
     # Load model
     device = select_device(device)
     model = DetectMultiBackend(weights, device=device, dnn=dnn, data=data, fp16=half)
@@ -188,6 +178,7 @@ def run_deepsort(
     # Run inference
     model.warmup(imgsz=(1 if pt or model.triton else bs, 3, *imgsz))  # warmup
     seen, windows, dt = 0, [], (Profile(), Profile(), Profile())
+    frame_counts = []  
     for path, im, im0s, vid_cap, s in dataset:
         with dt[0]:
             im = torch.from_numpy(im).to(model.device)
@@ -200,16 +191,15 @@ 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][1]
+            pred = pred[0][1]
 
         # NMS
         with dt[2]:
-            pred = pred[0][1] if isinstance(pred[0], list) else pred[0]  # single model or ensemble
             pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
 
         # Second-stage classifier (optional)
         # pred = utils.general.apply_classifier(pred, classifier_model, im, im0s)
-
+        counts = {}
         # Process predictions
         for i, det in enumerate(pred):  # per image
             seen += 1
@@ -233,6 +223,7 @@ def run_deepsort(
                 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
+                    counts[names[int(c)]] = n
                 xywh_bboxs = []
                 confs = []
                 oids = []
@@ -253,7 +244,7 @@ def run_deepsort(
                     classNameInt = int(cls)
                     oids.append(classNameInt)
                 xywhs = torch.tensor(xywh_bboxs)
-                confss = torch.tensor(confs)
+                confss = torch.tensor(confs)    
                 outputs = deepsort.update(xywhs, confss, oids, ims)
                 if len(outputs) > 0:
                     bbox_xyxy = outputs[:, :4]
@@ -287,9 +278,28 @@ def run_deepsort(
 
         # Print time (inference-only)
         LOGGER.info(f"{s}{'' if len(det) else '(no detections), '}{dt[1].dt * 1E3:.1f}ms")
+        frame_counts.append((frame, counts))  # Append the counts for each frame
+    transformed_data = []
+
+    # Iterate over frame_counts and transform each entry into a row in the DataFrame
+    for frame, counts_dict in frame_counts:
+        for label, count in counts_dict.items():
+            transformed_data.append((frame, label.capitalize(), count))
+
+    # Create a DataFrame from the transformed data
+    df = pd.DataFrame(transformed_data, columns=['frame', 'label', 'count'])
+
+    # Convert count column from tensors to integers
+    df['count'] = df['count'].apply(convert_to_int)
+
+    counts_df = pd.DataFrame(counts.items(), columns=['label', 'count'])
+    counts_df['count'] = counts_df['count'].apply(convert_to_int)
+    counts_df['label'] = counts_df['label'].astype(str)  
+
     if update:
         strip_optimizer(weights[0])  # update model (to fix SourceChangeWarning)
-    return save_path 
+    return save_path, counts_df, df 
+
 
 def parse_opt():
     parser = argparse.ArgumentParser()