import gradio as gr
from ultralytics import YOLO
import cv2
from deep_sort_realtime.deepsort_tracker import DeepSort
import logging

def initialize_tracker(max_age=30, n_init=3, nn_budget=100):
    return DeepSort(max_age=max_age, n_init=n_init, nn_budget=nn_budget)

def detect_people(model, frame, confidence_threshold=0.5):
    results = model(frame, device="cpu")  # Force CPU if CUDA is unavailable
    detections = []
    for result in results:
        for box, cls, conf in zip(result.boxes.xyxy, result.boxes.cls, result.boxes.conf):
            if result.names[int(cls)] == "person" and conf > confidence_threshold:
                x1, y1, x2, y2 = map(int, box)
                bbox = [x1, y1, x2 - x1, y2 - y1]
                detections.append((bbox, conf, "person"))
    return detections

def count_people_in_video(video_path, model_path="setosys_ppl_in_video_small_v1.pt", confidence_threshold=0.5):
    logging.basicConfig(level=logging.INFO)
    cap = cv2.VideoCapture(video_path)
    model = YOLO(model_path)  # Auto-detect device during model loading
    tracker = initialize_tracker()

    total_ids = set()
    frame_count = 0

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break

        frame_count += 1
        logging.info(f"Processing frame {frame_count}")
        detections = detect_people(model, frame, confidence_threshold)
        tracks = tracker.update_tracks(detections, frame=frame)

        for track in tracks:
            if track.is_confirmed():
                total_ids.add(track.track_id)

    cap.release()
    return len(total_ids)


# Initialize YOLO model
model = YOLO("setosys_ppl_in_video_small_v1.pt")  # Load  model
#tracker = DeepSort(max_age=30, n_init=3, nn_budget=100)

def count_people_in_video_old(video_path):
    cap = cv2.VideoCapture(video_path)  # Load video
    total_ids = set()  # Track unique IDs

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break

        # Run YOLO inference on the frame
        results = model(frame)
        detections = []

        # Parse YOLO detections
        for result in results:
            for box, cls, conf in zip(result.boxes.xyxy, result.boxes.cls, result.boxes.conf):
                if result.names[int(cls)] == "person" and conf > 0.5:  # Detect "person" class
                    x1, y1, x2, y2 = map(int, box)
                    bbox = [x1, y1, x2 - x1, y2 - y1]  # Convert to [x, y, width, height]
                    detections.append((bbox, conf, "person"))

        # Update DeepSORT tracker with detections
        tracks = tracker.update_tracks(detections, frame=frame)

        # Add unique IDs from confirmed tracks
        for track in tracks:
            if track.is_confirmed():
                total_ids.add(track.track_id)

    cap.release()
    return len(total_ids)

# Gradio Interface
def process_video(video_file):
    # The `video_file` is a path to the temporary file
    total_people = count_people_in_video(video_file)
    return f"Total unique people in the video: {total_people}"

interface = gr.Interface(
    fn=process_video,
    inputs=gr.Video(label="Upload a Video"),
    outputs="text",
    title="Person Counting in a Video",
    description="Upload a video to count the number of unique people using YOLOv8 and DeepSORT."
)

if __name__ == "__main__":
    interface.launch(share=True)