import os

import gradio as gr
from imgutils.detect import detection_visualize

from detect import _ALL_MODELS, _DEFAULT_MODEL, detect_text


def _gr_detect_text(image, model: str, threshold: float):
    return detection_visualize(image, detect_text(image, model, threshold))


if __name__ == '__main__':
    with gr.Blocks() as demo:
        with gr.Row():
            with gr.Column():
                gr_face_input_image = gr.Image(type='pil', label='Original Image')
                gr_face_model = gr.Dropdown(_ALL_MODELS, value=_DEFAULT_MODEL, label='Model')
                with gr.Row():
                    gr_face_score_threshold = gr.Slider(0.0, 1.0, 0.05, label='Score Threshold')

                gr_face_submit = gr.Button(value='Submit', variant='primary')

            with gr.Column():
                gr_face_output_image = gr.Image(type='pil', label="Labeled")

            gr_face_submit.click(
                _gr_detect_text,
                inputs=[
                    gr_face_input_image, gr_face_model, gr_face_score_threshold,
                ],
                outputs=[gr_face_output_image],
            )

    demo.queue(os.cpu_count()).launch()