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

import utils

_RESOLUTION = 224
_MODEL_URL = "https://tfhub.dev/sayakpaul/deit_tiny_patch16_224/1"


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_URL)

    logits, attention_scores_dict = hub_module(
        inputs
    )  # Second output in the tuple is a dictionary containing attention scores.

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


_MODEL = get_model()


def show_rollout(image):
    """Function to be called when user hits submit on the UI."""
    _, preprocessed_image = utils.preprocess_image(
        image, "deit_tiny_patch16_224"
    )
    _, attention_scores_dict = _MODEL.predict(preprocessed_image)
    result = utils.attention_rollout_map(
        image, attention_scores_dict, "deit_tiny_patch16_224"
    )
    return Image.fromarray(result)


title = "Generate Attention Rollout Plots"
article = "Attention Rollout was proposed by [Abnar et al.](https://arxiv.org/abs/2005.00928) to quantify the information that flows through self-attention layers. In the original ViT paper ([Dosovitskiy et al.](https://arxiv.org/abs/2010.11929)), the authors use it to investigate the representations learned by ViTs. The model used in the backend is `deit_tiny_patch16_224`. For more details about it, refer [here](https://tfhub.dev/sayakpaul/collections/deit/1). DeiT was proposed by [Touvron et al.](https://arxiv.org/abs/2012.12877)"

iface = gr.Interface(
    show_rollout,
    inputs=gr.inputs.Image(type="pil", label="Input Image"),
    outputs="image",
    title=title,
    article=article,
    allow_flagging="never",
    # examples=[["./car.jpeg", "./bulbul.jpeg"]],
)
iface.launch()