Update visualization.py

visualization.py

@@ -21,17 +21,10 @@ def visualize_prediction(
     fig, ax = plt.subplots(figsize=(12, 12))
     ax.imshow(pil_img)
     if display_mask and mask is not None:
-        # Convert the mask image to a numpy array
         mask_arr = np.asarray(mask)
-        
-        # Create a new mask with white objects and black background
         new_mask = np.zeros_like(mask_arr)
         new_mask[mask_arr > 0] = 255
-        
-        # Convert the numpy array back to a PIL Image
         new_mask = Image.fromarray(new_mask)
-        
-        # Display the new mask as a semi-transparent overlay
         ax.imshow(new_mask, alpha=0.5, cmap='viridis')
        
     colors = COLORS * 100
@@ -62,26 +55,20 @@ 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()
     
-    # 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")
@@ -91,20 +78,3 @@ def visualize_attention_map(pil_img, attention_map):
     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)