fix: output grid and caching.

app.py

@@ -1,3 +1,5 @@
+import glob
+
 import gradio as gr
 import matplotlib.pyplot as plt
 import timm
@@ -6,7 +8,6 @@ from timm import create_model
 from timm.models.layers import PatchEmbed
 from torchvision.models.feature_extraction import create_feature_extractor
 from torchvision.transforms import functional as F
-import glob
 
 CAIT_MODEL = create_model("cait_xxs24_224.fb_dist_in1k", pretrained=True).eval()
 TRANSFORM = timm.data.create_transform(
@@ -73,6 +74,7 @@ def generate_plot(processed_map):
     fig.tight_layout()
     return fig
 
+
 def serialize_images(processed_map):
     """Serializes attention maps."""
     print(f"Number of maps: {processed_map.shape[0]}")
@@ -94,7 +96,7 @@ def generate_class_attn_map(image, block_id=0):
 
     block_key = f"blocks_token_only.{block_id}.attn.softmax"
     processed_cls_attn_map = get_cls_attention_map(image_tensor, out, block_key)
-    
+
     serialize_images(processed_cls_attn_map)
     all_attn_img_paths = sorted(glob.glob("attention_map_*.png"))
     print(f"Number of images: {len(all_attn_img_paths)}")
@@ -107,10 +109,10 @@ article = "Class attention maps as investigated in [Going deeper with Image Tran
 iface = gr.Interface(
     generate_class_attn_map,
     inputs=[
-        gr.inputs.Image(type="pil", label="Input Image"),
+        gr.Image(type="pil", label="Input Image"),
         gr.Slider(0, 1, value=0, step=1, label="Block ID", info="Transformer Block ID"),
     ],
-    outputs=gr.Gallery().style(grid=[2], height="auto"),
+    outputs=gr.Gallery().style(columns=2, height="auto", object_fit="scale-down"),
     title=title,
     article=article,
     allow_flagging="never",

gradio_cached_examples/14/log.csv

@@ -1,2 +0,0 @@
-output,flag,username,timestamp
-/Users/sayakpaul/Downloads/class-attention-map/gradio_cached_examples/14/output/24ed4fad-3279-4814-ba76-b4c411c673a0,,,2023-06-11 11:55:03.515035

gradio_cached_examples/14/output/24ed4fad-3279-4814-ba76-b4c411c673a0/captions.json

@@ -1 +0,0 @@
-{"/Users/sayakpaul/Downloads/class-attention-map/gradio_cached_examples/14/output/24ed4fad-3279-4814-ba76-b4c411c673a0/76269f58a7c390191fe41c6e016b4904749cd456/attention_map_i.png": null}