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_webcam(start_button, stop_button, frame):
    global process, model
    if start_button and not process and frame is not None:
        process = True
        model = load_yolo_model()
        return process_frame(frame)
    elif stop_button and process:
        process = False
        return None, ""

def process_video(start_button, stop_button, uploaded_video):
    global process, model
    if start_button and not process and uploaded_video is not None:
        process = True
        model = load_yolo_model()
        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")
frame = gr.components.Video(label="Webcam Feed")
uploaded_video = gr.components.Video(label="Upload Video")

webcam_interface = gr.Interface(
    fn=process_webcam,
    inputs=[start_button, stop_button, frame],
    outputs=[gr.components.Image(), gr.components.Textbox()],
    live=True,
    title="YOLO Image Segmentation (Webcam)",
    description="This application uses the YOLO model to perform image segmentation on a webcam feed. Check 'Start' to begin. Check 'Stop' to end the process.",
)

video_interface = gr.Interface(
    fn=process_video,
    inputs=[start_button, stop_button, uploaded_video],
    outputs=[gr.components.Image(), gr.components.Textbox()],
    live=True,
    title="YOLO Image Segmentation (Video)",
    description="This application uses the YOLO model to perform image segmentation on an uploaded video. Upload a video and check 'Start' to begin. Check 'Stop' to end the process.",
)

gr.TabbedInterface(
    [webcam_interface, video_interface],
    tab_names=['Webcam inference', 'Video inference']
).launch(share=True)