Update app.py

app.py

@@ -127,11 +127,17 @@ def vit_attention_rollout(outputs):
 # Overlay Helper
 # ---------------------------
 def overlay_attention(pil_img, attention_map):
+    # Convert attention map to heatmap
     heatmap = (attention_map * 255).astype(np.uint8)
     heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
     heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)
-    heatmap = Image.fromarray(heatmap).resize(pil_img.size)
-    blended = Image.blend(pil_img.convert("RGB"), Image.fromarray(heatmap), alpha=0.4)
+    
+    # Resize heatmap to match image size
+    heatmap = cv2.resize(heatmap, pil_img.size)
+    
+    # Convert to PIL and blend
+    heatmap_pil = Image.fromarray(heatmap)
+    blended = Image.blend(pil_img.convert("RGB"), heatmap_pil, alpha=0.4)
     return blended