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:
        process = True
        if mode_selection == "Video" and uploaded_video:
            segment_video(uploaded_video)
        elif mode_selection == "Webcam":
            segment_webcam()
    elif stop_button:
        process = False

iface = gr.Interface(
    fn=process_inputs,
    inputs=[
        gr.components.Button(label="Start"),
        gr.components.Button(label="Stop"),
        gr.components.Radio(["Video", "Webcam"], label="Mode Selection"),
        gr.components.File(label="Upload Video")
    ],
    outputs=None,
    live=True
)

iface.launch()