Update visualization.py

visualization.py

@@ -45,20 +45,49 @@ def visualize_prediction(
 
 
 def visualize_attention_map(pil_img, attention_map):
+    # Get the attention map for the last layer
     attention_map = attention_map[-1].detach().cpu()
+    
+    # Get the number of heads
+    n_heads = attention_map.shape[1]
+    
+    # Calculate the average attention weight for each head
     avg_attention_weight = torch.mean(attention_map, dim=1).squeeze()
-    avg_attention_weight_resized = (
-        F.interpolate(
-            avg_attention_weight.unsqueeze(0).unsqueeze(0),
-            size=pil_img.size[::-1],
-            mode="bicubic",
-        )
-        .squeeze()
-        .numpy()
-    )
-
-    plt.imshow(pil_img)
-    plt.imshow(avg_attention_weight_resized, alpha=0.7, cmap="viridis")
-    plt.axis("off")
-    fig = plt.gcf()
+    
+    # Resize the attention map
+    resized_attention_weight = F.interpolate(
+        avg_attention_weight.unsqueeze(0).unsqueeze(0),
+        size=pil_img.size[::-1],
+        mode="bicubic",
+    ).squeeze().numpy()
+    
+    # Create a grid of subplots
+    fig, axes = plt.subplots(nrows=1, ncols=n_heads, figsize=(n_heads*4, 4))
+    
+    # Loop through the subplots and plot the attention for each head
+    for i, ax in enumerate(axes.flat):
+        ax.imshow(pil_img)
+        ax.imshow(attention_map[0,i,:,:].squeeze(), alpha=0.7, cmap="viridis")
+        ax.set_title(f"Head {i+1}")
+        ax.axis("off")
+    
+    plt.tight_layout()
+    
     return fig2img(fig)
+    # attention_map = attention_map[-1].detach().cpu()
+    # avg_attention_weight = torch.mean(attention_map, dim=1).squeeze()
+    # avg_attention_weight_resized = (
+    #     F.interpolate(
+    #         avg_attention_weight.unsqueeze(0).unsqueeze(0),
+    #         size=pil_img.size[::-1],
+    #         mode="bicubic",
+    #     )
+    #     .squeeze()
+    #     .numpy()
+    # )
+
+    # plt.imshow(pil_img)
+    # plt.imshow(avg_attention_weight_resized, alpha=0.7, cmap="viridis")
+    # plt.axis("off")
+    # fig = plt.gcf()
+    # return fig2img(fig)