torch_em.model.probabilistic_unet

  1"""@private
  2"""
  3
  4# This code is based on:
  5#   1. The original TensorFlow implementation: https://github.com/SimonKohl/probabilistic_unet
  6#   2. PyTorch adaptation from https://github.com/stefanknegt/Probabilistic-Unet-Pytorch
  7#   3. PhiSeg's benchmarking script from https://github.com/annawundram/glaucoma-diagnosis-pipeline.
  8
  9from typing import Optional, List
 10
 11import torch
 12import torch.nn as nn
 13import torch.nn.functional as F
 14from torch.distributions import Normal, Independent, kl
 15
 16from torch_em.model import UNet2d
 17from torch_em.loss.dice import DiceLoss
 18from torch_em.model.unet import get_norm_layer
 19
 20
 21def init_weights(m: nn.Module) -> None:
 22    """@private
 23    """
 24    if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):
 25        nn.init.kaiming_normal_(m.weight, mode='fan_in', nonlinearity='relu')
 26        if m.bias is not None:
 27            nn.init.trunc_normal_(m.bias, std=0.001)
 28
 29
 30def init_weights_orthogonal_normal(m: nn.Module) -> None:
 31    """@private
 32    """
 33    if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):
 34        nn.init.orthogonal_(m.weight)
 35        if m.bias is not None:
 36            nn.init.trunc_normal_(m.bias, std=0.001)
 37
 38
 39class Encoder(nn.Module):
 40    """Convolutional encoder for the prior and posterior networks.
 41
 42    Stacks len(num_filters) blocks of no_convs_per_block Conv-Norm-ReLU layers with average
 43    pooling between blocks. For the posterior, the segmentation mask is concatenated with the
 44    image along the channel axis before encoding.
 45
 46    Args:
 47        input_channels: Number of input image channels.
 48        num_filters: Number of filters for each encoder block.
 49        no_convs_per_block: Number of Conv-Norm-ReLU layers per block.
 50        posterior: If True, expects a concatenated segmentation mask as input.
 51        num_classes: Number of segmentation channels appended when posterior=True.
 52        norm: Normalization type applied after each convolution. None disables normalisation.
 53    """
 54    def __init__(
 55        self,
 56        input_channels: int,
 57        num_filters: List[int],
 58        no_convs_per_block: int,
 59        posterior: bool = False,
 60        num_classes: Optional[int] = None,
 61        norm: Optional[str] = None,
 62    ) -> None:
 63        super().__init__()
 64
 65        self.input_channels = input_channels
 66        self.num_filters = num_filters
 67
 68        if posterior:
 69            # To accommodate for the mask concatenated at the channel axis, increase input_channels.
 70            assert num_classes is not None
 71            self.input_channels += num_classes
 72
 73        layers = []
 74        output_dim = None
 75
 76        for i in range(len(self.num_filters)):
 77            # First block: input_channels -> num_filters[i]; subsequent: prev_output -> num_filters[i].
 78            input_dim = self.input_channels if i == 0 else output_dim
 79            output_dim = num_filters[i]
 80
 81            if i != 0:
 82                layers.append(nn.AvgPool2d(kernel_size=2, stride=2, padding=0, ceil_mode=True))
 83
 84            layers.append(nn.Conv2d(input_dim, output_dim, kernel_size=3))
 85            if norm is not None:
 86                layers.append(get_norm_layer(norm, 2, output_dim))
 87            layers.append(nn.ReLU(inplace=True))
 88
 89            for _ in range(no_convs_per_block - 1):
 90                layers.append(nn.Conv2d(output_dim, output_dim, kernel_size=3))
 91                if norm is not None:
 92                    layers.append(get_norm_layer(norm, 2, output_dim))
 93                layers.append(nn.ReLU(inplace=True))
 94
 95        self.layers = nn.Sequential(*layers)
 96        self.layers.apply(init_weights)
 97
 98    def forward(self, x: torch.Tensor) -> torch.Tensor:
 99        return self.layers(x)
100
101
102class AxisAlignedConvGaussian(nn.Module):
103    """Convolutional network that outputs a diagonal-covariance Gaussian distribution.
104
105    Encodes the input (and optionally a segmentation mask) into a global feature vector,
106    then predicts mu and log-sigma for each latent dimension. Returns an Independent(Normal)
107    distribution over the latent space.
108
109    Args:
110        input_channels: Number of input image channels.
111        num_filters: Number of filters per encoder block.
112        no_convs_per_block: Number of convolutions per encoder block.
113        latent_dim: Dimensionality of the latent space.
114        posterior: If True, encodes both image and segmentation (posterior q(z|x,y)).
115        num_classes: Number of segmentation channels when posterior=True.
116        use_onehot: If True, converts integer class labels to one-hot before concatenation.
117        num_raters: Number of segmentation maps per image for the joint posterior.
118        norm: Normalization type passed to the encoder. None disables normalisation.
119    """
120    def __init__(
121        self,
122        input_channels: int,
123        num_filters: List[int],
124        no_convs_per_block: int,
125        latent_dim: int,
126        posterior: bool = False,
127        num_classes: Optional[int] = None,
128        use_onehot: bool = False,
129        num_raters: int = 1,
130        norm: Optional[str] = None,
131    ) -> None:
132        super().__init__()
133        self.latent_dim = latent_dim
134        self.num_classes = num_classes
135        self.use_onehot = use_onehot
136
137        self.encoder = Encoder(
138            input_channels,
139            num_filters,
140            no_convs_per_block,
141            posterior=posterior,
142            num_classes=None if num_classes is None else num_classes * num_raters,
143            norm=norm,
144        )
145
146        self.conv_layer = nn.Conv2d(num_filters[-1], 2 * self.latent_dim, (1, 1), stride=1)
147        # Orthogonal init + truncated-normal bias per the original paper's training details.
148        nn.init.orthogonal_(self.conv_layer.weight, gain=1)
149        nn.init.trunc_normal_(self.conv_layer.bias, std=0.001)
150
151    def forward(self, patch: torch.Tensor, segm: Optional[torch.Tensor] = None) -> Independent:
152        # Posterior: encode segm and concatenate with image along channel axis.
153        # One-hot encodes class-index labels to remove spurious ordinal relationships.
154        # Centering (- 0.5) keeps inputs zero-mean in both the one-hot and raw binary cases.
155        if segm is not None:
156            if self.use_onehot:
157                segm = F.one_hot(segm.long(), self.num_classes).movedim(-1, 2).flatten(1, 2).float() - 0.5
158            else:
159                segm = segm.float() - 0.5
160            patch = torch.cat((patch, segm), dim=1)
161
162        encoding = self.encoder(patch)
163
164        # Global average pool to (B, C, 1, 1), then squeeze spatial dims to (B, C).
165        encoding = encoding.mean(dim=(2, 3), keepdim=True)
166        mu_log_sigma = self.conv_layer(encoding).squeeze(3).squeeze(2)
167
168        mu = mu_log_sigma[:, :self.latent_dim]
169        log_sigma = mu_log_sigma[:, self.latent_dim:]
170
171        # This is a multivariate normal with diagonal covariance matrix sigma
172        # https://github.com/pytorch/pytorch/pull/11178
173        dist = Independent(Normal(loc=mu, scale=torch.exp(log_sigma)), 1)
174        return dist
175
176
177class Fcomb(nn.Module):
178    """A sequence of 1x1 convolutions that combines the UNet feature map with a latent sample.
179
180    Broadcasts z to the spatial size of the feature map, concatenates along the channel axis,
181    and applies no_convs_fcomb pointwise convolutions to produce segmentation logits.
182
183    Args:
184        num_filters: Filter counts from the UNet encoder; num_filters[0] sets the fcomb width.
185        latent_dim: Dimensionality of the latent sample z.
186        num_output_channels: Number of output segmentation channels.
187        no_convs_fcomb: Total number of 1x1 conv layers (including the final projection).
188    """
189    def __init__(
190        self,
191        num_filters: List[int],
192        latent_dim: int,
193        num_output_channels: int,
194        no_convs_fcomb: int,
195    ) -> None:
196        super().__init__()
197
198        layers = []
199        layers.append(nn.Conv2d(num_filters[0] + latent_dim, num_filters[0], kernel_size=1))
200        layers.append(nn.ReLU(inplace=True))
201
202        for _ in range(no_convs_fcomb - 2):
203            layers.append(nn.Conv2d(num_filters[0], num_filters[0], kernel_size=1))
204            layers.append(nn.ReLU(inplace=True))
205
206        self.layers = nn.Sequential(*layers)
207        self.last_layer = nn.Conv2d(num_filters[0], num_output_channels, kernel_size=1)
208
209        self.layers.apply(init_weights_orthogonal_normal)
210        self.last_layer.apply(init_weights_orthogonal_normal)
211
212    def forward(self, feature_map: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
213        """Combine UNet feature map with a latent sample to produce a segmentation.
214
215        Args:
216            feature_map: UNet decoder output of shape (B, C, H, W).
217            z: Latent sample of shape (B, latent_dim), broadcast to (B, latent_dim, H, W)
218                before concatenation with feature_map.
219        """
220        H, W = feature_map.shape[2], feature_map.shape[3]
221        z = z.view(z.shape[0], z.shape[1], 1, 1).expand(-1, -1, H, W)
222        feature_map = torch.cat((feature_map, z), dim=1)
223        return self.last_layer(self.layers(feature_map))
224
225
226class ProbabilisticUNet(nn.Module):
227    """Network implementation for the Probabilistic UNet of Kohl et al. (https://arxiv.org/abs/1806.05034).
228    This generative segmentation heuristic uses UNet combined with a conditional variational
229    autoencoder enabling to efficiently produce an unlimited number of plausible hypotheses.
230
231    Labels have shape (B, R, H, W), where R is num_raters. Each rater provides a binary mask
232    or a multiclass map of integer class IDs. The joint posterior receives all raters as input channels.
233    Multiclass maps become R * C one-hot channels. Reconstruction loss averages over raters.
234    Prior samples have shape (B, C, H, W).
235
236    Args:
237        input_channels: Number of channels in the image (1 for grayscale and 3 for RGB). The default is set to 1.
238        output_channels: Number of channels to predict. The default is set to 1.
239        num_raters: Number of annotators per image for binary or multiclass labels. The default is set to 1.
240        num_filters: Number of filters per encoder level. The default is set to [32, 64, 128, 192].
241        latent_dim: Dimension of the latent space. The default is set to 6.
242        no_convs_per_block: Number of convolutions per block in the prior/posterior encoder. The default is set to 3.
243        no_convs_fcomb: Number of convolutions in the feature combination module. The default is set to 4.
244        norm: Normalization type for the prior and posterior encoder. The default is set to 'InstanceNorm'.
245        beta: Weighting factor for the KL divergence term in the ELBO (loss = reconstruction + beta * KL).
246            beta=1.0 is the standard VAE objective with no extra regularization. The original paper
247            used beta=10.0 for the LIDC-IDRI multi-rater chest X-ray task after task-specific tuning,
248            which over-regularizes the latent space for most other tasks. Set beta > 1 to encourage a
249            more structured latent space at the cost of reconstruction quality, or beta < 1 to
250            prioritize reconstruction. The default is set to 1.0.
251        consensus_masking: Whether to apply consensus masking in the reconstruction loss. The default is set to False.
252        rl_swap: Whether to use dice loss instead of BCE/CE for reconstruction. The default is set to False.
253        device: Device to place the model on. The default is set to None.
254    """
255
256    def __init__(
257        self,
258        input_channels: int = 1,
259        output_channels: int = 1,
260        num_raters: int = 1,
261        num_filters: List[int] = [32, 64, 128, 192],
262        latent_dim: int = 6,
263        no_convs_per_block: int = 3,
264        no_convs_fcomb: int = 4,
265        norm: Optional[str] = "InstanceNorm",
266        beta: float = 1.0,
267        consensus_masking: bool = False,
268        rl_swap: bool = False,
269        device: Optional[torch.device] = None,
270    ) -> None:
271        super().__init__()
272
273        if output_channels < 1 or num_raters < 1:
274            raise ValueError("output_channels and num_raters must be positive.")
275
276        self.input_channels = input_channels
277        self.output_channels = output_channels
278        self.num_raters = num_raters
279        self.num_filters = num_filters
280        self.latent_dim = latent_dim
281        self.no_convs_per_block = no_convs_per_block
282        self.no_convs_fcomb = no_convs_fcomb
283        self.beta = beta
284        self.consensus_masking = consensus_masking
285        self.rl_swap = rl_swap
286        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') if device is None else device
287
288        self.unet = UNet2d(
289            in_channels=self.input_channels,
290            out_channels=None,
291            depth=len(self.num_filters),
292            initial_features=num_filters[0],
293            norm=norm,
294        )
295
296        self.prior = AxisAlignedConvGaussian(
297            self.input_channels,
298            self.num_filters,
299            self.no_convs_per_block,
300            self.latent_dim,
301            norm=norm,
302        )
303
304        use_onehot = output_channels > 1
305        self.posterior = AxisAlignedConvGaussian(
306            self.input_channels,
307            self.num_filters,
308            self.no_convs_per_block,
309            self.latent_dim,
310            posterior=True,
311            num_classes=output_channels,
312            use_onehot=use_onehot,
313            num_raters=num_raters,
314            norm=norm,
315        )
316
317        self.fcomb = Fcomb(
318            self.num_filters,
319            self.latent_dim,
320            self.output_channels,
321            self.no_convs_fcomb,
322        )
323
324        if rl_swap:
325            self._criterion = DiceLoss()
326        elif output_channels == 1:
327            self._criterion = nn.BCEWithLogitsLoss(reduction="none")
328        else:
329            self._criterion = nn.CrossEntropyLoss(reduction="none")
330
331        self.to(self.device)
332
333    def _check_shape(self, patch: torch.Tensor) -> None:
334        spatial_shape = patch.shape[2:]
335        depth = len(self.num_filters)
336        factor = [2**depth] * len(spatial_shape)
337        if any(sh % fac != 0 for sh, fac in zip(spatial_shape, factor)):
338            msg = f"Invalid shape for Probabilistic U-Net: {spatial_shape} is not divisible by {factor}"
339            raise ValueError(msg)
340
341    def forward(self, patch: torch.Tensor, segm: Optional[torch.Tensor] = None) -> None:
342        """Run the image through the UNet and build the prior latent space.
343
344        Labels have shape (B, R, H, W). The joint posterior receives all R raters.
345        """
346        self._check_shape(patch)
347
348        if segm is not None:
349            self._check_labels(segm, patch)
350            self.posterior_latent_space = self.posterior(patch, segm)
351
352        self.prior_latent_space = self.prior(patch)
353        self.unet_features = self.unet(patch)
354
355    def _check_labels(self, segm: torch.Tensor, image: torch.Tensor) -> None:
356        expected_shape = (image.shape[0], self.num_raters, *image.shape[2:])
357        if segm.shape != expected_shape:
358            raise ValueError(f"Expected labels with shape {expected_shape}, got {tuple(segm.shape)}.")
359
360    def sample(self) -> torch.Tensor:
361        """Sample a segmentation from the prior and decode it through the UNet feature map.
362
363        Draws a latent vector z from the prior p(z|x) and passes it through fcomb together
364        with the UNet features to produce one segmentation hypothesis. Calling this N times
365        yields N diverse hypotheses for the same input image.
366        """
367        z_prior = self.prior_latent_space.sample()
368        return self.fcomb(self.unet_features, z_prior)
369
370    def reconstruct(
371        self,
372        use_posterior_mean: bool = False,
373        calculate_posterior: bool = False,
374        z_posterior: Optional[torch.Tensor] = None,
375    ) -> torch.Tensor:
376        """Decode a joint posterior sample into logits with shape (B, C, H, W).
377
378        Args:
379            use_posterior_mean: Use the posterior mean as z instead of sampling.
380            calculate_posterior: Draw a fresh reparametrized sample from the posterior.
381                Ignored when use_posterior_mean=True or z_posterior is provided.
382            z_posterior: Pre-computed posterior sample to decode. Used directly when
383                use_posterior_mean=False and calculate_posterior=False.
384        """
385        if use_posterior_mean:
386            z_posterior = self.posterior_latent_space.mean
387        else:
388            if calculate_posterior:
389                z_posterior = self.posterior_latent_space.rsample()
390        return self.fcomb(self.unet_features, z_posterior)
391
392    def kl_divergence(
393        self,
394        analytic: bool = True,
395        calculate_posterior: bool = False,
396        z_posterior: Optional[torch.Tensor] = None,
397    ) -> torch.Tensor:
398        """Compute KL(posterior || prior).
399
400        Args:
401            analytic: If True, compute the KL in closed form. If False, estimate via
402                a posterior sample (log q(z) - log p(z)).
403            calculate_posterior: Draw a fresh reparametrized sample when analytic=False.
404                Ignored when z_posterior is provided.
405            z_posterior: Pre-computed posterior sample for the Monte Carlo estimate.
406                Only used when analytic=False.
407        """
408        if analytic:
409            kl_div = kl.kl_divergence(self.posterior_latent_space, self.prior_latent_space)
410        else:
411            if calculate_posterior:
412                z_posterior = self.posterior_latent_space.rsample()
413            log_posterior_prob = self.posterior_latent_space.log_prob(z_posterior)
414            log_prior_prob = self.prior_latent_space.log_prob(z_posterior)
415            kl_div = log_posterior_prob - log_prior_prob
416        return kl_div
417
418    def elbo(
419        self,
420        segm: torch.Tensor,
421        consm: Optional[torch.Tensor] = None,
422        analytic_kl: bool = True,
423        reconstruct_posterior_mean: bool = False,
424    ) -> torch.Tensor:
425        """Compute the evidence lower bound -E[log p(y|x,z)] + beta * KL(q||p).
426
427        One joint posterior sample supplies the reconstruction and KL terms for each image.
428        Reconstruction loss averages over raters. The KL term contributes once per image.
429
430        Args:
431            segm: Binary labels or multiclass indices of shape (B, R, H, W).
432            consm: Optional shared mask (B, 1, H, W) or per-rater mask (B, R, H, W).
433                Applied as a multiplicative weight when consensus_masking=True.
434            analytic_kl: If True, compute the KL divergence in closed form.
435            reconstruct_posterior_mean: If True, decode the posterior mean instead of a sample.
436        """
437
438        self._check_labels(segm, self.unet_features)
439        mask = consm if self.consensus_masking else None
440        if mask is not None:
441            shared_shape = (segm.shape[0], 1, *segm.shape[2:])
442            if mask.shape != segm.shape and mask.shape != shared_shape:
443                raise ValueError(f"Expected consensus mask with shape {shared_shape} or {tuple(segm.shape)}.")
444            mask = mask.expand_as(segm)
445
446        z_posterior = self.posterior_latent_space.rsample()
447
448        kl_term = torch.mean(
449            self.kl_divergence(analytic=analytic_kl, calculate_posterior=False, z_posterior=z_posterior)
450        )
451
452        reconstruction = self.reconstruct(
453            use_posterior_mean=reconstruct_posterior_mean,
454            calculate_posterior=False,
455            z_posterior=z_posterior
456        )
457        losses = []
458        for rater in range(self.num_raters):
459            target = segm[:, rater:rater + 1]
460            rater_mask = None if mask is None else mask[:, rater:rater + 1]
461            losses.append(self._reconstruction_loss(reconstruction, target, rater_mask))
462        reconstruction_loss = torch.stack(losses).mean()
463        return -(reconstruction_loss + self.beta * kl_term)
464
465    def _reconstruction_loss(
466        self, reconstruction: torch.Tensor, segm: torch.Tensor, mask: Optional[torch.Tensor]
467    ) -> torch.Tensor:
468        use_dice = self.rl_swap
469        use_bce = isinstance(self._criterion, nn.BCEWithLogitsLoss)
470        if use_dice:
471            if self.output_channels == 1:
472                prediction = reconstruction.sigmoid()
473                target = segm.float()
474            else:
475                prediction = reconstruction.softmax(dim=1)
476                target = F.one_hot(segm.squeeze(1).long(), self.output_channels).movedim(-1, 1).float()
477            if mask is not None:
478                # Mask probabilities before Dice reduces the spatial dimensions.
479                prediction = prediction * mask
480                target = target * mask
481            reconstruction_loss = self._criterion(prediction, target)
482            if mask is not None:
483                reconstruction_loss = reconstruction_loss * mask.any()
484        else:
485            segm_t = segm.float() if use_bce else segm.squeeze(1).long()
486            reconstruction_loss = self._criterion(reconstruction, segm_t)
487            if mask is not None:
488                reconstruction_loss = reconstruction_loss * (mask if use_bce else mask.squeeze(1))
489
490        if use_dice:
491            return reconstruction_loss
492
493        # Sum over spatial dims, mean over batch - keeps loss scale independent of batch size.
494        reconstruction_loss = reconstruction_loss.sum(dim=tuple(range(1, reconstruction_loss.dim()))).mean()
495        return reconstruction_loss