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

# Load the YOLO model
model = 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):
    frame = resize_frame(frame)
    results = model(frame)
    segments = results[0].plot()
    return segments

frame = gr.components.Video(label="Webcam Feed")

iface = gr.Interface(
    fn=process_frame,
    inputs=[frame],
    outputs=[gr.components.Image()],
    live=True,
    title="YOLO Image Segmentation",
    description="This application uses the YOLO model to perform image segmentation on a webcam feed.",
)

iface.launch()