import gradio as gr
import cv2
import numpy as np
from ultralytics import YOLO
from sort import Sort  # SORT tracker

# Load YOLOv8 model (pre-trained on COCO dataset)
model = YOLO("yolov8x.pt")  # Highest accuracy version

# Class label for trucks (COCO dataset)
TRUCK_CLASS_ID = 7  # "truck" in COCO dataset

# Initialize SORT tracker
tracker = Sort()  

def count_trucks(video_path):
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        return "Error: Unable to open video file."

    frame_count = 0
    unique_truck_ids = set()
    frame_skip = 5  # Process every 5th frame for efficiency

    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 for efficiency

        # Run YOLOv8 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
                x1, y1, x2, y2 = map(int, box.xyxy[0])  # Get bounding box

                if class_id == TRUCK_CLASS_ID and confidence > 0.6:
                    detections.append([x1, y1, x2, y2, confidence])  # Append detection

        # Convert to numpy array for SORT input
        if len(detections) > 0:
            detections = np.array(detections)
        else:
            detections = np.empty((0, 5))  # Empty array when no trucks detected

        # Update tracker
        tracked_objects = tracker.update(detections)

        # Store unique truck IDs
        for obj in tracked_objects:
            truck_id = int(obj[4])  # SORT assigns unique IDs
            unique_truck_ids.add(truck_id)

    cap.release()

    return {
        "Total Unique Trucks in Video": len(unique_truck_ids)
    }

# Gradio UI function
def analyze_video(video_file):
    result = count_trucks(video_file)
    return "\n".join([f"{key}: {value}" for key, value in result.items()])

# Gradio Interface
interface = gr.Interface(
    fn=analyze_video,
    inputs=gr.Video(label="Upload Video"),
    outputs=gr.Textbox(label="Truck Counting Results"),
    title="YOLOv8-based Truck Counter with Object Tracking",
    description="Upload a video to detect and count unique trucks using YOLOv8 and SORT tracker."
)

# Launch app
interface.launch()