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


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


def process_video(model, cap):
    while 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):
    model = load_yolo_model()
    cap = cv.VideoCapture(uploaded_video.name)
    process_video(model, cap)


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


def wrapper_for_segment_video(uploaded_video):
    segment_video(uploaded_video)


iface = gr.Interface(
    inputs=[gr.inputs.Checkbox(label="Segment Video"),
            gr.inputs.Checkbox(label="Segment Webcam"),
            gr.inputs.File(label="Upload Video")],
    outputs=None,
    live=True
)


def update_fn(segment_video_checkbox, segment_webcam_checkbox, uploaded_video):
    if segment_video_checkbox and uploaded_video:
        iface.fn = wrapper_for_segment_video
        iface.args = [uploaded_video]
    elif segment_webcam_checkbox:
        iface.fn = segment_webcam
        iface.args = []
    else:
        iface.fn = None


iface.update(update_fn)
iface.launch()