A bayesian Segmentation Model That Flags Its Own Uncertain Pixels

Full proposed model code in PyTorch

The implementation below reproduces the framework end to end. It includes a Monte Carlo dropout BCNN backbone, user prior modulation through element wise masking, transformer based multi head query fusion across image and user tokens, a conditional VAE with a user guided latent prior, the combined ELBO and supervised loss, both inference modes described above including the iterative test time adaptation loop, and the pixel level uncertainty map. A runnable smoke test with synthetic data is included at the bottom so the whole pipeline can be sanity checked without any real dataset.

# ==============================================================================
# Bayesian Multiclass Segmentation for Remote Sensing
# Paper: https://doi.org/10.1109/TGRS.2026.3670205
# Authors: Yeongsu Kim, Haeyun Lee, Seo-Yeon Choi, Kyungsu Lee
# Journal: IEEE Trans. Geoscience and Remote Sensing, Vol. 64, 2026
# PyTorch 2.4+ implementation with CUDA support
# ==============================================================================
from __future__ import annotations
import math
from typing import List, Optional, Tuple
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor


# --- SECTION 1: Bayesian CNN Backbone (BCNN) with MC Dropout ----------------

class MCDropout(nn.Module):
    """
    Monte Carlo Dropout, active during both training AND inference.
    Standard nn.Dropout turns off at eval(); this module does not,
    enabling weight uncertainty sampling (Eq. 7-8 of the paper).

    Each forward pass with a different random seed samples a different
    weight realisation W_l ~ q(W_l | D), approximating Bayesian inference.
    """
    def __init__(self, p: float = 0.1) -> None:
        super().__init__()
        self.p = p

    def forward(self, x: Tensor) -> Tensor:
        # Always active, training=True forces dropout even at inference
        return F.dropout(x, p=self.p, training=True)


class BCNNBlock(nn.Module):
    """
    One Bayesian convolutional block: Conv2d -> BatchNorm -> ReLU -> MCDropout.
    Stacking these implements f_BCNN(I; {W_l}) where W_l are stochastic.
    """
    def __init__(self, c_in: int, c_out: int,
                 k: int = 3, s: int = 1, p: int = 1,
                 drop_p: float = 0.1) -> None:
        super().__init__()
        self.conv = nn.Conv2d(c_in, c_out, k, s, p, bias=False)
        self.bn   = nn.BatchNorm2d(c_out)
        self.act  = nn.ReLU(inplace=True)
        self.drop = MCDropout(drop_p)

    def forward(self, x: Tensor) -> Tensor:
        return self.drop(self.act(self.bn(self.conv(x))))


class BCNNEncoder(nn.Module):
    """
    Bayesian CNN encoder backbone.

    Implements the stochastic feature extraction:
        F = f_BCNN(I; {W_l}),  W_l ~ q(W_l | D)  for all l   [Eq. 7]

    The predictive distribution p(F|I) is approximated by M forward passes:
        {F^(m)}_{m=1}^M,  W_l^(m) ~ q(W_l | D)              [Eq. 8]

    Epistemic uncertainty = pixelwise variance across M passes:
        Var_m[F^(m)]

    Parameters
    ----------
    in_channels   : number of input spectral channels (C=3 for RGB)
    feature_dim   : output feature channel dimension
    drop_p        : MC dropout probability (applied every forward pass)
    """
    def __init__(self, in_channels: int = 3,
                 feature_dim: int = 256,
                 drop_p: float = 0.1) -> None:
        super().__init__()
        self.feature_dim = feature_dim
        self.encoder = nn.Sequential(
            BCNNBlock(in_channels, 64,  drop_p=drop_p),
            nn.MaxPool2d(2),
            BCNNBlock(64,  128, drop_p=drop_p),
            nn.MaxPool2d(2),
            BCNNBlock(128, 256, drop_p=drop_p),
            BCNNBlock(256, feature_dim, drop_p=drop_p),
        )

    def forward(self, x: Tensor) -> Tensor:
        """Single stochastic forward pass, sample one W_l realisation."""
        return self.encoder(x)

    def mc_forward(self, x: Tensor, M: int = 20) -> Tuple[Tensor, Tensor]:
        """
        Monte Carlo sampling: run M stochastic forward passes.

        Parameters
        ----------
        x : (B, C, H, W) input image batch
        M : number of MC samples (paper uses M=20)

        Returns
        -------
        F_mean : (B, feature_dim, H', W') mean feature map
        F_var  : (B, feature_dim, H', W') epistemic uncertainty (variance)
        """
        samples = torch.stack([self.forward(x) for _ in range(M)], dim=0)
        F_mean = samples.mean(dim=0)
        F_var  = samples.var(dim=0, unbiased=True)
        return F_mean, F_var


# --- SECTION 2: User Prior Modulation (Eq. 9) --------------------------------

class UserPriorModulation(nn.Module):
    """
    Inject user correction masks as spatial priors into the feature map.

    For each user correction mask M_user^(n) in {0,1}^{H x W},
    the element-wise product selects BCNN features at correction regions:
        F_prior^(n) = F (x) M_user^(n)                          [Eq. 9]

    The masks are upsampled to match the feature map spatial resolution
    before masking. This mechanism allows user domain knowledge to be
    encoded directly as feature-space modulations rather than output
    constraints, a soft, probabilistic form of user guidance.

    Parameters
    ----------
    feature_dim : number of feature channels (must match BCNN output)
    """
    def __init__(self, feature_dim: int = 256) -> None:
        super().__init__()
        self.feature_dim = feature_dim

    def forward(self, F: Tensor,
                masks: Optional[List[Tensor]] = None) -> List[Tensor]:
        """
        Parameters
        ----------
        F     : (B, C, Hf, Wf) BCNN feature map
        masks : list of N binary tensors, each (B, 1, H, W) or (B, H, W)
                If None, returns empty list (no-user-prior scenario)

        Returns
        -------
        F_priors : list of N (B, C, Hf, Wf) user-prior feature maps
        """
        if masks is None or len(masks) == 0:
            return []

        B, C, Hf, Wf = F.shape
        priors = []
        for m in masks:
            if m.dim() == 2:
                m = m.unsqueeze(0).unsqueeze(0)
            elif m.dim() == 3:
                m = m.unsqueeze(1)
            m_rs = F.interpolate(m.float(), size=(Hf, Wf), mode='nearest')
            priors.append(F * m_rs)
        return priors


# --- SECTION 3: Token Embedding and Transformer Query Fusion -----------------

class TokenEmbedding(nn.Module):
    """
    Project feature maps into query token sequences for the transformer.
    Implements h_F and h_prior_emb in Eq. 13.

    Spatially flattens the feature map (H' x W' -> N_pix tokens)
    after a 1x1 projection to embedding dimension d.
    """
    def __init__(self, in_dim: int, embed_dim: int = 256) -> None:
        super().__init__()
        self.proj = nn.Conv2d(in_dim, embed_dim, 1)

    def forward(self, F: Tensor) -> Tensor:
        """
        Parameters
        ----------
        F : (B, C, H, W) feature map

        Returns
        -------
        tokens : (B, H*W, embed_dim) flattened token sequence
        """
        x = self.proj(F)
        B, D, H, W = x.shape
        return x.flatten(2).permute(0, 2, 1)


class MultiHeadSelfAttention(nn.Module):
    """
    Multi-head self-attention over query tokens (Eq. 15 and 18).

    MHAttn(q, K, V) = sum_h softmax(q K_h^T / sqrt(d_k)) V_h

    All queries (image tokens + user-prior tokens) attend to each other,
    propagating user correction signals globally across the spatial graph.
    """
    def __init__(self, embed_dim: int = 256,
                 num_heads: int = 8, dropout: float = 0.1) -> None:
        super().__init__()
        self.num_heads = num_heads
        self.head_dim  = embed_dim // num_heads
        self.scale     = self.head_dim ** -0.5
        self.WQ = nn.Linear(embed_dim, embed_dim, bias=False)
        self.WK = nn.Linear(embed_dim, embed_dim, bias=False)
        self.WV = nn.Linear(embed_dim, embed_dim, bias=False)
        self.out = nn.Linear(embed_dim, embed_dim)
        self.drop = nn.Dropout(dropout)

    def forward(self, q: Tensor) -> Tensor:
        """
        Parameters
        ----------
        q : (B, M, embed_dim), M = N_img + N_user tokens

        Returns
        -------
        out : (B, M, embed_dim)
        """
        B, M, D = q.shape
        H, Dh = self.num_heads, self.head_dim

        def split_heads(x):
            return x.view(B, M, H, Dh).permute(0, 2, 1, 3)

        Q = split_heads(self.WQ(q))
        K = split_heads(self.WK(q))
        V = split_heads(self.WV(q))

        attn = (Q @ K.transpose(-2, -1)) * self.scale
        attn = self.drop(attn.softmax(dim=-1))
        out  = (attn @ V).permute(0, 2, 1, 3).reshape(B, M, D)
        return self.out(out)


class TransformerQueryFusionLayer(nn.Module):
    """One transformer layer: MHAttn + FFN + LayerNorm + residual."""
    def __init__(self, embed_dim: int = 256,
                 num_heads: int = 8, ffn_ratio: int = 4,
                 dropout: float = 0.1) -> None:
        super().__init__()
        self.norm1 = nn.LayerNorm(embed_dim)
        self.attn  = MultiHeadSelfAttention(embed_dim, num_heads, dropout)
        self.norm2 = nn.LayerNorm(embed_dim)
        self.ffn   = nn.Sequential(
            nn.Linear(embed_dim, embed_dim * ffn_ratio),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(embed_dim * ffn_ratio, embed_dim),
            nn.Dropout(dropout),
        )

    def forward(self, q: Tensor) -> Tensor:
        q = q + self.attn(self.norm1(q))
        q = q + self.ffn(self.norm2(q))
        return q


class TransformerQueryFusion(nn.Module):
    """
    Stack of L transformer layers that jointly refine image and user tokens.

    Implements Q^0 -> Q^* via L applications of Eq. 17:
        q_i^(l+1) = q_i^(l) + sum_j a_ij^(l) W_V q_j^(l)

    The final Q* = {q_i^*}_{i=1}^M encodes globally contextual,
    uncertainty-aware, and user-bias-aware representations.

    Parameters
    ----------
    embed_dim : token embedding dimension d
    num_heads : H attention heads
    num_layers: L transformer layers
    """
    def __init__(self, embed_dim: int = 256,
                 num_heads: int = 8, num_layers: int = 4,
                 dropout: float = 0.1) -> None:
        super().__init__()
        self.layers = nn.ModuleList([
            TransformerQueryFusionLayer(embed_dim, num_heads, dropout=dropout)
            for _ in range(num_layers)
        ])

    def forward(self, Q: Tensor) -> Tensor:
        """
        Parameters
        ----------
        Q : (B, M, embed_dim) initial query set Q^0 (Eq. 14)

        Returns
        -------
        Q_star : (B, M, embed_dim) refined query set Q* (Eq. 16)
        """
        for layer in self.layers:
            Q = layer(Q)
        return Q


# --- SECTION 4: Conditional VAE with User-Guided Prior -----------------------

class VAEPriorNet(nn.Module):
    """
    Parameterises the user-conditioned Gaussian prior (Eq. 20 and 22):
        p(z | I, {F_prior^(n)}) = N(z; mu_prior, Sigma_prior)
    where [mu_prior, Sigma_prior] = h_prior(GP(F), {GP(F_prior^(n))})

    GP(.) = global average pooling (sufficient statistic for the prior).
    Aggregated user features are concatenated with image features.
    """
    def __init__(self, feature_dim: int = 256,
                 latent_dim: int = 128,
                 max_user_priors: int = 10) -> None:
        super().__init__()
        self.max_n = max_user_priors
        self.mlp = nn.Sequential(
            nn.Linear(feature_dim * 2, 512),
            nn.ReLU(inplace=True),
            nn.Linear(512, latent_dim * 2),
        )
        self.user_agg = nn.Linear(feature_dim, feature_dim)

    def forward(self, F: Tensor,
                F_priors: Optional[List[Tensor]] = None) -> Tuple[Tensor, Tensor]:
        """
        Parameters
        ----------
        F        : (B, C, H, W) image feature map
        F_priors : list of N (B, C, H, W) user-prior feature maps (may be empty)

        Returns
        -------
        mu_prior  : (B, latent_dim)
        log_var_prior : (B, latent_dim)
        """
        gp_F = F.mean(dim=[2, 3])
        if F_priors and len(F_priors) > 0:
            user_stack = torch.stack([fp.mean(dim=[2, 3]) for fp in F_priors], dim=1)
            gp_user = user_stack.mean(dim=1)
            gp_user = self.user_agg(gp_user)
        else:
            gp_user = torch.zeros_like(gp_F)

        h = torch.cat([gp_F, gp_user], dim=1)
        params = self.mlp(h)
        mu, log_var = params.chunk(2, dim=-1)
        return mu, log_var


class VAEEncoder(nn.Module):
    """
    Variational posterior encoder: q_phi(z | I, {M_user^(n)}, Y).
    Takes the concatenation of image features and ground-truth mask embeddings
    during training to learn the approximate posterior.
    """
    def __init__(self, feature_dim: int = 256,
                 num_classes: int = 7,
                 latent_dim: int = 128) -> None:
        super().__init__()
        self.mask_embed = nn.Linear(num_classes, feature_dim)
        self.mlp = nn.Sequential(
            nn.Linear(feature_dim * 2, 512),
            nn.ReLU(inplace=True),
            nn.Linear(512, latent_dim * 2),
        )

    def forward(self, F: Tensor, Y_onehot: Tensor) -> Tuple[Tensor, Tensor]:
        """
        Parameters
        ----------
        F        : (B, C, H, W) feature map
        Y_onehot : (B, K, H, W) one-hot ground truth mask

        Returns
        -------
        mu_q  : (B, latent_dim)
        lv_q  : (B, latent_dim)  log variance
        """
        gp_F = F.mean(dim=[2, 3])
        gp_Y = Y_onehot.permute(0, 2, 3, 1).mean(dim=[1, 2])
        gp_Y = self.mask_embed(gp_Y)
        h = torch.cat([gp_F, gp_Y], dim=1)
        params = self.mlp(h)
        return params.chunk(2, dim=-1)


def reparameterise(mu: Tensor, log_var: Tensor) -> Tensor:
    """
    Reparameterisation trick: z = mu + sigma * eps,  eps ~ N(0, I)
    Enables gradient flow through the stochastic latent variable.
    """
    std = (0.5 * log_var).exp()
    eps = torch.randn_like(std)
    return mu + std * eps


class SegmentationDecoder(nn.Module):
    """
    Generative decoder: f_dec(Q*, z) -> Y_pred.          [Eq. 27]

    Takes the refined query representation Q* from the transformer
    and the sampled latent z, and decodes pixel-level class logits.
    The decoder combines global latent information with local spatial
    feature information via broadcast addition.

    Parameters
    ----------
    embed_dim    : transformer embedding dimension d
    latent_dim   : VAE latent dimension
    num_classes  : K land-cover categories
    feat_hw      : spatial resolution of the feature map
    """
    def __init__(self, embed_dim: int = 256,
                 latent_dim: int = 128,
                 num_classes: int = 7,
                 feat_hw: int = 128) -> None:
        super().__init__()
        self.feat_hw = feat_hw
        self.z_proj = nn.Linear(latent_dim, embed_dim)
        self.query_proj = nn.Linear(embed_dim, embed_dim)
        self.head = nn.Sequential(
            nn.Conv2d(embed_dim, 128, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.Upsample(scale_factor=4, mode='bilinear', align_corners=False),
            nn.Conv2d(128, num_classes, 1),
        )

    def forward(self, Q_star: Tensor, z: Tensor) -> Tensor:
        """
        Parameters
        ----------
        Q_star : (B, M, embed_dim)  transformer output queries
        z      : (B, latent_dim)    sampled latent variable

        Returns
        -------
        logits : (B, K, H, W)  pixel-level class logits
        """
        B, M, D = Q_star.shape
        H = W = self.feat_hw

        q_feat = self.query_proj(Q_star)
        n_spatial = H * W
        q_spatial = q_feat[:, :n_spatial, :]
        q_spatial = q_spatial.permute(0, 2, 1).reshape(B, D, H, W)

        z_bias = self.z_proj(z)[:, :, None, None]
        feat = q_spatial + z_bias

        return self.head(feat)


# --- SECTION 5: Full Bayesian Segmentation Model ------------------------------

class BayesSegNet(nn.Module):
    """
    Full Bayesian Multiclass Segmentation Framework.

    Architecture (Table I of the paper):
      1. Feature Extraction : F = f_BCNN(I; {W_l})                [Eq. 7]
      2. User Prior Modulation: F_prior^(n) = F (x) M_user^(n)     [Eq. 9]
      3. Token Embedding : T_F = h_F(F), T_prior = h_emb(F_prior)  [Eq. 13]
      4. Query Fusion    : Q* = Transformer(Q^0)                   [Eq. 16]
      5. Prior Parameterisation: [mu_prior, Sigma_prior] = h_prior(...)  [Eq. 20]
      6. Latent Sampling : z ~ N(mu_prior, Sigma_prior) or q_phi(z|...)
      7. Decoding        : Y_pred = f_dec(Q*, z)                   [Eq. 27]

    Parameters
    ----------
    num_classes      : K land-cover categories
    in_channels      : C spectral channels (3 for RGB)
    feature_dim      : BCNN output channels
    embed_dim        : transformer embedding dimension d
    latent_dim       : VAE latent space dimension
    num_heads        : H attention heads
    num_tf_layers    : L transformer layers
    mc_samples       : M Monte Carlo samples for uncertainty estimation
    drop_p           : MC dropout probability
    img_size         : assumed square input resolution (H=W)
    """
    def __init__(
        self,
        num_classes: int = 7,
        in_channels: int = 3,
        feature_dim: int = 256,
        embed_dim:   int = 256,
        latent_dim:  int = 128,
        num_heads:   int = 8,
        num_tf_layers: int = 4,
        mc_samples:  int = 20,
        drop_p:      float = 0.1,
        img_size:    int = 512,
    ) -> None:
        super().__init__()
        self.num_classes = num_classes
        self.mc_samples  = mc_samples
        feat_hw = img_size // 4

        self.bcnn       = BCNNEncoder(in_channels, feature_dim, drop_p)
        self.prior_mod  = UserPriorModulation(feature_dim)
        self.img_tok    = TokenEmbedding(feature_dim, embed_dim)
        self.user_tok   = TokenEmbedding(feature_dim, embed_dim)
        self.transformer = TransformerQueryFusion(embed_dim, num_heads, num_tf_layers)
        self.vae_prior  = VAEPriorNet(feature_dim, latent_dim)
        self.vae_enc    = VAEEncoder(feature_dim, num_classes, latent_dim)
        self.decoder    = SegmentationDecoder(embed_dim, latent_dim,
                                               num_classes, feat_hw)

    def forward(
        self,
        I: Tensor,
        user_masks: Optional[List[Tensor]] = None,
        Y_gt: Optional[Tensor] = None,
        training: bool = True,
    ) -> dict:
        """
        Full forward pass.

        Parameters
        ----------
        I          : (B, C, H, W) input image
        user_masks : list of N binary masks (B, H, W), user corrections
        Y_gt       : (B, H, W) ground truth for VAE encoder (training only)
        training   : if True, use posterior q_phi; else use prior p(z|I,U)

        Returns
        -------
        dict with keys: logits, mu_q, lv_q, mu_p, lv_p, z
        """
        F = self.bcnn(I)
        F_priors = self.prior_mod(F, user_masks)

        T_img = self.img_tok(F)
        T_usr = [self.user_tok(fp) for fp in F_priors]

        Q0_parts = [T_img] + T_usr
        Q0 = torch.cat(Q0_parts, dim=1)
        Q_star = self.transformer(Q0)

        mu_p, lv_p = self.vae_prior(F, F_priors)

        if training and Y_gt is not None:
            Y_onehot = F.one_hot(Y_gt.long(), self.num_classes)
            Y_onehot = Y_onehot.permute(0, 3, 1, 2).float()
            B, K, H_orig, W_orig = Y_onehot.shape
            H_feat, W_feat = F.shape[2], F.shape[3]
            Y_dn = F.interpolate(Y_onehot, size=(H_feat, W_feat), mode='nearest')
            mu_q, lv_q = self.vae_enc(F, Y_dn)
            z = reparameterise(mu_q, lv_q)
        else:
            mu_q, lv_q = mu_p, lv_p
            z = reparameterise(mu_p, lv_p)

        logits = self.decoder(Q_star, z)
        logits = F.interpolate(logits, size=I.shape[-2:],
                                mode='bilinear', align_corners=False)

        return {
            'logits': logits,
            'mu_q':  mu_q, 'lv_q': lv_q,
            'mu_p':  mu_p, 'lv_p': lv_p,
            'z':     z,
        }

    @torch.no_grad()
    def mc_predict(
        self,
        I: Tensor,
        user_masks: Optional[List[Tensor]] = None,
        M: Optional[int] = None,
    ) -> Tuple[Tensor, Tensor]:
        """
        Monte Carlo inference: M stochastic forward passes, mean + uncertainty.

        Implements Eqs. 27-28:
          S_bar = (1/M) sum_m S^(m)
          U(i,j) = (1/(M-1)) sum_m (S^(m)(i,j) - S_bar(i,j))^2

        Parameters
        ----------
        I          : (B, C, H, W) input image
        user_masks : list of N user correction masks
        M          : number of MC samples (default: self.mc_samples)

        Returns
        -------
        mean_pred  : (B, K, H, W) mean class probabilities
        uncertainty: (B, K, H, W) pixelwise epistemic uncertainty map
        """
        M = M or self.mc_samples
        samples = []
        for _ in range(M):
            out = self.forward(I, user_masks, training=False)
            prob = out['logits'].softmax(dim=1)
            samples.append(prob)

        stack = torch.stack(samples, dim=0)
        mean_pred   = stack.mean(dim=0)
        uncertainty = stack.var(dim=0, unbiased=True)
        return mean_pred, uncertainty


# --- SECTION 6: ELBO + Supervised Composite Loss ------------------------------

class BayesSegLoss(nn.Module):
    """
    Composite training loss (Eq. 24-26):
        L_total = L_sup + lambda_ELBO * L_ELBO

    L_sup  = pixelwise cross-entropy (or Dice + CE) over mean prediction
    L_ELBO = KL(q_phi(z|I,M_user,Y) || p(z|I,F_prior)) - E[log p_theta(Y|z,Q*)]

    Parameters
    ----------
    num_classes : K (for class weighting)
    lambda_elbo : balance weight lambda_ELBO between sup and ELBO losses
    """
    def __init__(self, num_classes: int = 7,
                 lambda_elbo: float = 0.1) -> None:
        super().__init__()
        self.ce  = nn.CrossEntropyLoss(ignore_index=255)
        self.lam = lambda_elbo

    def kl_divergence(self, mu_q: Tensor, lv_q: Tensor,
                       mu_p: Tensor, lv_p: Tensor) -> Tensor:
        """Analytical KL between two diagonal Gaussians."""
        kl = (0.5 * (
            lv_p - lv_q
            + (lv_q.exp() + (mu_q - mu_p).pow(2)) / lv_p.exp().clamp(1e-6)
            - 1
        )).sum(dim=-1).mean()
        return kl

    def dice_loss(self, pred: Tensor, target: Tensor,
                   smooth: float = 1.0) -> Tensor:
        """Soft Dice loss for class imbalance robustness."""
        B, K, H, W = pred.shape
        pred_flat   = pred.softmax(dim=1).view(B, K, -1)
        target_oh   = F.one_hot(target.long().clamp(0), K).permute(0, 3, 1, 2)
        target_flat = target_oh.float().view(B, K, -1)
        intersection = (pred_flat * target_flat).sum(-1)
        dice = 1 - (2 * intersection + smooth) / (
            pred_flat.sum(-1) + target_flat.sum(-1) + smooth
        )
        return dice.mean()

    def forward(self, out: dict, Y_gt: Tensor) -> Tuple[Tensor, dict]:
        """
        Parameters
        ----------
        out  : dict from BayesSegNet.forward() containing logits + VAE params
        Y_gt : (B, H, W) ground truth class indices

        Returns
        -------
        total_loss : scalar
        loss_dict  : breakdown for logging
        """
        logits = out['logits']

        ce_loss   = self.ce(logits, Y_gt.long())
        dice_loss = self.dice_loss(logits, Y_gt)
        L_sup = ce_loss + dice_loss

        kl = self.kl_divergence(
            out['mu_q'], out['lv_q'],
            out['mu_p'], out['lv_p']
        )

        recon = F.cross_entropy(logits, Y_gt.long(), ignore_index=255)
        L_ELBO = kl - recon

        total = L_sup + self.lam * L_ELBO

        return total, {
            'total': total.item(), 'ce': ce_loss.item(),
            'dice': dice_loss.item(), 'kl': kl.item(),
        }


# --- SECTION 7: Inference Pipeline ---------------------------------------------

@torch.no_grad()
def iterative_tta_inference(
    model: BayesSegNet,
    I: Tensor,
    user_masks_init: Optional[List[Tensor]] = None,
    t_max: int = 3,
    eps: float = 1e-4,
    mc_samples: int = 20,
    device: Optional[torch.device] = None,
) -> Tuple[Tensor, Tensor, List[Tensor]]:
    """
    Iterative Test-Time Adaptation inference loop (Eqs. 30-31).

    Supports both inference scenarios from the paper.
    Scenario 1, user provides correction masks upfront, single round.
    Scenario 2, user provides masks after inspecting uncertainty maps.

    Algorithm.
      1. Initial prediction (with or without user masks)
      2. Compute uncertainty map U(i,j)
      3. User inspects U and provides correction masks for high uncertainty regions
      4. Re-run prediction conditioned on new masks
      5. Repeat until the L1 change falls below eps, or t = T_max

    Parameters
    ----------
    model           : fitted BayesSegNet
    I               : (1, C, H, W) test image
    user_masks_init : initial user correction masks (may be None)
    t_max           : maximum refinement rounds T_max (paper: 3)
    eps             : convergence tolerance
    mc_samples      : M for uncertainty estimation

    Returns
    -------
    final_pred  : (1, K, H, W) final mean probability prediction
    uncertainty : (1, K, H, W) final pixelwise uncertainty map
    history     : list of T+1 prediction tensors (for visualisation)
    """
    if device is None:
        device = next(model.parameters()).device
    I = I.to(device)

    model.eval()
    current_masks = user_masks_init
    history = []
    prev_pred = None

    for t in range(t_max + 1):
        pred, uncertainty = model.mc_predict(I, current_masks, M=mc_samples)
        history.append(pred.clone())

        if t > 0 and prev_pred is not None:
            delta = (pred - prev_pred).abs().mean()
            if delta < eps:
                print(f"Converged at t={t} with delta {delta:.5f} under eps {eps}")
                break

        prev_pred = pred.clone()
        if t < t_max:
            unc_map = uncertainty.mean(dim=1, keepdim=True)
            threshold = unc_map.view(unc_map.shape[0], -1).quantile(0.85, dim=1)
            threshold = threshold[:, None, None, None]
            new_mask = (unc_map > threshold).float()
            current_masks = [new_mask.squeeze(1)]

    return pred, uncertainty, history


# --- SECTION 8: Training Loop ---------------------------------------------------

def train_one_epoch(
    model: BayesSegNet,
    loader,
    optimizer: torch.optim.Optimizer,
    criterion: BayesSegLoss,
    device: torch.device,
    scaler: Optional[torch.cuda.amp.GradScaler] = None,
) -> dict:
    """
    One full training epoch.

    During training, simulated user corrections are generated by sampling
    error prone regions from the ground truth, following the paper's
    protocol of N in {0, 1, 3, 5, 10} corrections drawn from misclassified
    regions.

    Parameters
    ----------
    model     : BayesSegNet
    loader    : DataLoader yielding (images, masks, [user_corrections]) tuples
    optimizer : AdamW with lr=1e-4 and polynomial decay (paper settings)
    criterion : BayesSegLoss
    device    : compute device
    scaler    : optional AMP grad scaler

    Returns
    -------
    dict with mean loss values for logging
    """
    model.train()
    totals = {'total': 0.0, 'ce': 0.0, 'dice': 0.0, 'kl': 0.0}
    n = 0

    for batch in loader:
        images = batch[0].to(device)
        masks  = batch[1].to(device)
        user_m = batch[2] if len(batch) > 2 else None
        if user_m is not None:
            user_m = [u.to(device) for u in user_m]

        optimizer.zero_grad()
        if scaler:
            with torch.cuda.amp.autocast():
                out = model(images, user_m, masks, training=True)
                loss, ld = criterion(out, masks)
            scaler.scale(loss).backward()
            scaler.unscale_(optimizer)
            nn.utils.clip_grad_norm_(model.parameters(), 5.0)
            scaler.step(optimizer); scaler.update()
        else:
            out = model(images, user_m, masks, training=True)
            loss, ld = criterion(out, masks)
            loss.backward()
            nn.utils.clip_grad_norm_(model.parameters(), 5.0)
            optimizer.step()

        for k in totals: totals[k] += ld.get(k, 0.0)
        n += 1

    return {k: v / max(n, 1) for k, v in totals.items()}


# --- SECTION 9: Evaluation Metrics -----------------------------------------------

def compute_miou(pred: Tensor, target: Tensor,
                  num_classes: int, ignore_idx: int = 255) -> float:
    """
    Mean Intersection over Union across K classes (Eq. 34).
    IoU_c = TP_c / (TP_c + FP_c + FN_c)
    mIOU  = (1/K) sum_c IoU_c

    Parameters
    ----------
    pred        : (B, K, H, W) logits or probabilities
    target      : (B, H, W)   ground truth class indices
    num_classes : K
    ignore_idx  : class index to exclude from evaluation
    """
    pred_cls = pred.argmax(dim=1).view(-1).cpu()
    true_cls = target.view(-1).cpu()
    valid = true_cls != ignore_idx
    pred_cls, true_cls = pred_cls[valid], true_cls[valid]

    ious = []
    for c in range(num_classes):
        tp = ((pred_cls == c) & (true_cls == c)).sum().float()
        fp = ((pred_cls == c) & (true_cls != c)).sum().float()
        fn = ((pred_cls != c) & (true_cls == c)).sum().float()
        denom = tp + fp + fn
        if denom > 0:
            ious.append((tp / denom).item())
    return float(sum(ious) / max(1, len(ious)))


# --- SECTION 10: Smoke Test -------------------------------------------------------

if __name__ == '__main__':
    print("============================================================")
    print("BayesSegNet smoke test, Bayesian BCNN plus VAE plus user priors")
    print("============================================================")
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Device: {device}")

    model = BayesSegNet(
        num_classes=7, in_channels=3, feature_dim=64,
        embed_dim=64, latent_dim=32, num_heads=4,
        num_tf_layers=2, mc_samples=5, img_size=64
    ).to(device)
    n_params = sum(p.numel() for p in model.parameters())
    print(f"Parameters: {n_params:,}")

    I = torch.randn(2, 3, 64, 64, device=device)
    Y = torch.randint(0, 7, (2, 64, 64), device=device)
    masks = [torch.randint(0, 2, (2, 64, 64), device=device).float()
             for _ in range(3)]

    model.train()
    out = model(I, masks, Y, training=True)
    print(f"Logits shape: {out['logits'].shape}")
    print(f"mu_q shape:   {out['mu_q'].shape}")
    print(f"z shape:      {out['z'].shape}")

    criterion = BayesSegLoss(num_classes=7, lambda_elbo=0.1)
    loss, ld = criterion(out, Y)
    print(f"Total loss: {loss.item():.4f} | CE={ld['ce']:.4f} | KL={ld['kl']:.4f}")

    model.eval()
    I_test = torch.randn(1, 3, 64, 64, device=device)
    mean_pred, unc = model.mc_predict(I_test, masks=None, M=5)
    print(f"MC pred shape: {mean_pred.shape}, unc shape: {unc.shape}")
    print(f"Mean uncertainty: {unc.mean():.4f}")

    final, unc_final, hist = iterative_tta_inference(
        model, I_test, user_masks_init=None,
        t_max=3, mc_samples=5, device=device
    )
    print(f"TTA history length: {len(hist)} rounds")

    Y_test = torch.randint(0, 7, (1, 64, 64))
    miou = compute_miou(final.cpu(), Y_test, num_classes=7)
    print(f"Smoke test mIoU on random data: {miou:.4f}")

    print("All checks passed.")