Update app.py

app.py

@@ -244,13 +244,8 @@ def create_cam_visualization_pil(image_pil, cam, alpha=0.6, vis_threshold=0.2):
     w, h = image_pil.size
     size = max(w, h)
 
-    # Resize CAM to match image
-    cam_resized = np.array(Image.fromarray(cam).resize((size, size), resample=Image.Resampling.BILINEAR))
-
     # Normalize CAM to [0, 1]
-    cam_norm = (cam_resized - cam_resized.min()) / (np.ptp(cam_resized) + 1e-8)
-
-    cam_norm = transforms.CenterCrop((h, w))(cam_norm)
+    cam_norm = (cam - cam.min()) / (np.ptp(cam) + 1e-8)
 
     # Create heatmap using matplotlib colormap
     colormap = cm.get_cmap('jet')
@@ -258,7 +253,13 @@ def create_cam_visualization_pil(image_pil, cam, alpha=0.6, vis_threshold=0.2):
     cam_alpha = (cam_norm >= vis_threshold).astype(np.float32) * alpha  # Alpha mask
 
     cam_rgba = np.dstack((cam_colored, cam_alpha))  # Shape: (H, W, 4)
-    cam_image = Image.fromarray((cam_rgba * 255).astype(np.uint8), mode="RGBA")
+    
+    # Resize CAM to match image
+    cam_resized = np.array(Image.fromarray(cam_rgba).resize((size, size), resample=Image.Resampling.BILINEAR))
+
+    cam_resized = transforms.CenterCrop((h, w))(torch.tensor(cam_resized)).numpy()
+
+    cam_image = Image.fromarray((cam_resized * 255).astype(np.uint8), mode="RGBA")
 
     # Composite over original
     composite = Image.alpha_composite(image_pil, cam_image)