import gradio as gr
import cv2
import spaces
import numpy as np
import tempfile
import os
from ultralytics import YOLO

@spaces.GPU()
def stream_object_detection(video_path):
    # Load the YOLO model
    model = YOLO("weights/best.pt")
    cap = cv2.VideoCapture(video_path)

    # Get video properties
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH) // 2)
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT) // 2)

    # Temporary file for processed video
    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
    temp_file_path = temp_file.name

    # VideoWriter to save processed frames
    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    out = cv2.VideoWriter(temp_file_path, fourcc, fps, (width, height))

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        frame = cv2.resize(frame, (width, height))

        # Run YOLO predictions
        results = model.predict(frame)

        # Annotate frame with detection results
        annotated_frame = results[0].plot()

        # Write annotated frame to the video file
        out.write(annotated_frame)

    cap.release()
    out.release()

    return temp_file_path

with gr.Blocks() as app:
    with gr.Row():
        with gr.Column():
            video_input = gr.Video(label="Upload Video")
            # conf_threshold = gr.Slider(
            #     label="Confidence Threshold",
            #     minimum=0.0,
            #     maximum=1.0,
            #     step=0.05,
            #     value=0.30,
            # )
        with gr.Column():
            video_output = gr.Video(label="Processed Video")
    with gr.Row():
        
        with gr.Column():
            detect_button = gr.Button("Start Detection", variant="primary")
    detect_button.click(
        fn=stream_object_detection,
        inputs=[video_input],
        outputs=video_output,
    )

if __name__ == "__main__":
    app.launch()