Update app.py

app.py

@@ -41,6 +41,8 @@ class DoubleConv(nn.Module):
     def forward(self, x):
         return self.conv(x)
 
+
+# Also, make sure your AttentionBlock.forward() returns the attention map:
 class AttentionBlock(nn.Module):
     def __init__(self, F_g, F_l, F_int):
         super(AttentionBlock, self).__init__()
@@ -67,7 +69,7 @@ class AttentionBlock(nn.Module):
         x1 = self.W_x(x)
         psi = self.relu(g1 + x1)
         psi = self.psi(psi)
-        return x * psi, psi
+        return x * psi, psi  # Return both attended features AND attention map
 
 class AttentionUNET(nn.Module):
     def __init__(self, in_channels=1, out_channels=1, features=[32, 64, 128, 256]):
@@ -261,17 +263,34 @@ def apply_tta(model, input_tensor):
     return torch.mean(torch.stack(predictions), dim=0)
 
 def generate_attention_heatmap(attention_maps):
-    """Generate attention heatmap"""
+    """Generate attention heatmap - Fixed version"""
     if not attention_maps:
         return np.zeros((256, 256, 3))
     
-    # Combine attention maps
-    combined_att = torch.mean(torch.stack(attention_maps), dim=0).squeeze().cpu().numpy()
-    combined_att = cv2.resize(combined_att, (256, 256))
+    # Resize all attention maps to the same size (256x256) before combining
+    resized_maps = []
+    target_size = (256, 256)
+    
+    for att_map in attention_maps:
+        # Convert to numpy and squeeze
+        att_np = att_map.squeeze().cpu().numpy()
+        
+        # Resize to target size
+        att_resized = cv2.resize(att_np, target_size)
+        resized_maps.append(att_resized)
+    
+    # Now we can safely average the maps since they're all the same size
+    combined_att = np.mean(resized_maps, axis=0)
+    
+    # Normalize to [0, 1]
     combined_att = (combined_att - combined_att.min()) / (combined_att.max() - combined_att.min() + 1e-8)
+    
+    # Apply colormap
     heatmap = cv2.applyColorMap((combined_att * 255).astype(np.uint8), cv2.COLORMAP_JET)
+    
     return heatmap
 
+
 def analyze_image(image, ground_truth, filename):
     """Main analysis function"""
     if model is None: