import gradio as gr
from ultralytics import YOLO
import cv2
import numpy as np
import tempfile  # For creating temporary files
import os

# Load the model (do this outside the function for efficiency)
try:
    model = YOLO("yolov8n.pt")  # Or your model path
except Exception as e:
    print(f"Error loading model: {e}")
    exit()  # Exit if the model fails to load

def detect_pedestrians(video):
    try:
        cap = cv2.VideoCapture(video)
        fps = cap.get(cv2.CAP_PROP_FPS)  # Get the video's FPS
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # Codec for MP4 (you can change this)

        # Create a temporary file for the processed video
        with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as temp_video:
            out = cv2.VideoWriter(temp_video.name, fourcc, fps, (width, height))
            temp_file_name = temp_video.name #store the file name

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

                results = model(frame)

                for result in results:
                    boxes = result.boxes
                    for box in boxes:
                        if result.names[int(box.cls)] == 'person':
                            x1, y1, x2, y2 = map(int, box.xyxy[0])
                            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
                            cv2.putText(frame, 'Person', (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
                out.write(frame)  # Write the processed frame to the video
            cap.release()
            out.release()

        return temp_file_name # Return the path to the temporary video file

    except Exception as e:
        print(f"Error in detection: {e}")
        return f"Error: {e}"  # Return the error message as a string


iface = gr.Interface(
    fn=detect_pedestrians,
    inputs=gr.Video(),  # No 'source' argument here
    outputs=gr.Video(),  # Output is now a video
    title="Pedestrian Detection",
    description="Upload a video to detect pedestrians.",
    allow_flagging="never",
)

iface.launch()