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"
@@ -16,41 +17,32 @@ TRUCK_CLASS_ID = 7  # "truck"
 tracker = Sort()
 
 # Minimum confidence threshold for detection
-CONFIDENCE_THRESHOLD = 0.5
+CONFIDENCE_THRESHOLD = 0.4  # Lowered for better detection
 
 # 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
     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 = {}
@@ -63,13 +55,12 @@ def count_unique_trucks(video_path):
 
     # Determine the dynamic time interval based on filename keywords
     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
+    frame_skip = min(fps * time_interval, total_frames // 2)  # Reduced skipping
 
     frame_count = 0
 
@@ -80,7 +71,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)
@@ -96,35 +87,41 @@ 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])
 
+        # Debugging: Check detections
+        print(f"Frame {frame_count}: Detections -> {detections}")
+
         if len(detections) > 0:
             detections = np.array(detections)
             tracked_objects = tracker.update(detections)
+        else:
+            tracked_objects = []  # Prevent tracker from resetting
 
-            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
+        # Debugging: Check tracked objects
+        print(f"Frame {frame_count}: Tracked Objects -> {tracked_objects}")
 
-                truck_center = (x1 + x2) / 2, (y1 + y2) / 2  # Calculate the center of the truck
+        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
 
-                # 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)
+            # If truck is already in history, check 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))
 
-                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 distance > DISTANCE_THRESHOLD:
+                    unique_truck_ids.add(truck_id)  # Add only if moved significantly
 
-    cap.release()
+            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
@@ -133,7 +130,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 +141,3 @@ iface = gr.Interface(
 # Launch the Gradio app
 if __name__ == "__main__":
     iface.launch()
-