import cv2
import gradio as gr
import supervision as sv
from ultralytics import YOLO
from PIL import Image
import torch
import time
import numpy as np
import uuid
model = YOLO("yolov8s.pt")

def stream_object_detection(video):
    cap = cv2.VideoCapture(video)
    # This means we will output mp4 videos
    video_codec = cv2.VideoWriter_fourcc(*"mp4v") # type: ignore
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    desired_fps = fps // SUBSAMPLE
    width  = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) // 2
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) // 2
    iterating, frame = cap.read()
    n_frames = 0
    output_video_name = f"output_{uuid.uuid4()}.mp4"
    output_video = cv2.VideoWriter(output_video_name, video_codec, desired_fps, (width, height)) # type: ignore

    while iterating:
        frame = cv2.resize( frame, (0,0), fx=0.5, fy=0.5)
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        result = model(Image.fromarray(frame))[0]
        detections = sv.Detections.from_ultralytics(result)
        outp = draw_box(frame,detections)
        frame = np.array(outp)
        # Convert RGB to BGR
        frame = frame[:, :, ::-1].copy()
        output_video.write(frame)
        batch = []
        output_video.release()
        yield output_video_name
        output_video_name = f"output_{uuid.uuid4()}.mp4"
        output_video = cv2.VideoWriter(output_video_name, video_codec, desired_fps, (width, height)) # type: ignore
        iterating, frame = cap.read()
        n_frames += 1

with gr.Blocks() as app:
    #inp = gr.Image(type="filepath")
    with gr.Row():
        with gr.Column():
            inp = gr.Video()
            btn = gr.Button()
        outp_v = gr.Video(label="Processed Video", streaming=True, autoplay=True)
    btn.click(stream_object_detection,inp,[outp_v])
app.queue(concurrency_limit=20).launch()