| import time |
| import gradio as gr |
| import cv2 as cv |
| from ultralytics import YOLO |
|
|
|
|
| def load_yolo_model(): |
| return YOLO() |
|
|
|
|
| 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() |
|
|
