| 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 wrapper_for_segment_video(uploaded_video): | |
| segment_video(uploaded_video) | |
| iface = gr.Interface( | |
| inputs=[gr.inputs.Checkbox(label="Segment Video"), | |
| gr.inputs.Checkbox(label="Segment Webcam"), | |
| gr.inputs.File(label="Upload Video")], | |
| outputs=None, | |
| live=True | |
| ) | |
| def update_fn(segment_video_checkbox, segment_webcam_checkbox, uploaded_video): | |
| if segment_video_checkbox and uploaded_video: | |
| iface.fn = wrapper_for_segment_video | |
| iface.args = [uploaded_video] | |
| elif segment_webcam_checkbox: | |
| iface.fn = segment_webcam | |
| iface.args = [] | |
| else: | |
| iface.fn = None | |
| iface.update(update_fn) | |
| iface.launch() | |