app.py

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_inputs(action, mode_selection, frame, uploaded_video):
    global process, model
    if action == "Start" and not process and mode_selection:
        if mode_selection == "Webcam" and frame is not None:
            process = True
            model = load_yolo_model()
            return process_frame(frame)
        elif mode_selection == "Video" and uploaded_video is not None:
            process = True
            model = load_yolo_model()
            return process_frame(uploaded_video)
    elif action == "Stop" and process:
        process = False

action = gr.inputs.Dropdown(choices=["Start", "Stop"], label="Action")
mode_selection = gr.inputs.Radio(["Webcam", "Video"], label="Mode Selection")
frame = gr.inputs.Video(shape=(720, 1280), label="Webcam Feed")
uploaded_video = gr.inputs.Video(shape=(720, 1280), label="Upload Video")

iface = gr.Interface(
    fn=process_inputs,
    inputs=[action, mode_selection, frame, uploaded_video],
    outputs=[gr.outputs.Image(), gr.outputs.Textbox()],
    live=True,
    title="YOLO Image Segmentation",
    description="This application uses the YOLO model to perform image segmentation on a webcam feed or an uploaded video. Select 'Webcam' or 'Video', upload a video (if applicable), and select 'Start' to begin. Select 'Stop' to end the process.",
    theme="huggingface"
)

iface.launch()