Update app.py

app.py

@@ -1,7 +1,6 @@
 import cv2
 import numpy as np
 import torch
-import gradio as gr
 from ultralytics import YOLO
 from sort import Sort
 
@@ -15,17 +14,23 @@ TRUCK_CLASS_ID = 7  # "truck"
 # Initialize SORT tracker
 tracker = Sort()
 
+# Minimum confidence threshold for detection
+CONFIDENCE_THRESHOLD = 0.5
+
+# Distance threshold to avoid duplicate counts
+DISTANCE_THRESHOLD = 50
+
 def count_unique_trucks(video_path):
     cap = cv2.VideoCapture(video_path)
     if not cap.isOpened():
         return "Error: Unable to open video file."
 
     unique_truck_ids = set()
-    frame_skip = 5  # Process every 5th frame for efficiency
-    truck_tracker_history = {}  # To store truck ID with frame count and position
-    min_distance_threshold = 50  # Minimum distance between truck positions to be counted as different
+    truck_history = {}
 
+    frame_skip = 5  # Process every 5th frame for efficiency
     frame_count = 0
+
     while True:
         ret, frame = cap.read()
         if not ret:
@@ -45,7 +50,7 @@ def count_unique_trucks(video_path):
                 confidence = float(box.conf.item())  # Get confidence score
 
                 # Track only trucks
-                if class_id == TRUCK_CLASS_ID and confidence > 0.5:
+                if class_id == TRUCK_CLASS_ID and confidence > CONFIDENCE_THRESHOLD:
                     x1, y1, x2, y2 = map(int, box.xyxy[0])  # Get bounding box
                     detections.append([x1, y1, x2, y2, confidence])
 
@@ -57,30 +62,24 @@ def count_unique_trucks(video_path):
                 truck_id = int(obj[4])  # Unique ID assigned by SORT
                 x1, y1, x2, y2 = obj[:4]  # Get the bounding box coordinates
 
-                # Compute the center of the truck for better tracking
-                truck_center = ((x1 + x2) / 2, (y1 + y2) / 2)
+                truck_center = (x1 + x2) / 2, (y1 + y2) / 2  # Calculate the center of the truck
 
-                # Only add truck ID if it's sufficiently far from previously tracked trucks
-                if truck_id not in truck_tracker_history:
-                    truck_tracker_history[truck_id] = {
-                        "frame_count": frame_count,
-                        "position": truck_center
-                    }
-                    unique_truck_ids.add(truck_id)  # Add the truck as a unique truck
-                else:
-                    last_truck = truck_tracker_history[truck_id]
-                    last_position = last_truck["position"]
-
-                    # Calculate the distance between the current and last positions
+                # If truck is already in history, check the movement distance
+                if truck_id in truck_history:
+                    last_position = truck_history[truck_id]["position"]
                     distance = np.linalg.norm(np.array(truck_center) - np.array(last_position))
 
-                    # If the distance between truck positions is greater than the threshold, it's a new truck detection
-                    if distance > min_distance_threshold:
+                    if distance > DISTANCE_THRESHOLD:
+                        # If the truck moved significantly, count as new
                         unique_truck_ids.add(truck_id)
-                        truck_tracker_history[truck_id] = {
-                            "frame_count": frame_count,
-                            "position": truck_center
-                        }
+
+                else:
+                    # If truck is not in history, add it
+                    truck_history[truck_id] = {
+                        "frame_count": frame_count,
+                        "position": truck_center
+                    }
+                    unique_truck_ids.add(truck_id)
 
     cap.release()
 
@@ -92,6 +91,7 @@ def analyze_video(video_file):
     return "\n".join([f"{key}: {value}" for key, value in result.items()])
 
 # Define Gradio interface
+import gradio as gr
 iface = gr.Interface(
     fn=analyze_video,
     inputs=gr.Video(label="Upload Video"),