neuro-sam 0.1.0__py3-none-any.whl

neuro_sam/punet/prob_unet_deepd3.py

@@ -0,0 +1,431 @@
+"""
+Probabilistic U-Net components integrated with DeepD3 model.
+
+Contains encoder, Gaussian latent spaces, and feature combination modules.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.distributions import Normal, Independent, kl
+import numpy as np
+
+from deepd3_model import DeepD3Model
+from unet_blocks import *
+from utils import init_weights, init_weights_orthogonal_normal, l2_regularisation
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+
+class Encoder(nn.Module):
+    """
+    Convolutional encoder with downsampling blocks.
+    
+    Args:
+        input_channels: Number of input channels
+        num_filters: List of filter counts per block
+        no_convs_per_block: Number of convolutions per block
+        initializers: Weight initialization config
+        segm_channels: Number of segmentation channels to concatenate
+        padding: Whether to use padding
+        posterior: Whether this is a posterior encoder (concatenates segmentation)
+    """
+    def __init__(
+        self, 
+        input_channels, 
+        num_filters, 
+        no_convs_per_block, 
+        initializers, 
+        segm_channels,
+        padding=True, 
+        posterior=False
+    ):
+        super(Encoder, self).__init__()
+        self.contracting_path = nn.ModuleList()
+        self.input_channels = input_channels
+        self.num_filters = num_filters
+
+        if posterior:
+            self.input_channels += segm_channels
+
+        layers = []
+        for i in range(len(self.num_filters)):
+            input_dim = self.input_channels if i == 0 else output_dim
+            output_dim = num_filters[i]
+            
+            if i != 0:
+                layers.append(nn.AvgPool2d(kernel_size=2, stride=2, padding=0, ceil_mode=True))
+            
+            layers.append(nn.Conv2d(input_dim, output_dim, kernel_size=3, padding=int(padding)))
+            layers.append(nn.BatchNorm2d(output_dim)) 
+            layers.append(nn.ReLU(inplace=True))
+
+            for _ in range(no_convs_per_block - 1):
+                layers.append(nn.Conv2d(output_dim, output_dim, kernel_size=3, padding=int(padding)))
+                layers.append(nn.BatchNorm2d(output_dim))
+                layers.append(nn.ReLU(inplace=True))
+
+        self.layers = nn.Sequential(*layers)
+        self.layers.apply(init_weights)
+
+    def forward(self, input):
+        return self.layers(input)
+
+
+class AxisAlignedConvGaussian(nn.Module):
+    """
+    Convolutional network that parametrizes a Gaussian distribution 
+    with diagonal covariance matrix.
+    
+    Args:
+        input_channels: Number of input channels
+        num_filters: List of filter counts per block
+        no_convs_per_block: Number of convolutions per block
+        latent_dim: Dimensionality of latent space
+        initializers: Weight initialization config
+        segm_channels: Number of segmentation channels
+        posterior: Whether this is posterior (uses segmentation) or prior
+    """
+    def __init__(
+        self, 
+        input_channels, 
+        num_filters, 
+        no_convs_per_block, 
+        latent_dim, 
+        initializers, 
+        segm_channels,
+        posterior=False
+    ):
+        super(AxisAlignedConvGaussian, self).__init__()
+        self.input_channels = input_channels
+        self.channel_axis = 1
+        self.num_filters = num_filters
+        self.no_convs_per_block = no_convs_per_block
+        self.latent_dim = latent_dim
+        self.segm_channels = segm_channels
+        self.posterior = posterior
+        self.name = 'Posterior' if self.posterior else 'Prior'
+        
+        self.encoder = Encoder(
+            self.input_channels, 
+            self.num_filters, 
+            self.no_convs_per_block, 
+            initializers, 
+            self.segm_channels,
+            posterior=self.posterior
+        )
+        self.conv_layer = nn.Conv2d(num_filters[-1], 2 * self.latent_dim, (1, 1), stride=1)
+        self.show_img = 0
+        self.show_seg = 0
+        self.show_concat = 0
+        self.show_enc = 0
+        self.sum_input = 0
+
+        nn.init.kaiming_normal_(self.conv_layer.weight, mode='fan_in', nonlinearity='relu')
+        nn.init.normal_(self.conv_layer.bias)
+
+    def forward(self, input, segm=None):
+        """
+        Forward pass through encoder to latent distribution.
+        
+        Args:
+            input: Input image
+            segm: Segmentation mask (for posterior only)
+            
+        Returns:
+            Multivariate normal distribution with diagonal covariance
+        """
+        if segm is not None:
+            self.show_img = input
+            self.show_seg = segm
+            input = torch.cat((input, segm), dim=1)
+            self.show_concat = input
+            self.sum_input = torch.sum(input)
+
+        encoding = self.encoder(input)
+        self.show_enc = encoding
+
+        encoding = torch.mean(encoding, dim=2, keepdim=True)
+        encoding = torch.mean(encoding, dim=3, keepdim=True)
+
+        mu_log_sigma = self.conv_layer(encoding)
+
+        mu_log_sigma = torch.squeeze(mu_log_sigma, dim=2)
+        mu_log_sigma = torch.squeeze(mu_log_sigma, dim=2)
+
+        mu = mu_log_sigma[:, :self.latent_dim]
+        log_sigma = mu_log_sigma[:, self.latent_dim:]
+        log_sigma = torch.clamp(log_sigma, min=-10, max=10)
+
+        dist = Independent(Normal(loc=mu, scale=torch.exp(log_sigma + 1e-6)), 1)
+        return dist
+
+
+class Fcomb(nn.Module):
+    """
+    Feature combination module.
+    
+    Combines latent sample with U-Net features via 1x1 convolutions.
+    
+    Args:
+        num_filters: Filter configuration
+        latent_dim: Latent space dimensionality
+        num_output_channels: Number of output channels
+        num_classes: Number of output classes
+        no_convs_fcomb: Number of 1x1 convolutions
+        initializers: Weight initialization config
+        use_tile: Whether to tile latent sample to match spatial dimensions
+    """
+    def __init__(
+        self, 
+        num_filters, 
+        latent_dim, 
+        num_output_channels, 
+        num_classes, 
+        no_convs_fcomb, 
+        initializers, 
+        use_tile=True
+    ):
+        super(Fcomb, self).__init__()
+        self.num_channels = num_output_channels
+        self.num_classes = num_classes
+        self.channel_axis = 1
+        self.spatial_axes = [2, 3]
+        self.num_filters = num_filters
+        self.latent_dim = latent_dim
+        self.use_tile = use_tile
+        self.no_convs_fcomb = no_convs_fcomb
+        self.name = 'Fcomb'
+
+        if self.use_tile:
+            layers = []
+
+            layers.append(nn.Conv2d(
+                self.num_filters[0] + self.latent_dim, 
+                self.num_filters[0], 
+                kernel_size=1
+            ))
+            layers.append(nn.ReLU(inplace=True))
+
+            for _ in range(no_convs_fcomb - 2):
+                layers.append(nn.Conv2d(
+                    self.num_filters[0], 
+                    self.num_filters[0], 
+                    kernel_size=1
+                ))
+                layers.append(nn.ReLU(inplace=True))
+
+            self.layers = nn.Sequential(*layers)
+            self.last_layer = nn.Conv2d(self.num_filters[0], self.num_classes, kernel_size=1)
+
+            if initializers['w'] == 'orthogonal':
+                self.layers.apply(init_weights_orthogonal_normal)
+                self.last_layer.apply(init_weights_orthogonal_normal)
+            else:
+                self.layers.apply(init_weights)
+                self.last_layer.apply(init_weights)
+
+    def tile(self, a, dim, n_tile):
+        """
+        Tile tensor along specified dimension.
+        Mimics TensorFlow's tf.tile behavior.
+        """
+        init_dim = a.size(dim)
+        repeat_idx = [1] * a.dim()
+        repeat_idx[dim] = n_tile
+        a = a.repeat(*(repeat_idx))
+        order_index = torch.LongTensor(
+            np.concatenate([init_dim * np.arange(n_tile) + i for i in range(init_dim)])
+        ).to(device)
+        if a.device.type == 'mps':
+            return torch.index_select(a.cpu(), dim, order_index.cpu()).to(a.device)
+        return torch.index_select(a, dim, order_index)
+
+    def forward(self, feature_map, z):
+        """
+        Combine feature map with latent sample.
+        
+        Args:
+            feature_map: Feature map from U-Net [B, C, H, W]
+            z: Latent sample [B, latent_dim]
+            
+        Returns:
+            Combined output logits
+        """
+        if self.use_tile:
+            z = torch.unsqueeze(z, 2)
+            z = self.tile(z, 2, feature_map.shape[self.spatial_axes[0]])
+            z = torch.unsqueeze(z, 3)
+            z = self.tile(z, 3, feature_map.shape[self.spatial_axes[1]])
+
+            feature_map = torch.cat((feature_map, z), dim=self.channel_axis)
+            output = self.layers(feature_map)
+            return self.last_layer(output)
+
+
+class ProbabilisticUnet(nn.Module):
+    """
+    Probabilistic U-Net integrated with custom DeepD3 dual-decoder architecture.
+    
+    This is a basic version with single latent space. For dual latent spaces,
+    see ProbabilisticUnetDualLatent in prob_unet_with_tversky.py.
+    """
+    def __init__(
+        self, 
+        input_channels=1, 
+        num_classes=1, 
+        num_filters=[32, 64, 128, 192], 
+        latent_dim=6, 
+        no_convs_fcomb=4, 
+        beta=1
+    ):
+        super(ProbabilisticUnet, self).__init__()
+        self.input_channels = input_channels
+        self.num_classes = num_classes
+        self.num_filters = num_filters
+        self.latent_dim = latent_dim
+        self.no_convs_per_block = 3
+        self.no_convs_fcomb = no_convs_fcomb
+        self.initializers = {'w': 'he_normal', 'b': 'normal'}
+        self.beta = beta
+        self.z_prior_sample = 0
+
+        self.unet = DeepD3Model(
+            in_channels=self.input_channels,
+            base_filters=self.num_filters[0],
+            num_layers=len(self.num_filters),
+            activation="swish",
+            use_batchnorm=True,
+            apply_last_layer=False
+        ).to(device)
+
+        self.fcomb_dendrites = Fcomb(
+            self.num_filters, 
+            self.latent_dim, 
+            self.input_channels, 
+            self.num_classes, 
+            self.no_convs_fcomb, 
+            {'w': 'orthogonal', 'b': 'normal'}, 
+            use_tile=True
+        ).to(device)
+        
+        self.fcomb_spines = Fcomb(
+            self.num_filters, 
+            self.latent_dim, 
+            self.input_channels, 
+            self.num_classes, 
+            self.no_convs_fcomb, 
+            {'w': 'orthogonal', 'b': 'normal'}, 
+            use_tile=True
+        ).to(device)
+
+        self.prior = AxisAlignedConvGaussian(
+            self.input_channels, 
+            self.num_filters, 
+            self.no_convs_per_block, 
+            self.latent_dim, 
+            self.initializers, 
+            posterior=False,
+            segm_channels=1
+        ).to(device)
+        
+        self.posterior = AxisAlignedConvGaussian(
+            self.input_channels, 
+            self.num_filters, 
+            self.no_convs_per_block, 
+            self.latent_dim, 
+            self.initializers, 
+            posterior=True,
+            segm_channels=2
+        ).to(device)
+
+    def forward(self, patch, segm, training=True):
+        """
+        Forward pass through prior/posterior and U-Net.
+        """
+        if training:
+            self.posterior_latent_space = self.posterior.forward(patch, segm)
+        
+        self.prior_latent_space = self.prior.forward(patch)
+
+        dendrite_features, spine_features = self.unet(patch)
+        self.dendrite_features = dendrite_features
+        self.spine_features = spine_features
+
+    def sample(self, testing=False):
+        """
+        Sample segmentation by fusing latent with U-Net features.
+        """
+        if not testing:
+            z_prior = self.prior_latent_space.rsample()
+        else:
+            z_prior = self.prior_latent_space.sample()
+        
+        self.z_prior_sample = z_prior
+
+        dendrites = self.fcomb_dendrites(self.dendrite_features, z_prior)
+        spines = self.fcomb_spines(self.spine_features, z_prior)
+        
+        return dendrites, spines
+
+    def reconstruct(self, use_posterior_mean=False, calculate_posterior=False, z_posterior=None, training=True):
+        """
+        Reconstruct segmentation from latent space.
+        """
+        if self.posterior_latent_space is not None:
+            if use_posterior_mean:
+                z_posterior = self.posterior_latent_space.loc
+            elif calculate_posterior:
+                z_posterior = self.posterior_latent_space.rsample()
+        else:
+            z_posterior = self.prior_latent_space.rsample()
+
+        dendrites = self.fcomb_dendrites(self.dendrite_features, z_posterior)
+        spines = self.fcomb_spines(self.spine_features, z_posterior)
+        
+        return dendrites, spines
+
+    def kl_divergence(self, analytic=True, calculate_posterior=False, z_posterior=None):
+        """
+        Calculate KL divergence between posterior and prior.
+        """
+        if analytic:
+            kl_div = kl.kl_divergence(self.posterior_latent_space, self.prior_latent_space)
+        else:
+            if calculate_posterior:
+                z_posterior = self.posterior_latent_space.rsample()
+            log_posterior_prob = self.posterior_latent_space.log_prob(z_posterior)
+            log_prior_prob = self.prior_latent_space.log_prob(z_posterior)
+            kl_div = log_posterior_prob - log_prior_prob
+        
+        return kl_div
+
+    def elbo(self, segm_d, segm_s, analytic_kl=True, reconstruct_posterior_mean=False):
+        """
+        Calculate evidence lower bound (negative log-likelihood).
+        """
+        criterion = nn.BCEWithLogitsLoss(reduction='sum')
+        z_posterior = self.posterior_latent_space.rsample()
+
+        self.kl = torch.mean(
+            self.kl_divergence(analytic=analytic_kl, calculate_posterior=False, z_posterior=z_posterior)
+        )
+        
+        dendrites_rec, spines_rec = self.reconstruct(
+            use_posterior_mean=reconstruct_posterior_mean,
+            calculate_posterior=False, 
+            z_posterior=z_posterior
+        )
+
+        segm_dendrites = segm_d
+        segm_spines = segm_s
+
+        loss_dendrites = criterion(dendrites_rec, segm_dendrites)
+        loss_spines = criterion(spines_rec, segm_spines)
+        reconstruction_loss = loss_dendrites + loss_spines
+
+        epsilon = 1e-7
+        self.reconstruction_loss = torch.sum(reconstruction_loss + epsilon)
+        self.mean_reconstruction_loss = torch.mean(reconstruction_loss + epsilon)
+
+        return -(self.reconstruction_loss + self.beta * self.kl)