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 not process:
        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:
            cap = cv.VideoCapture(uploaded_video.name)
            ret, frame = cap.read()
            if ret:
                return process_frame(frame)
    elif stop_button and process:
        process = False
        return None, ""

start_button = gr.components.Checkbox(label="Start")
stop_button = gr.components.Checkbox(label="Stop")
mode_selection = gr.components.Radio(["Webcam", "Video"], label="Mode Selection")
frame = gr.components.Video(label="Webcam Feed")
uploaded_video = gr.components.Video(label="Upload Video")

iface = gr.Interface(
    fn=process_inputs,
    inputs=[start_button, stop_button, mode_selection, frame, uploaded_video],
    outputs=[gr.components.Image(), gr.components.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 check 'Start' to begin. Check 'Stop' to end the process.",
)

iface.launch(share=True)