Upload app.py

app.py

@@ -0,0 +1,66 @@
+import cv2
+import gradio as gr
+import numpy as np
+import tensorflow as tf
+from huggingface_hub.keras_mixin import from_pretrained_keras
+from PIL import Image
+
+import utils
+
+_RESOLUTION = 224
+
+
+def get_model() -> tf.keras.Model:
+    """Initiates a tf.keras.Model from HF Hub."""
+    inputs = tf.keras.Input((_RESOLUTION, _RESOLUTION, 3))
+    hub_module = from_pretrained_keras(
+        "probing-vits/cait_xxs24_224_classification"
+    )
+
+    logits, sa_atn_score_dict, ca_atn_score_dict = hub_module(
+        inputs, training=False
+    )
+
+    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."""
+    original_image, preprocessed_image = utils.preprocess_image(
+        image, _RESOLUTION
+    )
+    _, _, ca_atn_score_dict = _MODEL.predict(preprocessed_image)
+
+    # Compute the saliency map and superimpose.
+    result_first_block = utils.get_cls_attention_map(
+        image, ca_atn_score_dict, block_key="ca_ffn_block_0_att"
+    )
+    heatmap = cv2.applyColorMap(
+        np.uint8(255 * result_first_block), cv2.COLORMAP_CIVIDIS
+    )
+    heatmap = np.float32(heatmap) / 255
+
+    original_image = original_image / 255.0
+    saliency_map = heatmap + original_image
+    saliency_map = saliency_map / np.max(saliency_map)
+    return Image.fromarray(saliency_map)
+
+
+title = "Generate Class Saliency Plots"
+article = "Class saliency 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()