debugging

app.py

@@ -79,6 +79,7 @@ def serialize_images(processed_map):
     for i in range(processed_map.shape[0]):
         plt.imshow(processed_map[i].numpy())
         plt.title(f"Attention head: {i}", fontsize=14)
+        plt.axis("off")
         plt.savefig(fname="attention_map_{i}.png")
 
 
@@ -95,7 +96,8 @@ def generate_class_attn_map(image, block_id=0):
     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"))
+    all_attn_img_paths = sorted(glob.glob("*.png"))
+    print(f"Number of images: {len(all_attn_img_paths)}")
     return all_attn_img_paths
 
 
@@ -115,4 +117,4 @@ iface = gr.Interface(
     cache_examples=True,
     examples=[["./bird.png", 0]],
 )
-iface.launch()
+iface.launch(debug=True)