Update app.py

app.py

@@ -4,6 +4,7 @@ import numpy as np
 import torch
 from ultralytics import YOLO
 from sort import Sort
+import gradio as gr
 
 # Load YOLOv12x model
 MODEL_PATH = "yolov12x.pt"
@@ -13,44 +14,32 @@ model = YOLO(MODEL_PATH)
 TRUCK_CLASS_ID = 7  # "truck"
 
 # Initialize SORT tracker
-tracker = Sort()
+tracker = Sort(max_age=20, min_hits=3, iou_threshold=0.3)  # Improved tracking stability
 
 # Minimum confidence threshold for detection
-CONFIDENCE_THRESHOLD = 0.5
+CONFIDENCE_THRESHOLD = 0.4  # Adjusted to capture more trucks
 
 # Distance threshold to avoid duplicate counts
 DISTANCE_THRESHOLD = 50
 
 # Dictionary to define keyword-based time intervals
 TIME_INTERVALS = {
-    "one": 1,
-    "two": 2,
-    "three": 3,
-    "four": 4,
-    "five": 5,
-    "six": 6,
-    "seven": 7,
-    "eight": 8,
-    "nine": 9,
-    "ten": 10,
-    "eleven": 11
+    "one": 1, "two": 2, "three": 3, "four": 4, "five": 5,
+    "six": 6, "seven": 7, "eight": 8, "nine": 9, "ten": 10, "eleven": 11
 }
 
-
 def determine_time_interval(video_filename):
-    print(f"Checking filename: {video_filename}")  # Debugging
+    """ Determines frame skip interval based on keywords in the filename. """
     for keyword, interval in TIME_INTERVALS.items():
         if keyword in video_filename:
-            print(f"Matched keyword: {keyword} -> Interval: {interval}")  # Debugging
             return interval
-    print("No keyword match, using default interval: 5")  # Debugging
-    return 5  # Default interval if no keyword matches
-    
+    return 5  # Default interval
 
 def count_unique_trucks(video_path):
+    """ Counts unique trucks in a video using YOLOv12x and SORT tracking. """
     cap = cv2.VideoCapture(video_path)
     if not cap.isOpened():
-        return "Error: Unable to open video file."
+        return {"Error": "Unable to open video file."}
 
     unique_truck_ids = set()
     truck_history = {}
@@ -62,15 +51,13 @@ def count_unique_trucks(video_path):
     video_filename = os.path.basename(video_path).lower()
 
     # Determine the dynamic time interval based on filename keywords
-    #time_interval = determine_time_interval(video_filename)
-    time_interval = 7
+    time_interval = determine_time_interval(video_filename)
+
     # Get total frames in the video
     total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
 
-    # Ensure frame_skip does not exceed total frames
-    frame_skip = min(fps * time_interval, total_frames)
-
-    #frame_skip = fps * time_interval  # Convert time interval to frame count
+    # Dynamically adjust frame skipping based on FPS and movement density
+    frame_skip = max(1, min(fps * time_interval // 2, total_frames // 10))
 
     frame_count = 0
 
@@ -81,7 +68,7 @@ def count_unique_trucks(video_path):
 
         frame_count += 1
         if frame_count % frame_skip != 0:
-            continue  # Skip frames to process only every 5 seconds
+            continue  # Skip frames based on interval
 
         # Run YOLOv12x inference
         results = model(frame, verbose=False)
@@ -97,35 +84,39 @@ def count_unique_trucks(video_path):
                     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)
+        # Convert detections to numpy array for SORT
+        detections = np.array(detections) if len(detections) > 0 else np.empty((0, 5))
 
-            for obj in tracked_objects:
-                truck_id = int(obj[4])  # Unique ID assigned by SORT
-                x1, y1, x2, y2 = obj[:4]  # Get the bounding box coordinates
+        # Update SORT tracker
+        tracked_objects = tracker.update(detections)
 
-                truck_center = (x1 + x2) / 2, (y1 + y2) / 2  # Calculate the center of the truck
+        # Track movement history to avoid duplicate counts
+        for obj in tracked_objects:
+            truck_id = int(obj[4])  # Unique ID assigned by SORT
+            x1, y1, x2, y2 = obj[:4]  # Get bounding box coordinates
 
-                # 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))
+            truck_center = (x1 + x2) / 2, (y1 + y2) / 2  # Calculate truck center
 
-                    if distance > DISTANCE_THRESHOLD:
-                        # If the truck moved significantly, count as new
-                        unique_truck_ids.add(truck_id)
+            # Entry-exit zone logic (e.g., bottom 20% of the frame)
+            frame_height, frame_width = frame.shape[:2]
+            entry_line = frame_height * 0.8  # Bottom 20% of the frame
+            exit_line = frame_height * 0.2  # Top 20% of the frame
 
-                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)
+            if truck_id not in truck_history:
+                # New truck detected
+                truck_history[truck_id] = {
+                    "position": truck_center,
+                    "crossed_entry": truck_center[1] > entry_line,
+                    "crossed_exit": False
+                }
+                continue
 
-    cap.release()
+            # If the truck crosses from entry to exit, count it
+            if truck_history[truck_id]["crossed_entry"] and truck_center[1] < exit_line:
+                truck_history[truck_id]["crossed_exit"] = True
+                unique_truck_ids.add(truck_id)
 
+    cap.release()
     return {"Total Unique Trucks": len(unique_truck_ids)}
 
 # Gradio UI function
@@ -134,7 +125,6 @@ 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"),
@@ -145,4 +135,4 @@ iface = gr.Interface(
 
 # Launch the Gradio app
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch()