Update app.py

app.py

@@ -3,24 +3,27 @@ import numpy as np
 import torch
 import gradio as gr
 from ultralytics import YOLO
+from sort import Sort  # SORT tracker
 
 # Load YOLOv12x model
-MODEL_PATH = "yolov12x.pt"  # Ensure the model is uploaded to the Hugging Face Space
+MODEL_PATH = "yolov12x.pt"
 model = YOLO(MODEL_PATH)
 
-# COCO dataset class IDs
-PERSON_CLASS_ID = 0  # "person"
-TRUCK_CLASS_ID = 7    # "truck"
+# COCO dataset class ID for truck
+TRUCK_CLASS_ID = 7  # "truck"
 
-def count_objects(video_path):
+# Initialize SORT tracker
+tracker = Sort()
+
+def count_unique_trucks(video_path):
     cap = cv2.VideoCapture(video_path)
     if not cap.isOpened():
         return "Error: Unable to open video file."
 
-    frame_count = 0
-    object_counts = {"people": [], "trucks": []}
+    unique_truck_ids = set()
     frame_skip = 5  # Process every 5th frame for efficiency
 
+    frame_count = 0
     while True:
         ret, frame = cap.read()
         if not ret:
@@ -33,31 +36,32 @@ def count_objects(video_path):
         # Run YOLOv12x inference
         results = model(frame, verbose=False)
 
-        people_count, truck_count = 0, 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
 
-                # Count objects based on their class IDs
-                if class_id == PERSON_CLASS_ID and confidence > 0.5:
-                    people_count += 1
-                elif class_id == TRUCK_CLASS_ID and confidence > 0.5:
-                    truck_count += 1
+                # Track only trucks
+                if class_id == TRUCK_CLASS_ID and confidence > 0.5:
+                    x1, y1, x2, y2 = map(int, box.xyxy[0])  # Get bounding box
+                    detections.append([x1, y1, x2, y2, confidence])
+
+        if len(detections) > 0:
+            detections = np.array(detections)
+            tracked_objects = tracker.update(detections)
 
-        object_counts["people"].append(people_count)
-        object_counts["trucks"].append(truck_count)
+            for obj in tracked_objects:
+                truck_id = int(obj[4])  # Unique ID assigned by SORT
+                unique_truck_ids.add(truck_id)
 
     cap.release()
 
-    return {
-        #"Max People in a Frame": int(np.max(object_counts["people"])) if object_counts["people"] else 0,
-        "Truck Count": int(np.max(object_counts["trucks"])) if object_counts["trucks"] else 0
-    }
+    return {"Total Unique Trucks": len(unique_truck_ids)}
 
 # Gradio UI function
 def analyze_video(video_file):
-    result = count_objects(video_file)
+    result = count_unique_trucks(video_file)
     return "\n".join([f"{key}: {value}" for key, value in result.items()])
 
 # Define Gradio interface
@@ -65,10 +69,10 @@ iface = gr.Interface(
     fn=analyze_video,
     inputs=gr.Video(label="Upload Video"),
     outputs=gr.Textbox(label="Analysis Result"),
-    title="YOLOv12x Object Counter",
-    description="Upload a video to count people and trucks using YOLOv12x."
+    title="YOLOv12x Unique Truck Counter",
+    description="Upload a video to count unique trucks using YOLOv12x and SORT tracking."
 )
 
 # Launch the Gradio app
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch()