Update app.py

app.py

@@ -23,6 +23,22 @@ def get_model() -> tf.keras.Model:
 _MODEL = get_model()
 
 
+def plot(attentions: np.ndarray):
+	"""Plots the attention maps from individual attention heads."""
+    fig, axes = plt.subplots(nrows=1, ncols=4, figsize=(13, 13))
+	img_count = 0
+
+	for i in range(attentions.shape[-1]):
+	    if img_count < attentions.shape[-1]:
+	        axes[i].imshow(attentions[:, :, img_count])
+	        axes[i].title.set_text(f"Attention head: {img_count}")
+	        axes[i].axis("off")
+	        img_count += 1
+
+	fig.tight_layout()
+	return fig
+
+
 def show_plot(image):
     """Function to be called when user hits submit on the UI."""
     _, preprocessed_image = utils.preprocess_image(
@@ -36,9 +52,7 @@ def show_plot(image):
     result_second_block = utils.get_cls_attention_map(
         preprocessed_image, ca_atn_score_dict, block_key="ca_ffn_block_1_att"
     )
-    return Image.fromarray(result_first_block), Image.fromarray(
-        result_second_block
-    )
+    return plot(result_first_block), plot(result_second_block)
 
 
 title = "Generate Class Attention Plots"
@@ -47,7 +61,7 @@ article = "Class attention maps as investigated in [Going deeper with Image Tran
 iface = gr.Interface(
     show_plot,
     inputs=gr.inputs.Image(type="pil", label="Input Image"),
-    outputs="image",
+    outputs="plot",
     title=title,
     article=article,
     allow_flagging="never",