"""Counterfactual healthy generation for v9. Goal ---- Given an encoder latent that's been decomposed into (z_anatomy, z_tumor) by the causal SCM head, generate the *counterfactual healthy* image: x_healthy = D(z_anatomy, do(z_tumor = 0)) This is the "what would this scan look like with no tumor?" generation that gives clinicians (and the LLM reporter) a visual baseline. The residual between the input scan and the counterfactual is a clean visualization of the tumor mass. Architecture ------------ A small generative decoder (1-2 M parameters) that takes the recomposed latent (with z_tumor zeroed) and produces a low-resolution healthy reconstruction, then upsamples to image resolution. We use a U-Net-style upsampling block with optional VQ-VAE codebook for sharper textures. Two output modes: - "L1 reconstruction": direct pixel regression of the healthy image. Simpler, faster, but blurry on textures. - "perceptual + L1": adds a perceptual loss term (LPIPS) for sharper output. Heavier compute. Use this for final paper results. Loss components for training ---------------------------- On HEALTHY scans (mask is empty): - Reconstruction loss: counterfactual should equal the input (the healthy scan IS its own healthy version). L1 + perceptual. - Identity loss: z_tumor for a healthy scan should be near zero (nothing to "remove"). Driven by SCM disentanglement loss. On TUMOR scans (mask is non-empty): - Reconstruction in the NON-TUMOR region: counterfactual should equal input outside the tumor mask. - In the TUMOR region: counterfactual should "fill in" with surrounding healthy tissue (inpainting prior). Driven by the L1 loss with mask weighting. This is a self-supervised loss formulation (no paired "tumor / healthy" ground truth exists in BraTS). The SCM decomposition + masked reconstruction provides the supervision signal. For v9 brain-2D scope this is a minimal but working implementation. v10 should swap in a latent diffusion decoder (Stable-Diffusion-Med or similar) for state-of-the-art generative quality. """ from __future__ import annotations from typing import Optional import torch import torch.nn as nn import torch.nn.functional as F class CounterfactualHealthyDecoder(nn.Module): """Generates the healthy counterfactual image from causal latent. Forward signature: x_input: (B, C, H, W) original MRI (used as spatial reference) z_recomposed: (B, D) latent from CausalRecompose with z_tumor=0 mask_pred: (B, 1, H, W) predicted tumor mask (for inpainting) When None, the decoder simply reconstructs the whole image. Returns: x_healthy: (B, C, H, W) counterfactual healthy version """ def __init__(self, latent_dim: int = 256, image_size: int = 384, in_channels: int = 3, base_channels: int = 32): super().__init__() self.latent_dim = latent_dim self.image_size = image_size self.in_channels = in_channels # Project the latent to a small spatial feature map: latent -> (B, C, 16, 16). # 16x16 is small enough to be cheap but enough resolution for global structure. spatial = max(8, image_size // 24) # ~16 at 384, ~12 at 256 self.spatial = spatial self.latent_to_feat = nn.Linear(latent_dim, base_channels * spatial * spatial) # Upsampling blocks: 16 -> 32 -> 64 -> 128 -> ... -> image_size. # ConvTranspose with stride 2 doubles spatial size each block. upsamples = [] cur_size = spatial cur_ch = base_channels while cur_size < image_size: next_ch = max(base_channels // 2, base_channels - 8) upsamples.append(nn.Sequential( nn.ConvTranspose2d(cur_ch, next_ch, kernel_size=4, stride=2, padding=1), nn.GroupNorm(min(8, next_ch), next_ch), nn.SiLU(inplace=True), nn.Conv2d(next_ch, next_ch, kernel_size=3, padding=1), nn.GroupNorm(min(8, next_ch), next_ch), nn.SiLU(inplace=True), )) cur_ch = next_ch cur_size *= 2 self.upsamples = nn.ModuleList(upsamples) # Final projection to image channels with skip from x_input (lets # the decoder copy unchanged regions from input cheaply). self.input_proj = nn.Conv2d(in_channels, cur_ch, kernel_size=1) self.final = nn.Sequential( nn.Conv2d(cur_ch, cur_ch, kernel_size=3, padding=1), nn.GroupNorm(min(8, cur_ch), cur_ch), nn.SiLU(inplace=True), nn.Conv2d(cur_ch, in_channels, kernel_size=1), nn.Tanh(), # output in [-1, 1], scaled to image space by caller ) def forward(self, x_input: torch.Tensor, z_recomposed: torch.Tensor, mask_pred: Optional[torch.Tensor] = None) -> torch.Tensor: B = x_input.size(0) # Latent -> spatial feature h = self.latent_to_feat(z_recomposed).view(B, -1, self.spatial, self.spatial) # Upsample for block in self.upsamples: h = block(h) # Crop / resize to exact image size (in case spatial doubling overshoots). if h.shape[-1] != self.image_size: h = F.interpolate(h, size=(self.image_size, self.image_size), mode="bilinear", align_corners=False) # Skip-add input projection (lets the decoder pass through anatomy # cheaply; tumor region is "fixed up" via the latent path). x_proj = self.input_proj(x_input) h = h + x_proj # Final image generation x_healthy = self.final(h) # Tanh output in [-1, 1]. Caller rescales to MRI intensity space. return x_healthy def reconstruction_loss(self, x_input: torch.Tensor, x_healthy: torch.Tensor, tumor_mask: torch.Tensor, lambda_outside: float = 1.0, lambda_inside: float = 0.5) -> torch.Tensor: """Masked reconstruction loss for counterfactual training. - Outside the tumor mask: counterfactual should equal input (high weight, lambda_outside = 1.0 default). - Inside the tumor mask: counterfactual should smoothly "inpaint" with surrounding tissue (lower weight, lambda_inside = 0.5). For healthy scans (empty mask), only the outside term is active, which means the counterfactual must reconstruct the whole input (lossless on healthy). """ outside = (tumor_mask < 0.5).float() # 1 outside tumor, 0 inside inside = (tumor_mask >= 0.5).float() diff = (x_healthy - x_input).abs() loss_outside = (diff * outside).sum() / outside.sum().clamp_min(1.0) loss_inside = (diff * inside).sum() / inside.sum().clamp_min(1.0) return lambda_outside * loss_outside + lambda_inside * loss_inside def tumor_residual(x_input: torch.Tensor, x_healthy: torch.Tensor) -> torch.Tensor: """Compute the tumor residual = |input - counterfactual_healthy|. For a perfect counterfactual decoder, this residual is: - High inside the tumor region (where the input has tumor and the counterfactual has been "filled in" with healthy tissue) - Low elsewhere (where input == counterfactual) This residual is a CLEAN visualization of the tumor mass, independent of the segmenter's mask. Use it as an auxiliary anomaly signal that the conformal head can be calibrated on (see hyperbolic_conformal.py). """ return (x_input - x_healthy).abs().mean(dim=1, keepdim=True) # (B, 1, H, W) __all__ = ["CounterfactualHealthyDecoder", "tumor_residual"]