import time
import gradio as gr
import cv2 as cv
from ultralytics import YOLO

# Global variables to control the process
process = False
model = None

def load_yolo_model():
    return YOLO('yolov8n-seg.pt')

def resize_frame(frame, width=1280):
    height, width = frame.shape[:2]
    new_height = int(height * width / float(width))
    return cv.resize(frame, (width, new_height))

def process_frame(frame):
    global model
    frame = resize_frame(frame)
    start = time.perf_counter()
    results = model(frame)
    end = time.perf_counter()
    segments = results[0].plot()
    return segments, f'FPS: {int(1 // (end - start))}'

def process_inputs(start_button, stop_button, mode_selection, frame, uploaded_video):
    global process, model
    if start_button and mode_selection:
        process = True
        model = load_yolo_model()
        if mode_selection == "Webcam" and frame is not None:
            return process_frame(frame)
        elif mode_selection == "Video" and uploaded_video is not None:
            return process_frame(uploaded_video)
    elif stop_button:
        process = False

start_button = gr.inputs.Button(label="Start")
stop_button = gr.inputs.Button(label="Stop")
mode_selection = gr.inputs.Radio(["Webcam", "Video"], label="Mode Selection")
frame = gr.inputs.Video(shape=(720, 1280), label="Webcam Feed")
uploaded_video = gr.inputs.Video(shape=(720, 1280), label="Upload Video")

iface = gr.Interface(
    fn=process_inputs,
    inputs=[start_button, stop_button, mode_selection, frame, uploaded_video],
    outputs=[gr.outputs.Image(), gr.outputs.Textbox()],
    live=True,
    title="YOLO Image Segmentation",
    description="This application uses the YOLO model to perform image segmentation on a webcam feed or an uploaded video. Select 'Webcam' or 'Video', upload a video (if applicable), and click 'Start' to begin. Click 'Stop' to end the process.",
    theme="huggingface"
)

iface.launch()