| import time |
| import gradio as gr |
| import cv2 as cv |
| from ultralytics import YOLO |
|
|
| |
| 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(start_button, stop_button, mode_selection, frame, uploaded_video): |
| global process, model |
| if start_button and not process: |
| process = True |
| model = load_yolo_model() |
| if mode_selection == "Webcam" and frame is not None: |
| return process_frame(frame) |
| elif mode_selection == "Video" and uploaded_video is not None: |
| 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") |
| mode_selection = gr.components.Radio(["Webcam", "Video"], label="Mode Selection") |
| frame = gr.components.Video(label="Webcam Feed") |
| uploaded_video = gr.components.Video(label="Upload Video") |
|
|
| iface = gr.Interface( |
| fn=process_inputs, |
| inputs=[start_button, stop_button, 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 check 'Start' to begin. Check 'Stop' to end the process.", |
| ) |
|
|
| iface.launch(share=True) |
|
|
