Update detect_strongsort.py

detect_strongsort.py

@@ -15,6 +15,7 @@ import torch
 import torch.backends.cudnn as cudnn
 from numpy import random
 from time import time
+import pandas as pd
 
 
 FILE = Path(__file__).resolve()
@@ -44,7 +45,6 @@ VID_FORMATS = 'asf', 'avi', 'gif', 'm4v', 'mkv', 'mov', 'mp4', 'mpeg', 'mpg', 't
 def plot_one_box(x, img, color=None, label=None, line_thickness=3):
     # Plots one bounding box on image img
     tl = line_thickness or round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1  # line/font thickness
-    #color = color or [random.randint(0, 255) for _ in range(3)]
     c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))
     cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)
     if label:
@@ -55,6 +55,10 @@ def plot_one_box(x, img, color=None, label=None, line_thickness=3):
         cv2.putText(img, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)
 
 
+
+def convert_to_int(tensor):
+    return tensor.type(torch.int16).item()
+
 @smart_inference_mode()
 def run_strongsort(
         source='0',
@@ -163,6 +167,8 @@ def run_strongsort(
     model.warmup(imgsz=(1 if pt or model.triton else bs, 3, *imgsz))  # warmup
     seen, windows, dt,sdt = 0, [], (Profile(), Profile(), Profile(), Profile()),[0.0, 0.0, 0.0, 0.0]
     curr_frames, prev_frames = [None] * bs, [None] * bs
+    frame_counts = []
+
     for frame_idx, (path, im, im0s, vid_cap, s) in enumerate(dataset):
         # s = ''
         t1 = time_sync()
@@ -179,20 +185,19 @@ def run_strongsort(
         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]
         t3 = time_sync()
         sdt[1] += t3 - t2
 
         # Apply 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)
         sdt[2] += time_sync() - t3
         
         # Second-stage classifier (optional)
         # pred = utils.general.apply_classifier(pred, classifier_model, im, im0s)
 
-
+        counts = {}
         # Process detections
         for i, det in enumerate(pred):  # detections per image
             seen += 1
@@ -227,6 +232,7 @@ def run_strongsort(
             imc = im0.copy() if save_crop else im0  # for save_crop
             annotator = Annotator(im0, line_width=line_thickness, example=str(names))
 
+
             if cfg.STRONGSORT.ECC:  # camera motion compensation
                 strongsort_list[i].tracker.camera_update(prev_frames[i], curr_frames[i])
 
@@ -238,7 +244,7 @@ def run_strongsort(
                 for c in det[:, -1].unique():
                     n = (det[:, -1] == c).sum()  # detections per class
                     s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string
-
+                    counts[names[int(c)]] = n
                 xywhs = xyxy2xywh(det[:, 0:4])
                 confs = det[:, 4]
                 clss = det[:, 5]
@@ -248,12 +254,13 @@ def run_strongsort(
                 outputs[i] = strongsort_list[i].update(xywhs.cpu(), confs.cpu(), clss.cpu(), im0)
                 t5 = time_sync()
                 sdt[3] += t5 - t4
-
+            
                 # Write results
                 for j, (output, conf) in enumerate(zip(outputs[i], confs)):
                     xyxy = output[0:4]
                     id = output[4]
                     cls = output[5]
+                    label = names[int(cls)]
                 # for *xyxy, conf, cls in reversed(det):
                     if save_txt:  # Write to file
                         xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh
@@ -269,7 +276,7 @@ def run_strongsort(
                     if save_crop:
                         save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True)
 
-
+                frame_counts.append({'frame': frame_idx, 'counts': counts.copy()})
                 # # draw boxes for visualization
                 # if len(outputs[i]) > 0:
                 #     for j, (output, conf) in enumerate(zip(outputs[i], confs)):
@@ -305,6 +312,8 @@ def run_strongsort(
 
             # Stream results
             im0 = annotator.result()
+
+
             if view_img:
                 if platform.system() == 'Linux' and p not in windows:
                     windows.append(p)
@@ -334,8 +343,17 @@ def run_strongsort(
 
             prev_frames[i] = curr_frames[i]
 
-    # Print time (inference-only)
-    LOGGER.info(f"{s}{'' if len(det) else '(no detections), '}{dt[1].dt * 1E3:.1f}ms")
+            
+        # Print time (inference-only)
+        LOGGER.info(f"{s}{'' if len(det) else '(no detections), '}{dt[1].dt * 1E3:.1f}ms")
+    
+    flattened_counts = [
+        {'frame': entry['frame'], 'label': label, 'count': count}
+        for entry in frame_counts for label, count in entry['counts'].items()
+    ]
+    frame_counts_df = pd.DataFrame(flattened_counts)
+    frame_counts_df['count'] = frame_counts_df['count'].apply(convert_to_int)
+    counts_df = None
     # Print results
     LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape, %.1fms StrongSORT' % tuple(1E3 * x / seen for x in sdt))
     if save_txt or save_img:
@@ -343,7 +361,7 @@ def run_strongsort(
         LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")
     if update:
         strip_optimizer(yolo_weights[0])  # update model (to fix SourceChangeWarning)
-    return save_path
+    return save_path, counts_df, frame_counts_df
 def parse_opt():
     parser = argparse.ArgumentParser()
     parser.add_argument('--yolo-weights', nargs='+', type=str, default=WEIGHTS / 'yolov9.pt', help='model.pt path(s)')