Update app.py

app.py

@@ -2,28 +2,22 @@ import os
 import cv2
 import numpy as np
 import torch
-import logging
 from ultralytics import YOLO
 from sort import Sort
 import gradio as gr
 
-# Configure logging
-logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
-
 # Load YOLOv12x model
 MODEL_PATH = "yolov12x.pt"
-if not os.path.exists(MODEL_PATH):
-    raise FileNotFoundError(f"Model file '{MODEL_PATH}' not found.")
 model = YOLO(MODEL_PATH)
 
 # COCO dataset class ID for truck
 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.4  # Adjust based on performance
+CONFIDENCE_THRESHOLD = 0.4  # Adjusted to capture more trucks
 
 # Distance threshold to avoid duplicate counts
 DISTANCE_THRESHOLD = 50
@@ -36,39 +30,38 @@ TIME_INTERVALS = {
 
 def determine_time_interval(video_filename):
     """ Determines frame skip interval based on keywords in the filename. """
-    logging.info(f"Checking filename: {video_filename}")
     for keyword, interval in TIME_INTERVALS.items():
         if keyword in video_filename:
-            logging.info(f"Matched keyword: {keyword} -> Interval: {interval}")
             return interval
-    logging.info("No keyword match, using default interval: 5")
     return 5  # Default interval
 
 def count_unique_trucks(video_path):
     """ Counts unique trucks in a video using YOLOv12x and SORT tracking. """
-    if not os.path.exists(video_path):
-        return {"Error": "Video file not found."}
-    
     cap = cv2.VideoCapture(video_path)
     if not cap.isOpened():
         return {"Error": "Unable to open video file."}
 
     unique_truck_ids = set()
     truck_history = {}
-    
-    # Get FPS and total frames
-    fps = int(cap.get(cv2.CAP_PROP_FPS)) or 30  # Default to 30 if retrieval fails
-    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) or 1
-    
-    # Extract filename and determine time interval
+
+    # Get FPS of the video
+    fps = int(cap.get(cv2.CAP_PROP_FPS))
+
+    # Extract filename from the path and convert to lowercase
     video_filename = os.path.basename(video_path).lower()
-    time_interval = determine_time_interval(video_filename)
 
-    # Ensure frame_skip does not exceed total frames
-    frame_skip = min(fps * time_interval, max(1, total_frames // 2))
+    # Determine the dynamic time interval based on filename keywords
+    #time_interval = determine_time_interval(video_filename)
+    time_interval = 7
+    # 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))
+
     frame_count = 0
 
-    while cap.isOpened():
+    while True:
         ret, frame = cap.read()
         if not ret:
             break  # End of video
@@ -83,29 +76,44 @@ def count_unique_trucks(video_path):
         detections = []
         for result in results:
             for box in result.boxes:
-                class_id = int(box.cls.item())
-                confidence = float(box.conf.item())
+                class_id = int(box.cls.item())  # Get class ID
+                confidence = float(box.conf.item())  # Get confidence score
 
+                # Track only trucks
                 if class_id == TRUCK_CLASS_ID and confidence > CONFIDENCE_THRESHOLD:
-                    x1, y1, x2, y2 = map(int, box.xyxy[0])
+                    x1, y1, x2, y2 = map(int, box.xyxy[0])  # Get bounding box
                     detections.append([x1, y1, x2, y2, confidence])
 
-        if detections:
-            tracked_objects = tracker.update(np.array(detections))
-        else:
-            tracked_objects = []
+        # Convert detections to numpy array for SORT
+        detections = np.array(detections) if len(detections) > 0 else np.empty((0, 5))
+
+        # Update SORT tracker
+        tracked_objects = tracker.update(detections)
 
+        # Track movement history to avoid duplicate counts
         for obj in tracked_objects:
-            truck_id = int(obj[4])
-            truck_center = ((obj[0] + obj[2]) / 2, (obj[1] + obj[3]) / 2)
-
-            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 distance > DISTANCE_THRESHOLD:
-                    unique_truck_ids.add(truck_id)
-            else:
-                truck_history[truck_id] = {"position": truck_center}
+            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 truck center
+
+            # 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 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)
 
     cap.release()
@@ -113,9 +121,6 @@ def count_unique_trucks(video_path):
 
 # Gradio UI function
 def analyze_video(video_file):
-    if not video_file:
-        return "Error: No video file uploaded."
-    
     result = count_unique_trucks(video_file)
     return "\n".join([f"{key}: {value}" for key, value in result.items()])