Update app.py

app.py

@@ -1,31 +1,22 @@
-import matplotlib
-matplotlib.use('Agg')  # Use the 'Agg' backend for headless environments
-
 import gradio as gr
 import cv2
 import numpy as np
+import torch
 from ultralytics import YOLO
-from sort import Sort  # SORT tracker
-
-
-#from sort import Sort  # SORT tracker
 
 # Load YOLOv8 model (pre-trained on COCO dataset)
-model = YOLO("yolov8x.pt")  # Highest accuracy version
+model = YOLO("yolov8x.pt")  # Use "yolov8x.pt" for highest accuracy
 
 # 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
-    unique_truck_ids = set()
+    truck_count_per_frame = []
     frame_skip = 5  # Process every 5th frame for efficiency
 
     while True:
@@ -35,39 +26,27 @@ def count_trucks(video_path):
 
         frame_count += 1
         if frame_count % frame_skip != 0:
-            continue  # Skip frames for efficiency
+            continue  # Skip frames to improve efficiency
 
         # Run YOLOv8 inference
         results = model(frame, verbose=False)
 
-        detections = []
+        truck_count = 0
         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:
-                    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 += 1  # Count trucks with confidence > 0.6
 
-        # Store unique truck IDs
-        for obj in tracked_objects:
-            truck_id = int(obj[4])  # SORT assigns unique IDs
-            unique_truck_ids.add(truck_id)
+        truck_count_per_frame.append(truck_count)
 
     cap.release()
 
     return {
-        "Total Unique Trucks in Video": len(unique_truck_ids)
+        "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
     }
 
 # Gradio UI function
@@ -80,9 +59,9 @@ interface = gr.Interface(
     fn=analyze_video,
     inputs=gr.Video(label="Upload Video"),
     outputs=gr.Textbox(label="Truck Counting Results"),
-    title="YOLOv8-based Truck Counter with Object Tracking",
-    description="Upload a video to detect and count unique trucks using YOLOv8 and SORT tracker."
+    title="YOLOv8-based Truck Counter",
+    description="Upload a video to detect and count trucks using YOLOv8."
 )
 
 # Launch app
-interface.launch()
+interface.launch()