File size: 2,191 Bytes
fd10dcd
 
 
 
 
5abf1f3
 
 
 
fd10dcd
 
f1913e7
fd10dcd
5abf1f3
 
 
f1913e7
 
 
 
 
 
 
 
 
 
 
 
 
 
 
fd10dcd
 
5abf1f3
f1913e7
 
5abf1f3
fd10dcd
 
5abf1f3
f1913e7
 
5abf1f3
fd10dcd
9ec2149
5abf1f3
3858777
5abf1f3
2dbfa23
9ec2149
 
 
5abf1f3
 
fd10dcd
2dbfa23
 
3858777
 
 
fd10dcd
f1913e7
3858777
fd10dcd
3858777
 
 
 
fd10dcd
 
ebaceb4
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
import time
import gradio as gr
import cv2 as cv
from ultralytics import YOLO

# Global variables to control the process
process = False
model = None
cap = None

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

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

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

def process_inputs(start_button, stop_button, mode_selection, uploaded_video):
    global process
    if start_button and mode_selection:
        process = True
        if mode_selection == "Video" and uploaded_video is not None:
            segment_video(uploaded_video)
        elif mode_selection == "Webcam":
            segment_webcam()
    elif stop_button:
        process = False

start_button = gr.components.Button(label="Start")
stop_button = gr.components.Button(label="Stop")
mode_selection = gr.components.Radio(["Video", "Webcam"], label="Mode Selection")
uploaded_video = gr.components.File(label="Upload Video")

iface = gr.Interface(
    fn=process_inputs,
    inputs=[start_button, stop_button, mode_selection, uploaded_video],
    outputs=None,
    live=True,
    title="YOLO Image Segmentation",
    description="This application uses the YOLO model to perform image segmentation on a video or webcam feed. Select a mode, upload a video (if applicable), and click 'Start' to begin. Click 'Stop' to end the process.",
    theme="huggingface"
)

iface.launch()