from typing import *

import gradio as gr

from predict import predict_fn
from utils import populate_examples

description = """
    Anomaly detection models are trained with only <span style="color:lime;font-weight:bold">normal</span> images,
    and aimed to segment <span style="color:red;font-weight:bold">anomalies (deviations)</span> in input images.

    Scroll to bottom of this demo for a list of pretrained examples.
"""

def launch():
    input_image = gr.Image(label="Input image")
    threshold = gr.Slider(value=1, step=0.1, label="Threshold")
    devices = gr.Radio(
        label="Device",
        choices=["AUTO", "CPU", "GPU"],
        value="CPU",
        interactive=False
    )
    model = gr.Text(label="Model", interactive=False)

    output_image = gr.Image(label="Output image")
    output_heatmap = gr.Image(label="Heatmap")

    intf = gr.Interface(
        title="Anomaly Detection",
        description=description,
        fn=predict_fn,
        inputs=[input_image, threshold, devices, model],
        outputs=[output_image, output_heatmap],
        examples=populate_examples(),
        allow_flagging="never"
    )
    intf.launch()

if __name__ == "__main__":
    launch()