Update app.py

app.py

@@ -276,29 +276,18 @@ def create_cam_visualization_pil(cam, alpha=0.6, vis_threshold=0.2):
     # Normalize CAM to [0, 1]
     cam_norm = (cam_resized - cam_resized.min()) / (np.ptp(cam_resized) + 1e-8)
 
-    # Apply threshold mask
-    mask = cam_norm >= vis_threshold
-
     # Create heatmap using matplotlib colormap
     colormap = cm.get_cmap('jet')
-    heatmap_rgba = colormap(cam_norm)  # shape: (H, W, 4), values in [0, 1]
-    heatmap_rgb = (heatmap_rgba[:, :, :3] * 255).astype(np.uint8)
-
-    # Convert heatmap to PIL image
-    heatmap_pil = Image.fromarray(heatmap_rgb).convert("RGB")
+    cam_colored = colormap(cam_norm)[:, :, :3]  # RGB
+    cam_alpha = (cam_norm >= vis_threshold).astype(np.float32) * alpha  # Alpha mask
 
-    # Convert images to NumPy for blending
-    base_np = np.array(image_pil).astype(np.float32)
-    heat_np = np.array(heatmap_pil).astype(np.float32)
+    cam_rgba = np.dstack((cam_colored, cam_alpha))  # Shape: (H, W, 4)
+    cam_image = Image.fromarray((cam_rgba * 255).astype(np.uint8), mode="RGBA")
 
-    # Blend only where mask is True
-    blended_np = base_np.copy()
-    blended_np[mask] = base_np[mask] * (1 - alpha) + heat_np[mask] * alpha
-    blended_np = np.clip(blended_np, 0, 255).astype(np.uint8)
+    # Composite over original
+    composite = Image.alpha_composite(image_pil, cam_image)
 
-    # Convert back to PIL image
-    blended_img = Image.fromarray(blended_np)
-    return blended_img
+    return composite
 
 
 with gr.Blocks(css=".output-class { display: none; }") as demo: