careamics 0.0.1__py3-none-any.whl → 0.0.3__py3-none-any.whl

careamics/models/unet.py

@@ -0,0 +1,443 @@
+"""
+UNet model.
+
+A UNet encoder, decoder and complete model.
+"""
+
+from typing import Any, List, Tuple, Union
+
+import torch
+import torch.nn as nn
+
+from ..config.support import SupportedActivation
+from .activation import get_activation
+from .layers import Conv_Block, MaxBlurPool
+
+
+class UnetEncoder(nn.Module):
+    """
+    Unet encoder pathway.
+
+    Parameters
+    ----------
+    conv_dim : int
+        Number of dimension of the convolution layers, 2 for 2D or 3 for 3D.
+    in_channels : int, optional
+        Number of input channels, by default 1.
+    depth : int, optional
+        Number of encoder blocks, by default 3.
+    num_channels_init : int, optional
+        Number of channels in the first encoder block, by default 64.
+    use_batch_norm : bool, optional
+        Whether to use batch normalization, by default True.
+    dropout : float, optional
+        Dropout probability, by default 0.0.
+    pool_kernel : int, optional
+        Kernel size for the max pooling layers, by default 2.
+    n2v2 : bool, optional
+        Whether to use N2V2 architecture, by default False.
+    groups : int, optional
+        Number of blocked connections from input channels to output
+        channels, by default 1.
+    """
+
+    def __init__(
+        self,
+        conv_dim: int,
+        in_channels: int = 1,
+        depth: int = 3,
+        num_channels_init: int = 64,
+        use_batch_norm: bool = True,
+        dropout: float = 0.0,
+        pool_kernel: int = 2,
+        n2v2: bool = False,
+        groups: int = 1,
+    ) -> None:
+        """
+        Constructor.
+
+        Parameters
+        ----------
+        conv_dim : int
+            Number of dimension of the convolution layers, 2 for 2D or 3 for 3D.
+        in_channels : int, optional
+            Number of input channels, by default 1.
+        depth : int, optional
+            Number of encoder blocks, by default 3.
+        num_channels_init : int, optional
+            Number of channels in the first encoder block, by default 64.
+        use_batch_norm : bool, optional
+            Whether to use batch normalization, by default True.
+        dropout : float, optional
+            Dropout probability, by default 0.0.
+        pool_kernel : int, optional
+            Kernel size for the max pooling layers, by default 2.
+        n2v2 : bool, optional
+            Whether to use N2V2 architecture, by default False.
+        groups : int, optional
+            Number of blocked connections from input channels to output
+            channels, by default 1.
+        """
+        super().__init__()
+
+        self.pooling = (
+            getattr(nn, f"MaxPool{conv_dim}d")(kernel_size=pool_kernel)
+            if not n2v2
+            else MaxBlurPool(dim=conv_dim, kernel_size=3, max_pool_size=pool_kernel)
+        )
+
+        encoder_blocks = []
+
+        for n in range(depth):
+            out_channels = num_channels_init * (2**n) * groups
+            in_channels = in_channels if n == 0 else out_channels // 2
+            encoder_blocks.append(
+                Conv_Block(
+                    conv_dim,
+                    in_channels=in_channels,
+                    out_channels=out_channels,
+                    dropout_perc=dropout,
+                    use_batch_norm=use_batch_norm,
+                    groups=groups,
+                )
+            )
+            encoder_blocks.append(self.pooling)
+        self.encoder_blocks = nn.ModuleList(encoder_blocks)
+
+    def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
+        """
+        Forward pass.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input tensor.
+
+        Returns
+        -------
+        List[torch.Tensor]
+            Output of each encoder block (skip connections) and final output of the
+            encoder.
+        """
+        encoder_features = []
+        for module in self.encoder_blocks:
+            x = module(x)
+            if isinstance(module, Conv_Block):
+                encoder_features.append(x)
+        features = [x, *encoder_features]
+        return features
+
+
+class UnetDecoder(nn.Module):
+    """
+    Unet decoder pathway.
+
+    Parameters
+    ----------
+    conv_dim : int
+        Number of dimension of the convolution layers, 2 for 2D or 3 for 3D.
+    depth : int, optional
+        Number of decoder blocks, by default 3.
+    num_channels_init : int, optional
+        Number of channels in the first encoder block, by default 64.
+    use_batch_norm : bool, optional
+        Whether to use batch normalization, by default True.
+    dropout : float, optional
+        Dropout probability, by default 0.0.
+    n2v2 : bool, optional
+        Whether to use N2V2 architecture, by default False.
+    groups : int, optional
+        Number of blocked connections from input channels to output
+        channels, by default 1.
+    """
+
+    def __init__(
+        self,
+        conv_dim: int,
+        depth: int = 3,
+        num_channels_init: int = 64,
+        use_batch_norm: bool = True,
+        dropout: float = 0.0,
+        n2v2: bool = False,
+        groups: int = 1,
+    ) -> None:
+        """
+        Constructor.
+
+        Parameters
+        ----------
+        conv_dim : int
+            Number of dimension of the convolution layers, 2 for 2D or 3 for 3D.
+        depth : int, optional
+            Number of decoder blocks, by default 3.
+        num_channels_init : int, optional
+            Number of channels in the first encoder block, by default 64.
+        use_batch_norm : bool, optional
+            Whether to use batch normalization, by default True.
+        dropout : float, optional
+            Dropout probability, by default 0.0.
+        n2v2 : bool, optional
+            Whether to use N2V2 architecture, by default False.
+        groups : int, optional
+            Number of blocked connections from input channels to output
+            channels, by default 1.
+        """
+        super().__init__()
+
+        upsampling = nn.Upsample(
+            scale_factor=2, mode="bilinear" if conv_dim == 2 else "trilinear"
+        )
+        in_channels = out_channels = num_channels_init * groups * (2 ** (depth - 1))
+
+        self.n2v2 = n2v2
+        self.groups = groups
+
+        self.bottleneck = Conv_Block(
+            conv_dim,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            intermediate_channel_multiplier=2,
+            use_batch_norm=use_batch_norm,
+            dropout_perc=dropout,
+            groups=self.groups,
+        )
+
+        decoder_blocks: List[nn.Module] = []
+        for n in range(depth):
+            decoder_blocks.append(upsampling)
+            in_channels = (num_channels_init * 2 ** (depth - n)) * groups
+            out_channels = in_channels // 2
+            decoder_blocks.append(
+                Conv_Block(
+                    conv_dim,
+                    in_channels=(
+                        in_channels + in_channels // 2 if n > 0 else in_channels
+                    ),
+                    out_channels=out_channels,
+                    intermediate_channel_multiplier=2,
+                    dropout_perc=dropout,
+                    activation="ReLU",
+                    use_batch_norm=use_batch_norm,
+                    groups=groups,
+                )
+            )
+
+        self.decoder_blocks = nn.ModuleList(decoder_blocks)
+
+    def forward(self, *features: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass.
+
+        Parameters
+        ----------
+        *features :  List[torch.Tensor]
+            List containing the output of each encoder block(skip connections) and final
+            output of the encoder.
+
+        Returns
+        -------
+        torch.Tensor
+            Output of the decoder.
+        """
+        x: torch.Tensor = features[0]
+        skip_connections: Tuple[torch.Tensor, ...] = features[-1:0:-1]
+
+        x = self.bottleneck(x)
+
+        for i, module in enumerate(self.decoder_blocks):
+            x = module(x)
+            if isinstance(module, nn.Upsample):
+                # divide index by 2 because of upsampling layers
+                skip_connection: torch.Tensor = skip_connections[i // 2]
+                if self.n2v2:
+                    if x.shape != skip_connections[-1].shape:
+                        x = self._interleave(x, skip_connection, self.groups)
+                else:
+                    x = self._interleave(x, skip_connection, self.groups)
+        return x
+
+    @staticmethod
+    def _interleave(A: torch.Tensor, B: torch.Tensor, groups: int) -> torch.Tensor:
+        """Interleave two tensors.
+
+        Splits the tensors `A` and `B` into equally sized groups along the channel
+        axis (axis=1); then concatenates the groups in alternating order along the
+        channel axis, starting with the first group from tensor A.
+
+        Parameters
+        ----------
+        A : torch.Tensor
+            First tensor.
+        B : torch.Tensor
+            Second tensor.
+        groups : int
+            The number of groups.
+
+        Returns
+        -------
+        torch.Tensor
+            Interleaved tensor.
+
+        Raises
+        ------
+        ValueError:
+            If either of `A` or `B`'s channel axis is not divisible by `groups`.
+        """
+        if (A.shape[1] % groups != 0) or (B.shape[1] % groups != 0):
+            raise ValueError(f"Number of channels not divisible by {groups} groups.")
+
+        m = A.shape[1] // groups
+        n = B.shape[1] // groups
+
+        A_groups: List[torch.Tensor] = [
+            A[:, i * m : (i + 1) * m] for i in range(groups)
+        ]
+        B_groups: List[torch.Tensor] = [
+            B[:, i * n : (i + 1) * n] for i in range(groups)
+        ]
+
+        interleaved = torch.cat(
+            [
+                tensor_list[i]
+                for i in range(groups)
+                for tensor_list in [A_groups, B_groups]
+            ],
+            dim=1,
+        )
+
+        return interleaved
+
+
+class UNet(nn.Module):
+    """
+    UNet model.
+
+    Adapted for PyTorch from:
+    https://github.com/juglab/n2v/blob/main/n2v/nets/unet_blocks.py.
+
+    Parameters
+    ----------
+    conv_dims : int
+        Number of dimensions of the convolution layers (2 or 3).
+    num_classes : int, optional
+        Number of classes to predict, by default 1.
+    in_channels : int, optional
+        Number of input channels, by default 1.
+    depth : int, optional
+        Number of downsamplings, by default 3.
+    num_channels_init : int, optional
+        Number of filters in the first convolution layer, by default 64.
+    use_batch_norm : bool, optional
+        Whether to use batch normalization, by default True.
+    dropout : float, optional
+        Dropout probability, by default 0.0.
+    pool_kernel : int, optional
+        Kernel size of the pooling layers, by default 2.
+    final_activation : Optional[Callable], optional
+        Activation function to use for the last layer, by default None.
+    n2v2 : bool, optional
+        Whether to use N2V2 architecture, by default False.
+    independent_channels : bool
+        Whether to train the channels independently, by default True.
+    **kwargs : Any
+        Additional keyword arguments, unused.
+    """
+
+    def __init__(
+        self,
+        conv_dims: int,
+        num_classes: int = 1,
+        in_channels: int = 1,
+        depth: int = 3,
+        num_channels_init: int = 64,
+        use_batch_norm: bool = True,
+        dropout: float = 0.0,
+        pool_kernel: int = 2,
+        final_activation: Union[SupportedActivation, str] = SupportedActivation.NONE,
+        n2v2: bool = False,
+        independent_channels: bool = True,
+        **kwargs: Any,
+    ) -> None:
+        """
+        Constructor.
+
+        Parameters
+        ----------
+        conv_dims : int
+            Number of dimensions of the convolution layers (2 or 3).
+        num_classes : int, optional
+            Number of classes to predict, by default 1.
+        in_channels : int, optional
+            Number of input channels, by default 1.
+        depth : int, optional
+            Number of downsamplings, by default 3.
+        num_channels_init : int, optional
+            Number of filters in the first convolution layer, by default 64.
+        use_batch_norm : bool, optional
+            Whether to use batch normalization, by default True.
+        dropout : float, optional
+            Dropout probability, by default 0.0.
+        pool_kernel : int, optional
+            Kernel size of the pooling layers, by default 2.
+        final_activation : Optional[Callable], optional
+            Activation function to use for the last layer, by default None.
+        n2v2 : bool, optional
+            Whether to use N2V2 architecture, by default False.
+        independent_channels : bool
+            Whether to train parallel independent networks for each channel, by
+            default True.
+        **kwargs : Any
+            Additional keyword arguments, unused.
+        """
+        super().__init__()
+
+        groups = in_channels if independent_channels else 1
+
+        self.encoder = UnetEncoder(
+            conv_dims,
+            in_channels=in_channels,
+            depth=depth,
+            num_channels_init=num_channels_init,
+            use_batch_norm=use_batch_norm,
+            dropout=dropout,
+            pool_kernel=pool_kernel,
+            n2v2=n2v2,
+            groups=groups,
+        )
+
+        self.decoder = UnetDecoder(
+            conv_dims,
+            depth=depth,
+            num_channels_init=num_channels_init,
+            use_batch_norm=use_batch_norm,
+            dropout=dropout,
+            n2v2=n2v2,
+            groups=groups,
+        )
+        self.final_conv = getattr(nn, f"Conv{conv_dims}d")(
+            in_channels=num_channels_init * groups,
+            out_channels=num_classes,
+            kernel_size=1,
+            groups=groups,
+        )
+        self.final_activation = get_activation(final_activation)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass.
+
+        Parameters
+        ----------
+        x :  torch.Tensor
+            Input tensor.
+
+        Returns
+        -------
+        torch.Tensor
+            Output of the model.
+        """
+        encoder_features = self.encoder(x)
+        x = self.decoder(*encoder_features)
+        x = self.final_conv(x)
+        x = self.final_activation(x)
+        return x

careamics/prediction_utils/__init__.py

@@ -0,0 +1,10 @@
+"""Package to house various prediction utilies."""
+
+__all__ = [
+    "stitch_prediction",
+    "stitch_prediction_single",
+    "convert_outputs",
+]
+
+from .prediction_outputs import convert_outputs
+from .stitch_prediction import stitch_prediction, stitch_prediction_single

careamics/prediction_utils/lvae_prediction.py

@@ -0,0 +1,158 @@
+"""Module containing pytorch implementations for obtaining predictions from an LVAE."""
+
+from typing import Any, Optional
+
+import torch
+
+from careamics.models.lvae import LadderVAE as LVAE
+from careamics.models.lvae.likelihoods import LikelihoodModule
+
+# TODO: convert these functions to lightning module `predict_step`
+#   -> mmse_count will have to be an instance attribute?
+
+
+# This function is needed because the output of the datasets (input here) can include
+#   auxillary items, such as the TileInformation. This function allows for easier reuse
+#   between lvae_predict_single_sample and lvae_predict_mmse.
+def lvae_predict_single_sample(
+    model: LVAE,
+    likelihood_obj: LikelihoodModule,
+    input: torch.Tensor,
+) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+    """
+    Generate a single sample prediction from an LVAE model, for a given input.
+
+    Parameters
+    ----------
+    model : LVAE
+        Trained LVAE model.
+    likelihood_obj : LikelihoodModule
+        Instance of a likelihood class.
+    input : torch.tensor
+        Input to generate prediction for. Expected shape is (S, C, Y, X).
+
+    Returns
+    -------
+    tuple of (torch.tensor, optional torch.tensor)
+        The first element is the sample prediction, and the second element is the
+        log-variance. The log-variance will be None if `model.predict_logvar is None`.
+    """
+    model.eval()  # Not in original predict code: effects batch_norm and dropout layers
+    with torch.no_grad():
+        output: torch.Tensor
+        output, _ = model(input)  # 2nd item is top-down data dict
+
+    # presently, get_mean_lv just splits the output in 2 if predict_logvar=True,
+    #   optionally clips the logvavr if logvar_lowerbound is not None
+    # TODO: consider refactoring to remove use of the likelihood object
+    sample_prediction, log_var = likelihood_obj.get_mean_lv(output)
+
+    # TODO: output denormalization using target stats that will be saved in data config
+    # -> Don't think we need this, saw it in a random bit of code somewhere.
+
+    return sample_prediction, log_var
+
+
+def lvae_predict_tiled_batch(
+    model: LVAE,
+    likelihood_obj: LikelihoodModule,
+    input: tuple[Any],
+) -> tuple[tuple[Any], Optional[tuple[Any]]]:
+    # TODO: fix docstring return types, ... too many output options
+    """
+    Generate a single sample prediction from an LVAE model, for a given input.
+
+    Parameters
+    ----------
+    model : LVAE
+        Trained LVAE model.
+    likelihood_obj : LikelihoodModule
+        Instance of a likelihood class.
+    input : torch.tensor | tuple of (torch.tensor, Any, ...)
+        Input to generate prediction for. This can include auxilary inputs such as
+        `TileInformation`, but the model input is always the first item of the tuple.
+        Expected shape of the model input is (S, C, Y, X).
+
+    Returns
+    -------
+    tuple of ((torch.tensor, Any, ...), optional tuple of (torch.tensor, Any, ...))
+        The first element is the sample prediction, and the second element is the
+        log-variance. The log-variance will be None if `model.predict_logvar is None`.
+        Any auxillary data included in the input will also be include with both the
+        sample prediction and the log-variance.
+    """
+    x: torch.Tensor
+    aux: list[Any]
+    x, *aux = input
+
+    sample_prediction, log_var = lvae_predict_single_sample(
+        model=model, likelihood_obj=likelihood_obj, input=x
+    )
+
+    log_var_output = (log_var, *aux) if log_var is not None else None
+    return (sample_prediction, *aux), log_var_output
+
+
+def lvae_predict_mmse_tiled_batch(
+    model: LVAE,
+    likelihood_obj: LikelihoodModule,
+    input: tuple[Any],
+    mmse_count: int,
+) -> tuple[tuple[Any], tuple[Any], Optional[tuple[Any]]]:
+    # TODO: fix docstring return types, ... hard to make readable
+    """
+    Generate the MMSE (minimum mean squared error) prediction, for a given input.
+
+    This is calculated from the mean of multiple single sample predictions.
+
+    Parameters
+    ----------
+    model : LVAE
+        Trained LVAE model.
+    likelihood_obj : LikelihoodModule
+        Instance of a likelihood class.
+    input : torch.tensor | tuple of (torch.tensor, Any, ...)
+        Input to generate prediction for. This can include auxilary inputs such as
+        `TileInformation`, but the model input is always the first item of the tuple.
+        Expected shape of the model input is (S, C, Y, X).
+    mmse_count : int
+        Number of samples to generate to calculate MMSE (minimum mean squared error).
+
+    Returns
+    -------
+    tuple of (tuple of (torch.Tensor[Any], Any, ...))
+        A tuple of 3 elements. The first element contains the MMSE prediction, the
+        second contains the standard deviation of the samples used to create the MMSE
+        prediction. Finally the last element contains the log-variance of the
+        likelihood, this will be `None` if `likelihood.predict_logvar` is `None`.
+        Any auxillary data included in the input will also be include with all of the
+        MMSE prediction, the standard deviation, and the log-variance.
+    """
+    if mmse_count <= 0:
+        raise ValueError("MMSE count must be greater than zero.")
+
+    x: torch.Tensor
+    aux: list[Any]
+    x, *aux = input
+
+    input_shape = x.shape
+    output_shape = (input_shape[0], model.target_ch, *input_shape[2:])
+    log_var: Optional[torch.Tensor] = None
+    # pre-declare empty array to fill with individual sample predictions
+    sample_predictions = torch.zeros(size=(mmse_count, *output_shape))
+    for mmse_idx in range(mmse_count):
+        sample_prediction, lv = lvae_predict_single_sample(
+            model=model, likelihood_obj=likelihood_obj, input=x
+        )
+        # only keep the log variance of the first sample prediction
+        if mmse_idx == 0:
+            log_var = lv
+
+        # store sample predictions
+        sample_predictions[mmse_idx, ...] = sample_prediction
+
+    mmse_prediction = torch.mean(sample_predictions, dim=0)
+    mmse_prediction_std = torch.std(sample_predictions, dim=0)
+
+    log_var_output = (log_var, *aux) if log_var is not None else None
+    return (mmse_prediction, *aux), (mmse_prediction_std, *aux), log_var_output