import os
from typing import Mapping

import gradio as gr
from PIL import Image

from monochrome import _DEFAULT_MONOCHROME_CKPT, _MONOCHROME_CKPTS, get_monochrome_score


def monochrome_detect(image: Image.Image, ckpt: str) -> Mapping[str, float]:
    return get_monochrome_score(image, ckpt)


if __name__ == '__main__':
    with gr.Blocks() as demo:
        with gr.Row():
            with gr.Column():
                gr_input_image = gr.Image(type='pil', label='Original Image')
                with gr.Row():
                    gr_model = gr.Dropdown(_MONOCHROME_CKPTS, value=_DEFAULT_MONOCHROME_CKPT, label='Model')

                gr_btn_submit = gr.Button(value='Detection', variant='primary')

            with gr.Column():
                gr_tags = gr.Label(label='Detection')

        gr_btn_submit.click(
            monochrome_detect,
            inputs=[gr_input_image, gr_model],
            outputs=[gr_tags],
        )
    demo.queue(os.cpu_count()).launch()