careamics 0.0.1__py3-none-any.whl → 0.0.2__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/prediction_outputs.py

@@ -0,0 +1,135 @@
+"""Module containing functions to convert prediction outputs to desired form."""
+
+from typing import Any, List, Literal, Tuple, Union, overload
+
+import numpy as np
+from numpy.typing import NDArray
+
+from ..config.tile_information import TileInformation
+from .stitch_prediction import stitch_prediction
+
+
+def convert_outputs(predictions: List[Any], tiled: bool) -> list[NDArray]:
+    """
+    Convert the Lightning trainer outputs to the desired form.
+
+    This method allows stitching back together tiled predictions.
+
+    Parameters
+    ----------
+    predictions : list
+        Predictions that are output from `Trainer.predict`.
+    tiled : bool
+        Whether the predictions are tiled.
+
+    Returns
+    -------
+    list of numpy.ndarray or numpy.ndarray
+        List of arrays with the axes SC(Z)YX. If there is only 1 output it will not
+        be in a list.
+    """
+    if len(predictions) == 0:
+        return predictions
+
+    # this layout is to stop mypy complaining
+    if tiled:
+        predictions_comb = combine_batches(predictions, tiled)
+        predictions_output = stitch_prediction(*predictions_comb)
+    else:
+        predictions_output = combine_batches(predictions, tiled)
+
+    return predictions_output
+
+
+# for mypy
+@overload
+def combine_batches(  # numpydoc ignore=GL08
+    predictions: List[Any], tiled: Literal[True]
+) -> Tuple[List[NDArray], List[TileInformation]]: ...
+
+
+# for mypy
+@overload
+def combine_batches(  # numpydoc ignore=GL08
+    predictions: List[Any], tiled: Literal[False]
+) -> List[NDArray]: ...
+
+
+# for mypy
+@overload
+def combine_batches(  # numpydoc ignore=GL08
+    predictions: List[Any], tiled: Union[bool, Literal[True], Literal[False]]
+) -> Union[List[NDArray], Tuple[List[NDArray], List[TileInformation]]]: ...
+
+
+def combine_batches(
+    predictions: List[Any], tiled: bool
+) -> Union[List[NDArray], Tuple[List[NDArray], List[TileInformation]]]:
+    """
+    If predictions are in batches, they will be combined.
+
+    Parameters
+    ----------
+    predictions : list
+        Predictions that are output from `Trainer.predict`.
+    tiled : bool
+        Whether the predictions are tiled.
+
+    Returns
+    -------
+    (list of numpy.ndarray) or tuple of (list of numpy.ndarray, list of TileInformation)
+        Combined batches.
+    """
+    if tiled:
+        return _combine_tiled_batches(predictions)
+    else:
+        return _combine_array_batches(predictions)
+
+
+def _combine_tiled_batches(
+    predictions: List[Tuple[NDArray, List[TileInformation]]]
+) -> Tuple[List[NDArray], List[TileInformation]]:
+    """
+    Combine batches from tiled output.
+
+    Parameters
+    ----------
+    predictions : list of (numpy.ndarray, list of TileInformation)
+        Predictions that are output from `Trainer.predict`. For tiled batches, this is
+        a list of tuples. The first element of the tuples is the prediction output of
+        tiles with dimension (B, C, (Z), Y, X), where B is batch size. The second
+        element of the tuples is a list of TileInformation objects of length B.
+
+    Returns
+    -------
+    tuple of (list of numpy.ndarray, list of TileInformation)
+        Combined batches.
+    """
+    # turn list of lists into single list
+    tile_infos = [
+        tile_info for _, tile_info_list in predictions for tile_info in tile_info_list
+    ]
+    prediction_tiles: List[NDArray] = _combine_array_batches(
+        [preds for preds, _ in predictions]
+    )
+    return prediction_tiles, tile_infos
+
+
+def _combine_array_batches(predictions: List[NDArray]) -> List[NDArray]:
+    """
+    Combine batches of arrays.
+
+    Parameters
+    ----------
+    predictions : list
+        Prediction arrays that are output from `Trainer.predict`. A list of arrays that
+        have dimensions (B, C, (Z), Y, X), where B is batch size.
+
+    Returns
+    -------
+    list of numpy.ndarray
+        A list of arrays with dimensions (1, C, (Z), Y, X).
+    """
+    prediction_concat: NDArray = np.concatenate(predictions, axis=0)
+    prediction_split = np.split(prediction_concat, prediction_concat.shape[0], axis=0)
+    return prediction_split

careamics/prediction_utils/stitch_prediction.py

@@ -0,0 +1,98 @@
+"""Prediction utility functions."""
+
+import builtins
+from typing import List, Union
+
+import numpy as np
+from numpy.typing import NDArray
+
+from careamics.config.tile_information import TileInformation
+
+
+# TODO: why not allow input and output of torch.tensor ?
+def stitch_prediction(
+    tiles: List[np.ndarray],
+    tile_infos: List[TileInformation],
+) -> List[np.ndarray]:
+    """
+    Stitch tiles back together to form a full image(s).
+
+    Tiles are of dimensions SC(Z)YX, where C is the number of channels and can be a
+    singleton dimension.
+
+    Parameters
+    ----------
+    tiles : list of numpy.ndarray
+        Cropped tiles and their respective stitching coordinates. Can contain tiles
+        from multiple images.
+    tile_infos : list of TileInformation
+        List of information and coordinates obtained from
+        `dataset.tiled_patching.extract_tiles`.
+
+    Returns
+    -------
+    list of numpy.ndarray
+        Full image(s).
+    """
+    # Find where to split the lists so that only info from one image is contained.
+    # Do this by locating the last tiles of each image.
+    last_tiles = [tile_info.last_tile for tile_info in tile_infos]
+    last_tile_position = np.where(last_tiles)[0]
+    image_slices = [
+        slice(
+            None if i == 0 else last_tile_position[i - 1] + 1, last_tile_position[i] + 1
+        )
+        for i in range(len(last_tile_position))
+    ]
+    image_predictions = []
+    # slice the lists and apply stitch_prediction_single to each in turn.
+    for image_slice in image_slices:
+        image_predictions.append(
+            stitch_prediction_single(tiles[image_slice], tile_infos[image_slice])
+        )
+    return image_predictions
+
+
+def stitch_prediction_single(
+    tiles: List[NDArray],
+    tile_infos: List[TileInformation],
+) -> NDArray:
+    """
+    Stitch tiles back together to form a full image.
+
+    Tiles are of dimensions SC(Z)YX, where C is the number of channels and can be a
+    singleton dimension.
+
+    Parameters
+    ----------
+    tiles : list of numpy.ndarray
+        Cropped tiles and their respective stitching coordinates.
+    tile_infos : list of TileInformation
+        List of information and coordinates obtained from
+        `dataset.tiled_patching.extract_tiles`.
+
+    Returns
+    -------
+    numpy.ndarray
+        Full image, with dimensions SC(Z)YX.
+    """
+    # retrieve whole array size
+    input_shape = (1, *tile_infos[0].array_shape)  # add S dim
+    predicted_image = np.zeros(input_shape, dtype=np.float32)
+
+    for tile, tile_info in zip(tiles, tile_infos):
+
+        # Compute coordinates for cropping predicted tile
+        crop_slices: tuple[Union[builtins.ellipsis, slice], ...] = (
+            ...,
+            *[slice(c[0], c[1]) for c in tile_info.overlap_crop_coords],
+        )
+
+        # Crop predited tile according to overlap coordinates
+        cropped_tile = tile[crop_slices]
+
+        # Insert cropped tile into predicted image using stitch coordinates
+        image_slices = (..., *[slice(c[0], c[1]) for c in tile_info.stitch_coords])
+        predicted_image[image_slices] = cropped_tile.astype(np.float32)
+
+    return predicted_image

careamics/transforms/__init__.py

@@ -0,0 +1,20 @@
+"""Transforms that are used to augment the data."""
+
+__all__ = [
+    "get_all_transforms",
+    "N2VManipulate",
+    "XYFlip",
+    "XYRandomRotate90",
+    "ImageRestorationTTA",
+    "Denormalize",
+    "Normalize",
+    "Compose",
+]
+
+
+from .compose import Compose, get_all_transforms
+from .n2v_manipulate import N2VManipulate
+from .normalize import Denormalize, Normalize
+from .tta import ImageRestorationTTA
+from .xy_flip import XYFlip
+from .xy_random_rotate90 import XYRandomRotate90