import gradio as gr
from utils import load_model, load_image_from_url, inference


TITLE = """
<p align="right">
<img alt="Static Badge" src="https://img.shields.io/badge/SEA.AI-beta-blue">
</p>
<div align="center">
<h1>RGB Detection Demo</h1>
Give it a try! Upload an image or enter a URL to an image and click submit.
</div>
"""

NOTICE = """
Don't like what you see? Help us improve the model by flagging
incorrect predictions. Click on `Flag` to submit the image to
us for review and we'll use it to improve the A.I.
"""


model = load_model("SEA-AI/yolov5n6-RGB", img_size=1280)
model.conf = 0.25
model.iou = 0.4
model.max_det = 100
model.agnostic = True  # NMS class-agnostic

# This callback will be used to flag images
callback = gr.CSVLogger()

css = """
h1 {
    text-align: center;
    display: block;
}
"""

with gr.Blocks(css=css) as demo:
    gr.Markdown(value=TITLE)

    with gr.Row():
        with gr.Column():
            img_input = gr.Image(label="input", interactive=True)
            img_url = gr.Textbox(
                lines=1,
                placeholder="or enter URL to image here",
                label="input_url",
                show_label=False,
            )
            with gr.Row():
                clear = gr.ClearButton()
                submit = gr.Button("Submit", variant="primary")
        with gr.Column():
            img_output = gr.Image(
                label="output", interactive=False, show_share_button=True
            )
            flag = gr.Button("Flag", visible=False)
            notice = gr.Markdown(value=NOTICE, visible=False)

    gr.Examples(
        examples=[
            "https://images.pexels.com/photos/273886/pexels-photo-273886.jpeg?auto=compress&cs=tinysrgb&w=1260&h=750&dpr=2",
            "https://images.pexels.com/photos/913111/pexels-photo-913111.jpeg?auto=compress&cs=tinysrgb&w=1260&h=750&dpr=2",
            "https://images.pexels.com/photos/88517/pexels-photo-88517.jpeg?auto=compress&cs=tinysrgb&w=1260&h=750&dpr=2",
        ],
        inputs=img_input,
        outputs=img_output,
        fn=lambda image: inference(model, image),
        cache_examples=True,
    )

    # add components to clear
    clear.add([img_input, img_url, img_output])

    # event listeners
    img_url.change(load_image_from_url, [img_url], img_input)
    submit.click(lambda *args: inference(model, args), [img_input], img_output)

    @img_output.change(inputs=[img_output], outputs=[flag, notice])
    def show_hide(img_output):
        visible = img_output is not None
        return {
            flag: gr.Button("Flag", visible=visible),
            notice: gr.Markdown(value=NOTICE, visible=visible),
        }

    # This needs to be called prior to the first call to callback.flag()
    callback.setup([img_input, img_url, img_output], "flagged")

    # We can choose which components to flag (in this case, we'll flag all)
    flag.click(
        lambda *args: callback.flag(args),
        [img_input, img_url, img_output],
        None,
        preprocess=False,
    ).then(lambda: gr.Info("Thank you for contributing!"))


demo.launch()