app.py

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 process_inputs(start_button, stop_button, mode_selection, uploaded_video):
    if start_button:
        if mode_selection == "Video" and uploaded_video:
            segment_video(uploaded_video)
        elif mode_selection == "Webcam":
            segment_webcam()


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(share=True)