Update app.py

app.py

@@ -4,7 +4,6 @@ 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"
@@ -14,32 +13,44 @@ model = YOLO(MODEL_PATH)
 TRUCK_CLASS_ID = 7  # "truck"
 
 # Initialize SORT tracker
-tracker = Sort(max_age=20, min_hits=3, iou_threshold=0.3)  # Improved tracking stability
+tracker = Sort()
 
 # Minimum confidence threshold for detection
-CONFIDENCE_THRESHOLD = 0.4  # Adjusted to capture more trucks
+CONFIDENCE_THRESHOLD = 0.5
 
 # 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):
-    """ Determines frame skip interval based on keywords in the filename. """
+    print(f"Checking filename: {video_filename}")  # Debugging
     for keyword, interval in TIME_INTERVALS.items():
         if keyword in video_filename:
+            print(f"Matched keyword: {keyword} -> Interval: {interval}")  # Debugging
             return interval
-    return 5  # Default interval
+    print("No keyword match, using default interval: 5")  # Debugging
+    return 5  # Default interval if no keyword matches
+    
 
 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 = {}
@@ -51,13 +62,14 @@ 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))
 
-    # Dynamically adjust frame skipping based on FPS and movement density
-    frame_skip = max(1, min(fps * time_interval // 2, total_frames // 10))
+    # 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
 
     frame_count = 0
 
@@ -68,7 +80,7 @@ def count_unique_trucks(video_path):
 
         frame_count += 1
         if frame_count % frame_skip != 0:
-            continue  # Skip frames based on interval
+            continue  # Skip frames to process only every 5 seconds
 
         # Run YOLOv12x inference
         results = model(frame, verbose=False)
@@ -84,39 +96,35 @@ 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])
 
-        # Convert detections to numpy array for SORT
-        detections = np.array(detections) if len(detections) > 0 else np.empty((0, 5))
+        if len(detections) > 0:
+            detections = np.array(detections)
+            tracked_objects = tracker.update(detections)
 
-        # Update SORT tracker
-        tracked_objects = tracker.update(detections)
+            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
 
-        # 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
+                truck_center = (x1 + x2) / 2, (y1 + y2) / 2  # Calculate the center of the truck
 
-            truck_center = (x1 + x2) / 2, (y1 + y2) / 2  # Calculate truck center
+                # 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))
 
-            # 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
+                    if distance > DISTANCE_THRESHOLD:
+                        # If the truck moved significantly, count as new
+                        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
-
-            # 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)
+                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()
+
     return {"Total Unique Trucks": len(unique_truck_ids)}
 
 # Gradio UI function
@@ -125,6 +133,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"),
@@ -135,4 +144,4 @@ iface = gr.Interface(
 
 # Launch the Gradio app
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch()