File size: 2,545 Bytes
fd10dcd 5abf1f3 fd10dcd f1913e7 fd10dcd abfdd6f 04c39e7 abfdd6f 04c39e7 abfdd6f 04c39e7 5abf1f3 04c39e7 fd10dcd 04dda20 3858777 04c39e7 4886b30 3858777 04c39e7 fd10dcd 252a43a |
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 73 74 75 76 77 |
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_webcam(start_button, stop_button, frame):
global process, model
if start_button and not process and frame is not None:
process = True
model = load_yolo_model()
return process_frame(frame)
elif stop_button and process:
process = False
return None, ""
def process_video(start_button, stop_button, uploaded_video):
global process, model
if start_button and not process and uploaded_video is not None:
process = True
model = load_yolo_model()
cap = cv.VideoCapture(uploaded_video.name)
ret, frame = cap.read()
if ret:
return process_frame(frame)
elif stop_button and process:
process = False
return None, ""
start_button = gr.components.Checkbox(label="Start")
stop_button = gr.components.Checkbox(label="Stop")
frame = gr.components.Video(label="Webcam Feed")
uploaded_video = gr.components.Video(label="Upload Video")
webcam_interface = gr.Interface(
fn=process_webcam,
inputs=[start_button, stop_button, frame],
outputs=[gr.components.Image(), gr.components.Textbox()],
live=True,
title="YOLO Image Segmentation (Webcam)",
description="This application uses the YOLO model to perform image segmentation on a webcam feed. Check 'Start' to begin. Check 'Stop' to end the process.",
)
video_interface = gr.Interface(
fn=process_video,
inputs=[start_button, stop_button, uploaded_video],
outputs=[gr.components.Image(), gr.components.Textbox()],
live=True,
title="YOLO Image Segmentation (Video)",
description="This application uses the YOLO model to perform image segmentation on an uploaded video. Upload a video and check 'Start' to begin. Check 'Stop' to end the process.",
)
gr.TabbedInterface(
[webcam_interface, video_interface],
tab_names=['Webcam inference', 'Video inference']
).launch(share=True)
|