Upload app.py

app.py

@@ -0,0 +1,57 @@
+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()