Spaces:
Sleeping
Sleeping
| import shutil | |
| from pathlib import Path | |
| from typing import Dict, Union, Optional | |
| import numpy as np | |
| import gradio as gr | |
| from ultralytics import YOLO | |
| from utils import ( | |
| download_model, | |
| detect_image, | |
| detect_video, | |
| get_csv_annotate, | |
| get_matplotlib_fig, | |
| ) | |
| # ======================= MODEL =================================== | |
| MODELS_DIR = Path('models') | |
| MODELS_DIR.mkdir(exist_ok=True) | |
| MODELS = { | |
| 'yolov11n.pt': 'https://github.com/ultralytics/assets/releases/download/v8.3.0/yolo11n.pt', | |
| 'yolov11s.pt': 'https://github.com/ultralytics/assets/releases/download/v8.3.0/yolo11s.pt', | |
| 'yolov11m.pt': 'https://github.com/ultralytics/assets/releases/download/v8.3.0/yolo11m.pt', | |
| 'yolov11l.pt': 'https://github.com/ultralytics/assets/releases/download/v8.3.0/yolo11l.pt', | |
| 'yolov11x.pt': 'https://github.com/ultralytics/assets/releases/download/v8.3.0/yolo11x.pt', | |
| } | |
| MODEL_NAMES = list(MODELS.keys()) | |
| model_path = download_model(MODEL_NAMES[0], MODELS_DIR, MODELS) | |
| default_model = YOLO(model_path) | |
| IMAGE_EXTENSIONS = ['.jpg', '.jpeg', '.png'] | |
| VIDEO_EXTENSIONS = ['.mp4', '.avi'] | |
| # =================== ADDITIONAL INTERFACE FUNCTIONS ======================== | |
| def change_model(model_state: Dict[str, YOLO], model_name: str): | |
| progress = gr.Progress() | |
| progress(0.3, desc='Downloading the model') | |
| model_path = download_model(model_name, MODELS_DIR, MODELS) | |
| progress(0.7, desc='Model initialization') | |
| model_state['model'] = YOLO(model_path) | |
| return f'Model {model_name} initialized' | |
| def detect(file_path: str, file_link: str, model_state: Dict[str, YOLO], conf: float, iou: float): | |
| model = model_state['model'] | |
| if file_link: | |
| file_path = file_link | |
| file_ext = f'.{file_path.rsplit(".")[-1]}' | |
| if file_ext in IMAGE_EXTENSIONS: | |
| np_image = detect_image(file_path, model, conf, iou) | |
| return np_image, "Detection complete, opening image..." | |
| elif file_ext in VIDEO_EXTENSIONS or 'youtube.com' in file_link: | |
| video_path = detect_video(file_path, model, conf, iou) | |
| return video_path, "Detection complete, converting and opening video..." | |
| else: | |
| gr.Info('Invalid image or video format...') | |
| return None, None | |
| # =================== INTERFACE COMPONENTS ============================ | |
| def get_output_media_components(detect_result: Optional[Union[np.ndarray, str, Path]] = None): | |
| visible = isinstance(detect_result, np.ndarray) | |
| image_output = gr.Image( | |
| value=detect_result if visible else None, | |
| type="numpy", | |
| width=640, | |
| height=480, | |
| visible=visible, | |
| label='Output', | |
| ) | |
| visible = isinstance(detect_result, (str, Path)) | |
| video_output = gr.Video( | |
| value=detect_result if visible else None, | |
| width=640, | |
| height=480, | |
| visible=visible, | |
| label='Output', | |
| ) | |
| clear_btn = gr.Button( | |
| value='Clear', | |
| scale=0, | |
| visible=detect_result is not None, | |
| ) | |
| return image_output, video_output, clear_btn | |
| def get_download_csv_btn(csv_annotations_path: Optional[Path] = None): | |
| download_csv_btn = gr.DownloadButton( | |
| label='Download csv annotations for video', | |
| value=csv_annotations_path, | |
| scale=0, | |
| visible=csv_annotations_path is not None, | |
| ) | |
| return download_csv_btn | |
| # =================== APPINTERFACE ========================== | |
| css = '''.gradio-container {width: 70% !important}''' | |
| with gr.Blocks(css=css) as demo: | |
| with gr.Tab('Detection image / video'): | |
| gr.HTML("""<h3 style='text-align: center'>YOLOv11 Detector</h3>""") | |
| model_state = gr.State({'model': default_model}) | |
| detect_result = gr.State(None) | |
| csv_annotations_path = gr.State(None) | |
| with gr.Row(): | |
| with gr.Column(): | |
| file_path = gr.File(file_types=['image', 'video'], file_count='single', label='Select image or video') | |
| # file_link = gr.Textbox(label='Direct link to image or YouTube link') | |
| file_link = gr.State(None) | |
| model_name = gr.Radio(choices=MODEL_NAMES, value=MODEL_NAMES[0], label='Select YOLO model') | |
| conf = gr.Slider(0, 1, value=0.5, step=0.05, label='Confidence') | |
| iou = gr.Slider(0, 1, value=0.7, step=0.1, label='IOU') | |
| status_message = gr.Textbox(value='Ready to go', label='Status') | |
| detect_btn = gr.Button('Detect', interactive=True) | |
| with gr.Column(): | |
| image_output, video_output, clear_btn = get_output_media_components() | |
| download_csv_btn = get_download_csv_btn() | |
| model_name.change( | |
| fn=lambda: gr.update(interactive=False), | |
| inputs=None, | |
| outputs=[detect_btn], | |
| ).then( | |
| fn=change_model, | |
| inputs=[model_state, model_name], | |
| outputs=[status_message], | |
| ).success( | |
| fn=lambda: gr.update(interactive=True), | |
| inputs=None, | |
| outputs=[detect_btn], | |
| ) | |
| detect_btn.click( | |
| fn=detect, | |
| inputs=[file_path, file_link, model_state, conf, iou], | |
| outputs=[detect_result, status_message], | |
| ).success( | |
| fn=get_output_media_components, | |
| inputs=[detect_result], | |
| outputs=[image_output, video_output, clear_btn], | |
| ).then( | |
| fn=lambda: 'Ready to go', | |
| inputs=None, | |
| outputs=[status_message], | |
| ).then( | |
| fn=get_csv_annotate, | |
| inputs=[detect_result], | |
| outputs=[csv_annotations_path], | |
| ).success( | |
| fn=get_download_csv_btn, | |
| inputs=[csv_annotations_path], | |
| outputs=[download_csv_btn], | |
| ) | |
| def clear_results_dir(detect_result): | |
| if isinstance(detect_result, Path): | |
| shutil.rmtree(detect_result.parent, ignore_errors=True) | |
| clear_components = [image_output, video_output, clear_btn, download_csv_btn] | |
| clear_btn.click( | |
| fn=lambda: [gr.update(visible=False) for _ in range(len(clear_components))], | |
| inputs=None, | |
| outputs=clear_components, | |
| ).then( | |
| fn=clear_results_dir, | |
| inputs=[detect_result], | |
| outputs=None, | |
| ).then( | |
| fn=lambda: (None, None), | |
| inputs=None, | |
| outputs=[detect_result, csv_annotations_path] | |
| ) | |
| gr.HTML("""<h3 style='text-align: center'> | |
| <a href="https://github.com/sergey21000/yolo-detector" target='_blank'>GitHub Repository</a></h3> | |
| """) | |
| with gr.Tab('Show detection video results'): | |
| show_results_btn = gr.Button('Show detection results', scale=1) | |
| gr_info = 'To display the results, perform video detection on the first tab' | |
| show_results_btn.click( | |
| fn=lambda csv_path: get_matplotlib_fig(csv_path) if csv_path is not None else gr.Info(gr_info), | |
| inputs=[csv_annotations_path], | |
| outputs=gr.Plot(), | |
| ) | |
| demo.launch(server_name='0.0.0.0') # debug=True |