Update app.py

app.py

@@ -1,22 +1,25 @@
 import gradio as gr
 import cv2
 import numpy as np
-import torch
 from ultralytics import YOLO
+from sort import Sort  # SORT tracker
 
 # Load YOLOv8 model (pre-trained on COCO dataset)
-model = YOLO("yolov8x.pt")  # Use "yolov8x.pt" for highest accuracy
+model = YOLO("yolov8x.pt")  # Highest accuracy version
 
 # Class label for trucks (COCO dataset)
 TRUCK_CLASS_ID = 7  # "truck" in COCO dataset
 
+# Initialize SORT tracker
+tracker = Sort()  
+
 def count_trucks(video_path):
     cap = cv2.VideoCapture(video_path)
     if not cap.isOpened():
         return "Error: Unable to open video file."
 
     frame_count = 0
-    truck_count_per_frame = []
+    unique_truck_ids = set()
     frame_skip = 5  # Process every 5th frame for efficiency
 
     while True:
@@ -26,27 +29,39 @@ def count_trucks(video_path):
 
         frame_count += 1
         if frame_count % frame_skip != 0:
-            continue  # Skip frames to improve efficiency
+            continue  # Skip frames for efficiency
 
         # Run YOLOv8 inference
         results = model(frame, verbose=False)
 
-        truck_count = 0
+        detections = []
         for result in results:
             for box in result.boxes:
                 class_id = int(box.cls.item())  # Get class ID
                 confidence = float(box.conf.item())  # Get confidence score
+                x1, y1, x2, y2 = map(int, box.xyxy[0])  # Get bounding box
 
                 if class_id == TRUCK_CLASS_ID and confidence > 0.6:
-                    truck_count += 1  # Count trucks with confidence > 0.6
+                    detections.append([x1, y1, x2, y2, confidence])  # Append detection
+
+        # Convert to numpy array for SORT input
+        if len(detections) > 0:
+            detections = np.array(detections)
+        else:
+            detections = np.empty((0, 5))  # Empty array when no trucks detected
+
+        # Update tracker
+        tracked_objects = tracker.update(detections)
 
-        truck_count_per_frame.append(truck_count)
+        # Store unique truck IDs
+        for obj in tracked_objects:
+            truck_id = int(obj[4])  # SORT assigns unique IDs
+            unique_truck_ids.add(truck_id)
 
     cap.release()
 
     return {
-        "Total Trucks in Video": int(np.max(truck_count_per_frame)) if truck_count_per_frame else 0 #,
-        #"Avg Trucks Per Frame": round(np.mean(truck_count_per_frame), 2) if truck_count_per_frame else 0
+        "Total Unique Trucks in Video": len(unique_truck_ids)
     }
 
 # Gradio UI function
@@ -59,8 +74,8 @@ interface = gr.Interface(
     fn=analyze_video,
     inputs=gr.Video(label="Upload Video"),
     outputs=gr.Textbox(label="Truck Counting Results"),
-    title="YOLOv8-based Truck Counter",
-    description="Upload a video to detect and count trucks using YOLOv8."
+    title="YOLOv8-based Truck Counter with Object Tracking",
+    description="Upload a video to detect and count unique trucks using YOLOv8 and SORT tracker."
 )
 
 # Launch app