app.py

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