Update detect.py

detect.py

@@ -3,7 +3,7 @@ import os
 import platform
 import sys
 from pathlib import Path
-
+import pandas as pd
 import torch
 
 FILE = Path(__file__).resolve()
@@ -20,6 +20,9 @@ from utils.plots import Annotator, colors, save_one_box
 from utils.torch_utils import select_device, smart_inference_mode
 
 
+def convert_to_int(tensor):
+    return tensor.type(torch.int16).item()
+
 @smart_inference_mode()
 def run(
         weights=ROOT / 'yolo.pt',  # model path or triton URL
@@ -84,6 +87,7 @@ def run(
     # 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)
@@ -106,6 +110,7 @@ def run(
          
         # 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
@@ -131,6 +136,7 @@ def run(
                 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
 
                 # Write results
                 for *xyxy, conf, cls in reversed(det):
@@ -146,7 +152,7 @@ def run(
                         annotator.box_label(xyxy, label, color=colors(c, True))
                     if save_crop:
                         save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True)
-
+                    label_name = names[int(cls)]       
             # Stream results
             im0 = annotator.result()
             if view_img:
@@ -177,8 +183,25 @@ def run(
                     vid_writer[i].write(im0)
 
         # 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)  
     #vid_writer.release()
     # Print results
     t = tuple(x.t / seen * 1E3 for x in dt)  # speeds per image
@@ -188,7 +211,7 @@ def run(
         LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")
     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()