import gradio as gr
import tensorflow as tf
import tensorflow_hub as hub
from PIL import Image

import utils

_RESOLUTION = 224
_MODEL_PATH = "gs://cait-tf/cait_xxs24_224"


def get_model() -> tf.keras.Model:
    """Initiates a tf.keras.Model from TF-Hub."""
    inputs = tf.keras.Input((_RESOLUTION, _RESOLUTION, 3))
    hub_module = hub.KerasLayer(_MODEL_PATH)

    logits, sa_atn_score_dict, ca_atn_score_dict = hub_module(inputs)

    return tf.keras.Model(
        inputs, [logits, sa_atn_score_dict, ca_atn_score_dict]
    )


_MODEL = get_model()


def show_plot(image):
    """Function to be called when user hits submit on the UI."""
    _, preprocessed_image = utils.preprocess_image(
        image, "deit_tiny_patch16_224"
    )
    _, _, ca_atn_score_dict = _MODEL.predict(preprocessed_image)

    result_first_block = utils.get_cls_attention_map(
        image, ca_atn_score_dict, block_key="ca_ffn_block_0_att"
    )
    result_second_block = utils.get_cls_attention_map(
        image, ca_atn_score_dict, block_key="ca_ffn_block_1_att"
    )
    return Image.fromarray(result_first_block), Image.fromarray(
        result_second_block
    )


title = "Generate Class Attention Plots"
article = "Class attention maps as investigated in [Going deeper with Image Transformers](https://arxiv.org/abs/2103.17239) (Touvron et al.)."

iface = gr.Interface(
    show_plot,
    inputs=gr.inputs.Image(type="pil", label="Input Image"),
    outputs="image",
    title=title,
    article=article,
    allow_flagging="never",
    examples=[["./butterfly.jpg"]],
)
iface.launch()