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

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

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

def process_video():
    global process, model, cap
    while process and cap.isOpened():
        ret, image = cap.read()
        if not ret:
            break
        start = time.perf_counter()
        results = model(image)
        end = time.perf_counter()
        segments = results[0].plot()
        cv.putText(segments, f'FPS: {int(1 // (end - start))}', (10, 30),
                   cv.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
        cv.imshow('Image Segmentation', segments)
        key = cv.waitKey(1)
        if key & 0xFF == ord('q'):
            break
    cap.release()
    cv.destroyAllWindows()

def segment_video(uploaded_video):
    global model, cap
    model = load_yolo_model()
    cap = cv.VideoCapture(uploaded_video.name)
    process_video()

def segment_webcam():
    global model, cap
    model = load_yolo_model()
    cap = cv.VideoCapture(0)
    process_video()

def process_inputs(start_button, stop_button, mode_selection, uploaded_video):
    global process
    if start_button and mode_selection:
        process = True
        if mode_selection == "Video" and uploaded_video:
            segment_video(uploaded_video)
        elif mode_selection == "Webcam":
            segment_webcam()
    elif stop_button:
        process = False

start_button = gr.components.Button(label="Start", color="green", disabled=True)
stop_button = gr.components.Button(label="Stop", color="red", disabled=True)
mode_selection = gr.components.Radio(["Video", "Webcam"], label="Mode Selection")
uploaded_video = gr.components.File(label="Upload Video")

def on_mode_change(_, __, mode):
    if mode == "Webcam":
        uploaded_video.disable()
    else:
        uploaded_video.enable()
    start_button.enable()

mode_selection.on_change(on_mode_change)

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

iface.launch()