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()
|