Update app.py

app.py

@@ -1,4 +1,3 @@
-import cv2
 import gradio as gr
 import numpy as np
 import tensorflow as tf
@@ -14,13 +13,9 @@ _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"
-    )
+    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
-    )
+    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]
@@ -38,17 +33,15 @@ def show_plot(image):
     _, _, ca_atn_score_dict = _MODEL.predict(preprocessed_image)
 
     # Compute the saliency map and superimpose.
-    result_first_block = utils.get_cls_attention_map(
+    saliency_attention = utils.get_cls_attention_map(
         preprocessed_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)
+    fig = plt.figure()
+    plt.imshow(original_image.astype("int32"))
+    plt.imshow(saliency_attention.squeeze(), cmap="cividis", alpha=0.9)
+    plt.axis("off")
+    return fig
 
-    saliency_map = heatmap + original_image
-    saliency_map = np.clip(saliency_map, 0.0, 255.0).astype(np.uint8)
-    return Image.fromarray(saliency_map)
 
 
 title = "Generate Class Saliency Plots"
@@ -57,10 +50,10 @@ article = "Class saliency maps as investigated in [Going deeper with Image Trans
 iface = gr.Interface(
     show_plot,
     inputs=gr.inputs.Image(type="pil", label="Input Image"),
-    outputs="image",
+    outputs=gr.outputs.Plot(type="auto"),
     title=title,
     article=article,
     allow_flagging="never",
-    examples=[["./butterfly.jpg"]],
+    examples=[["./butterfly_cropped.png"]],
 )
-iface.launch()
+iface.launch(debug=True)