#!/usr/bin/env python

from __future__ import annotations

import pathlib

import gradio as gr

from model import Model

DESCRIPTION = "# [CBNetV2](https://github.com/VDIGPKU/CBNetV2)"

model = Model()

with gr.Blocks(css="style.css") as demo:
    gr.Markdown(DESCRIPTION)

    with gr.Row():
        with gr.Column():
            with gr.Row():
                input_image = gr.Image(label="Input Image", type="numpy")
            with gr.Row():
                detector_name = gr.Dropdown(
                    label="Detector", choices=list(model.models.keys()), value=model.model_name
                )
            with gr.Row():
                detect_button = gr.Button("Detect")
                detection_results = gr.Variable()
        with gr.Column():
            with gr.Row():
                detection_visualization = gr.Image(label="Detection Result", type="numpy")
            with gr.Row():
                visualization_score_threshold = gr.Slider(
                    label="Visualization Score Threshold", minimum=0, maximum=1, step=0.05, value=0.3
                )
            with gr.Row():
                redraw_button = gr.Button("Redraw")

    with gr.Row():
        paths = sorted(pathlib.Path("images").rglob("*.jpg"))
        gr.Examples(examples=[[path.as_posix()] for path in paths], inputs=input_image)

    detector_name.change(fn=model.set_model_name, inputs=[detector_name], outputs=None)
    detect_button.click(
        fn=model.detect_and_visualize,
        inputs=[
            input_image,
            visualization_score_threshold,
        ],
        outputs=[
            detection_results,
            detection_visualization,
        ],
    )
    redraw_button.click(
        fn=model.visualize_detection_results,
        inputs=[
            input_image,
            detection_results,
            visualization_score_threshold,
        ],
        outputs=[detection_visualization],
    )
demo.queue(max_size=10).launch()