feat: full jet colormap composite heatmap (matches medical Grad-CAM style)

main.py

@@ -1161,50 +1161,71 @@ def _compute_xray_pseudo_cam(model, img_tensor: torch.Tensor, target_idx: int):
         handle.remove()
 
 
-def _heatmap_to_png_base64(heatmap, percentile: float = 70.0) -> str:
-    """Convert a (H, W) [0,1] numpy array to a jet-colored RGBA base64 PNG.
-
-    Applies percentile thresholding so only the most salient regions
-    (top ~30% of attention) are visible. This produces focused heatmaps
-    instead of uniform red tint when the raw CAM is diffuse.
+def _cam_to_jet_rgb(cam):
+    """Map a (H, W) [0, 1] heatmap to a (H, W, 3) uint8 RGB jet colormap.
+
+    Matches matplotlib's jet colormap:
+        0.00 → dark blue   (0, 0, 128)
+        0.25 → blue/cyan
+        0.50 → green
+        0.75 → yellow
+        1.00 → dark red    (128, 0, 0)
+    """
+    import numpy as np
 
-    Colormap: jet-like gradient
-        low  → transparent
-        mid  → green / yellow
-        high → red
+    cam = np.clip(cam, 0.0, 1.0)
+    r = np.clip(1.5 - np.abs(4.0 * cam - 3.0), 0, 1)
+    g = np.clip(1.5 - np.abs(4.0 * cam - 2.0), 0, 1)
+    b = np.clip(1.5 - np.abs(4.0 * cam - 1.0), 0, 1)
+
+    rgb = np.zeros((*cam.shape, 3), dtype=np.uint8)
+    rgb[..., 0] = (r * 255).astype(np.uint8)
+    rgb[..., 1] = (g * 255).astype(np.uint8)
+    rgb[..., 2] = (b * 255).astype(np.uint8)
+    return rgb
+
+
+def _composite_heatmap_on_xray(
+    cam,
+    xray_gray: "np.ndarray",
+    alpha: float = 0.5,
+    gamma: float = 0.7,
+) -> str:
+    """Composite a jet heatmap onto a grayscale X-ray and return a base64 PNG.
+
+    Produces the classic "medical Grad-CAM" look:
+      - Full jet colormap across the image (blue→cyan→green→yellow→red)
+      - Base X-ray always visible underneath via alpha blending
+      - Gamma < 1 boosts contrast of mid-range attention
+
+    Args:
+        cam: (H, W) float array in [0, 1]
+        xray_gray: (H, W) uint8 grayscale X-ray matching CAM dimensions
+        alpha: weight of the heatmap (0 = only X-ray, 1 = only heatmap)
+        gamma: gamma correction on the heatmap (<1 brightens mid values)
     """
     import numpy as np
 
-    # Percentile threshold — only show top (100 - percentile)% of attention
-    thresh = np.percentile(heatmap, percentile)
-    focused = np.clip(heatmap - thresh, 0, None)
-    if focused.max() > 0:
-        focused = focused / focused.max()
-    else:
-        focused = np.zeros_like(heatmap)
-
-    h, w = focused.shape
-    rgba = np.zeros((h, w, 4), dtype=np.uint8)
-
-    # Jet-like colormap computed in pure numpy
-    # 0.00 → blue (0,0,255)
-    # 0.25 → cyan (0,255,255)
-    # 0.50 → green (0,255,0)
-    # 0.75 → yellow (255,255,0)
-    # 1.00 → red (255,0,0)
-    r = np.clip(1.5 - np.abs(4 * focused - 3), 0, 1)
-    g = np.clip(1.5 - np.abs(4 * focused - 2), 0, 1)
-    b = np.clip(1.5 - np.abs(4 * focused - 1), 0, 1)
-
-    rgba[..., 0] = (r * 255).astype(np.uint8)
-    rgba[..., 1] = (g * 255).astype(np.uint8)
-    rgba[..., 2] = (b * 255).astype(np.uint8)
-    # Alpha ramps up quickly so low-attention areas are transparent
-    rgba[..., 3] = (np.sqrt(focused) * 220).astype(np.uint8)
-
-    overlay = Image.fromarray(rgba, mode="RGBA")
+    # Smooth and sharpen the CAM
+    cam = np.clip(cam, 0.0, 1.0)
+    if cam.max() > 0:
+        cam = (cam - cam.min()) / (cam.max() - cam.min() + 1e-8)
+    cam = np.power(cam, gamma)  # gamma correction for contrast
+
+    # Jet colormap
+    heatmap_rgb = _cam_to_jet_rgb(cam)  # (H, W, 3) uint8
+
+    # Grayscale X-ray → RGB
+    xray_rgb = np.stack([xray_gray, xray_gray, xray_gray], axis=-1).astype(np.float32)
+    heatmap_f = heatmap_rgb.astype(np.float32)
+
+    # Alpha blend: composite = alpha * heatmap + (1 - alpha) * xray
+    composite = alpha * heatmap_f + (1.0 - alpha) * xray_rgb
+    composite = np.clip(composite, 0, 255).astype(np.uint8)
+
+    overlay = Image.fromarray(composite, mode="RGB")
     buf = io.BytesIO()
-    overlay.save(buf, format="PNG")
+    overlay.save(buf, format="PNG", optimize=True)
     b64 = base64.b64encode(buf.getvalue()).decode("ascii")
     return f"data:image/png;base64,{b64}"
 
@@ -1270,7 +1291,15 @@ async def predict_xray(
                     cam = _compute_xray_pseudo_cam(model, tensor_cam, top_idx)
                 else:
                     cam = _compute_xray_gradcam(model, tensor_cam, top_idx)
-                heatmap_b64 = _heatmap_to_png_base64(cam)
+
+                # Reconstruct the 224x224 grayscale X-ray for compositing.
+                # torchxrayvision normalizes to [-1024, 1024]; invert the mapping.
+                xray_float = (arr[0] + 1024.0) / 2048.0 * 255.0
+                xray_uint8 = np.clip(xray_float, 0, 255).astype(np.uint8)
+
+                heatmap_b64 = _composite_heatmap_on_xray(
+                    cam, xray_uint8, alpha=0.5, gamma=0.7
+                )
             except Exception as cam_err:
                 print(f"Heatmap computation failed: {cam_err}")