File size: 1,965 Bytes
fd10dcd 5abf1f3 fd10dcd f1913e7 fd10dcd abfdd6f cfc42c6 abfdd6f cfc42c6 abfdd6f cfc42c6 abfdd6f cfc42c6 abfdd6f cfc42c6 5abf1f3 fd10dcd cfc42c6 abfdd6f 31f9d33 3858777 fd10dcd f1913e7 cfc42c6 4886b30 3858777 cfc42c6 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 |
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(label="Webcam Feed")
uploaded_video = gr.inputs.Video(label="Upload Video")
iface = gr.Interface(
fn=process_inputs,
inputs=[action, mode_selection, frame, uploaded_video],
outputs=[gr.components.Image(), gr.components.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()
|