Spaces:
Sleeping
Sleeping
File size: 3,858 Bytes
69545c2 6fab089 69545c2 19f2c02 0a9b429 4e8a148 d211767 0a9b429 9bb7cec 6fab089 9bb7cec 4e8a148 d211767 0a9b429 6fab089 0a9b429 4e8a148 0a9b429 4e8a148 b0801de 4e8a148 974967f 4e8a148 0a9b429 4e8a148 a210028 4e5fb8f d211767 a210028 d211767 a210028 b0801de d211767 4e5fb8f d211767 0a9b429 974967f 4e8a148 3503d68 6fab089 b0801de 3503d68 6fab089 3503d68 6fab089 b0801de 6fab089 3503d68 611534d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 |
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()
# Minimum confidence threshold for detection
CONFIDENCE_THRESHOLD = 0.5
# Distance threshold to avoid duplicate counts
DISTANCE_THRESHOLD = 50
def count_unique_trucks(video_path, time_interval):
cap = cv2.VideoCapture(video_path)
if not cap.isOpened():
return "Error: Unable to open video file."
unique_truck_ids = set()
truck_history = {}
# Get FPS of the video
fps = cap.get(cv2.CAP_PROP_FPS)
# Ensure FPS is valid, fallback to default if needed
if fps is None or fps <= 0:
fps = 30 # Default fallback FPS
# Set frame skip dynamically based on FPS and time_interval
frame_skip = max(int(fps * time_interval), 1) # Ensure frame_skip is at least 1
frame_count = 0
while True:
ret, frame = cap.read()
if not ret:
break # End of video
frame_count += 1
if frame_count % frame_skip != 0:
continue # Skip frames dynamically
# Run YOLOv12x inference
results = model(frame, verbose=False)
detections = []
for result in results:
for box in result.boxes:
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]) # Get bounding box
detections.append([x1, y1, x2, y2, confidence])
if len(detections) > 0:
detections = np.array(detections)
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
truck_center = (x1 + x2) / 2, (y1 + y2) / 2 # Calculate the center of the truck
# 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))
if distance > DISTANCE_THRESHOLD:
# If the truck moved significantly, count as new
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
def analyze_video(video_file, time_interval=7):
result = count_unique_trucks(video_file, time_interval)
return "\n".join([f"{key}: {value}" for key, value in result.items()])
# Define Gradio interface with time_interval as a dynamic input
iface = gr.Interface(
fn=analyze_video,
inputs=[
gr.Video(label="Upload Video"),
gr.Slider(minimum=1, maximum=20, step=1, value=7, label="Time Interval (Seconds)")
],
outputs=gr.Textbox(label="Analysis Result"),
title="YOLOv12x Unique Truck Counter",
description="Upload a video and adjust the time interval to count unique trucks using YOLOv12x and SORT tracking."
)
# Launch the Gradio app
if __name__ == "__main__":
iface.launch()
|