Update app.py

app.py

@@ -1,21 +1,4 @@
-import cv2
-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"
-model = YOLO(MODEL_PATH)
-
-# COCO dataset class ID for truck
-TRUCK_CLASS_ID = 7  # "truck"
-
-# Initialize SORT tracker
-tracker = Sort()
-
-def count_unique_trucks(video_path, video_type, confidence_threshold, distance_threshold, frame_skip_seconds):
+def count_unique_trucks(video_path, frame_skip_factor=2):
     cap = cv2.VideoCapture(video_path)
     if not cap.isOpened():
         return "Error: Unable to open video file."
@@ -25,7 +8,7 @@ def count_unique_trucks(video_path, video_type, confidence_threshold, distance_t
 
     # Get FPS of the video
     fps = int(cap.get(cv2.CAP_PROP_FPS))
-    frame_skip = max(1, int(fps * frame_skip_seconds))  # Dynamic frame skipping
+    frame_skip = fps * frame_skip_factor  # Skip frames based on the dynamic factor
 
     frame_count = 0
 
@@ -48,7 +31,7 @@ def count_unique_trucks(video_path, video_type, confidence_threshold, distance_t
                 confidence = float(box.conf.item())  # Get confidence score
 
                 # Track only trucks
-                if class_id == TRUCK_CLASS_ID and confidence > confidence_threshold:
+                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])
 
@@ -67,7 +50,7 @@ def count_unique_trucks(video_path, video_type, confidence_threshold, distance_t
                     last_position = truck_history[truck_id]["position"]
                     distance = np.linalg.norm(np.array(truck_center) - np.array(last_position))
 
-                    if distance > distance_threshold:
+                    if distance > DISTANCE_THRESHOLD:
                         # If the truck moved significantly, count as new
                         unique_truck_ids.add(truck_id)
 
@@ -80,25 +63,25 @@ def count_unique_trucks(video_path, video_type, confidence_threshold, distance_t
                     unique_truck_ids.add(truck_id)
 
     cap.release()
+
     return {"Total Unique Trucks": len(unique_truck_ids)}
 
 # Gradio UI function
-def analyze_video(video_file, video_type, confidence, distance, frame_skip):
-    return count_unique_trucks(video_file, video_type, confidence, distance, frame_skip)
+def analyze_video(video_file, frame_skip_factor):
+    result = count_unique_trucks(video_file, frame_skip_factor)
+    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"),
-        gr.Radio(["drone", "fixed"], label="Video Type"),
-        gr.Slider(0.3, 0.9, 0.5, label="Confidence Threshold"),
-        gr.Slider(10, 100, 50, label="Distance Threshold"),
-        gr.Slider(1, 10, 2, label="Frame Skip (Seconds)"),
+        gr.Slider(minimum=1, maximum=10, step=1, default=2, label="Frame Skip Factor"),
     ],
-    outputs=gr.JSON(label="Analysis Result"),
-    title="YOLOv12x Dynamic Truck Counter",
-    description="Upload a video, adjust parameters, and analyze unique trucks."
+    outputs=gr.Textbox(label="Analysis Result"),
+    title="YOLOv12x Unique Truck Counter",
+    description="Upload a video to count unique trucks using YOLOv12x and SORT tracking. Adjust the frame skip factor to control processing speed."
 )
 
 # Launch the Gradio app