careamics 0.1.0rc4__py3-none-any.whl → 0.1.0rc6__py3-none-any.whl

careamics/models/layers.py

@@ -3,6 +3,7 @@ Layer module.
 
 This submodule contains layers used in the CAREamics models.
 """
+
 from typing import List, Optional, Tuple, Union
 
 import torch
@@ -161,6 +162,18 @@ def _unpack_kernel_size(
     """Unpack kernel_size to a tuple of ints.
 
     Inspired by Kornia implementation. TODO: link
+
+    Parameters
+    ----------
+    kernel_size : Union[Tuple[int, ...], int]
+        Kernel size.
+    dim : int
+        Number of dimensions.
+
+    Returns
+    -------
+    Tuple[int, ...]
+        Kernel size tuple.
     """
     if isinstance(kernel_size, int):
         kernel_dims = tuple([kernel_size for _ in range(dim)])
@@ -172,7 +185,20 @@ def _unpack_kernel_size(
 def _compute_zero_padding(
     kernel_size: Union[Tuple[int, ...], int], dim: int
 ) -> Tuple[int, ...]:
-    """Utility function that computes zero padding tuple."""
+    """Utility function that computes zero padding tuple.
+
+    Parameters
+    ----------
+    kernel_size : Union[Tuple[int, ...], int]
+        Kernel size.
+    dim : int
+        Number of dimensions.
+
+    Returns
+    -------
+    Tuple[int, ...]
+        Zero padding tuple.
+    """
     kernel_dims = _unpack_kernel_size(kernel_size, dim)
     return tuple([(kd - 1) // 2 for kd in kernel_dims])
 
@@ -190,14 +216,19 @@ def get_pascal_kernel_1d(
 
     Parameters
     ----------
-    kernel_size: height and width of the kernel.
-    norm: if to normalize the kernel or not. Default: False.
-    device: tensor device
-    dtype: tensor dtype
+    kernel_size : int
+        Kernel size.
+    norm : bool
+        Normalize the kernel, by default False.
+    device : Optional[torch.device]
+        Device of the tensor, by default None.
+    dtype : Optional[torch.dtype]
+        Data type of the tensor, by default None.
 
     Returns
     -------
-    kernel shaped as :math:`(kernel_size,)`
+    torch.Tensor
+        Pascal kernel.
 
     Examples
     --------
@@ -244,19 +275,28 @@ def _get_pascal_kernel_nd(
 ) -> torch.Tensor:
     """Generate pascal filter kernel by kernel size.
 
+    If kernel_size is an integer the kernel will be shaped as (kernel_size, kernel_size)
+    otherwise the kernel will be shaped as kernel_size
+
     Inspired by Kornia implementation.
 
     Parameters
     ----------
-    kernel_size: height and width of the kernel.
-    norm: if to normalize the kernel or not. Default: True.
-    device: tensor device
-    dtype: tensor dtype
+    kernel_size : Union[Tuple[int, int], int]
+        Kernel size for the pascal kernel.
+    norm : bool
+        Normalize the kernel, by default True.
+    dim : int
+        Number of dimensions, by default 2.
+    device : Optional[torch.device]
+        Device of the tensor, by default None.
+    dtype : Optional[torch.dtype]
+        Data type of the tensor, by default None.
 
     Returns
     -------
-    if kernel_size is an integer the kernel will be shaped as (kernel_size, kernel_size)
-    otherwise the kernel will be shaped as kernel_size
+    torch.Tensor
+        Pascal kernel.
 
     Examples
     --------
@@ -302,6 +342,24 @@ def _max_blur_pool_by_kernel2d(
     """Compute max_blur_pool by a given :math:`CxC_(out, None)xNxN` kernel.
 
     Inspired by Kornia implementation.
+
+    Parameters
+    ----------
+    x : torch.Tensor
+        Input tensor.
+    kernel : torch.Tensor
+        Kernel tensor.
+    stride : int
+        Stride.
+    max_pool_size : int
+        Maximum pool size.
+    ceil_mode : bool
+        Ceil mode, by default False. Set to True to match output size of conv2d.
+
+    Returns
+    -------
+    torch.Tensor
+        Output tensor.
     """
     # compute local maxima
     x = F.max_pool2d(
@@ -322,6 +380,24 @@ def _max_blur_pool_by_kernel3d(
     """Compute max_blur_pool by a given :math:`CxC_(out, None)xNxNxN` kernel.
 
     Inspired by Kornia implementation.
+
+    Parameters
+    ----------
+    x : torch.Tensor
+        Input tensor.
+    kernel : torch.Tensor
+        Kernel tensor.
+    stride : int
+        Stride.
+    max_pool_size : int
+        Maximum pool size.
+    ceil_mode : bool
+        Ceil mode, by default False. Set to True to match output size of conv2d.
+
+    Returns
+    -------
+    torch.Tensor
+        Output tensor.
     """
     # compute local maxima
     x = F.max_pool3d(
@@ -342,21 +418,16 @@ class MaxBlurPool(nn.Module):
 
     Parameters
     ----------
-    dim: int
-        Toggles between 2D and 3D
-    kernel_size: Union[Tuple[int, int], int]
+    dim : int
+        Toggles between 2D and 3D.
+    kernel_size : Union[Tuple[int, int], int]
         Kernel size for max pooling.
-    stride: int
+    stride : int
         Stride for pooling.
-    max_pool_size: int
+    max_pool_size : int
         Max kernel size for max pooling.
-    ceil_mode: bool
-        Should be true to match output size of conv2d with same kernel size.
-
-    Returns
-    -------
-    torch.Tensor
-        The pooled and blurred tensor.
+    ceil_mode : bool
+        Ceil mode, by default False. Set to True to match output size of conv2d.
     """
 
     def __init__(
@@ -367,6 +438,21 @@ class MaxBlurPool(nn.Module):
         max_pool_size: int = 2,
         ceil_mode: bool = False,
     ) -> None:
+        """Constructor.
+
+        Parameters
+        ----------
+        dim : int
+            Dimension of the convolution.
+        kernel_size : Union[Tuple[int, int], int]
+            Kernel size for max pooling.
+        stride : int, optional
+            Stride, by default 2.
+        max_pool_size : int, optional
+            Maximum pool size, by default 2.
+        ceil_mode : bool, optional
+            Ceil mode, by default False. Set to True to match output size of conv2d.
+        """
         super().__init__()
         self.dim = dim
         self.kernel_size = kernel_size
@@ -376,7 +462,18 @@ class MaxBlurPool(nn.Module):
         self.kernel = _get_pascal_kernel_nd(kernel_size, norm=True, dim=self.dim)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """Forward pass of the function."""
+        """Forward pass of the function.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input tensor.
+
+        Returns
+        -------
+        torch.Tensor
+            Output tensor.
+        """
         self.kernel = torch.as_tensor(self.kernel, device=x.device, dtype=x.dtype)
         if self.dim == 2:
             return _max_blur_pool_by_kernel2d(

careamics/models/model_factory.py

@@ -3,6 +3,7 @@ Model factory.
 
 Model creation factory functions.
 """
+
 from typing import Union
 
 import torch
@@ -26,7 +27,7 @@ def model_factory(
     Parameters
     ----------
     model_configuration : Union[UNetModel, VAEModel]
-        Model configuration
+        Model configuration.
 
     Returns
     -------

careamics/models/unet.py

@@ -3,7 +3,8 @@ UNet model.
 
 A UNet encoder, decoder and complete model.
 """
-from typing import Any, List, Union
+
+from typing import Any, List, Tuple, Union
 
 import torch
 import torch.nn as nn
@@ -33,6 +34,11 @@ class UnetEncoder(nn.Module):
         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__(
@@ -45,6 +51,7 @@ class UnetEncoder(nn.Module):
         dropout: float = 0.0,
         pool_kernel: int = 2,
         n2v2: bool = False,
+        groups: int = 1,
     ) -> None:
         """
         Constructor.
@@ -65,6 +72,11 @@ class UnetEncoder(nn.Module):
             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__()
 
@@ -77,7 +89,7 @@ class UnetEncoder(nn.Module):
         encoder_blocks = []
 
         for n in range(depth):
-            out_channels = num_channels_init * (2**n)
+            out_channels = num_channels_init * (2**n) * groups
             in_channels = in_channels if n == 0 else out_channels // 2
             encoder_blocks.append(
                 Conv_Block(
@@ -86,6 +98,7 @@ class UnetEncoder(nn.Module):
                     out_channels=out_channels,
                     dropout_perc=dropout,
                     use_batch_norm=use_batch_norm,
+                    groups=groups,
                 )
             )
             encoder_blocks.append(self.pooling)
@@ -131,6 +144,11 @@ class UnetDecoder(nn.Module):
         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__(
@@ -141,6 +159,7 @@ class UnetDecoder(nn.Module):
         use_batch_norm: bool = True,
         dropout: float = 0.0,
         n2v2: bool = False,
+        groups: int = 1,
     ) -> None:
         """
         Constructor.
@@ -157,15 +176,21 @@ class UnetDecoder(nn.Module):
             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 * 2 ** (depth - 1)
+        in_channels = out_channels = num_channels_init * groups * (2 ** (depth - 1))
 
         self.n2v2 = n2v2
+        self.groups = groups
 
         self.bottleneck = Conv_Block(
             conv_dim,
@@ -174,34 +199,32 @@ class UnetDecoder(nn.Module):
             intermediate_channel_multiplier=2,
             use_batch_norm=use_batch_norm,
             dropout_perc=dropout,
+            groups=self.groups,
         )
 
-        decoder_blocks = []
+        decoder_blocks: List[nn.Module] = []
         for n in range(depth):
             decoder_blocks.append(upsampling)
-            in_channels = (
-                num_channels_init ** (depth - n)
-                if (self.n2v2 and n == depth - 1)
-                else num_channels_init * 2 ** (depth - n)
-            )
+            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,
+                    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: List[torch.Tensor]) -> torch.Tensor:
+    def forward(self, *features: torch.Tensor) -> torch.Tensor:
         """
         Forward pass.
 
@@ -217,20 +240,73 @@ class UnetDecoder(nn.Module):
             Output of the decoder.
         """
         x: torch.Tensor = features[0]
-        skip_connections: torch.Tensor = features[1:][::-1]
+        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 = torch.cat([x, skip_connections[i // 2]], axis=1)
+                        x = self._interleave(x, skip_connection, self.groups)
                 else:
-                    x = torch.cat([x, skip_connections[i // 2]], axis=1)
+                    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):
     """
@@ -257,8 +333,14 @@ class UNet(nn.Module):
         Dropout probability, by default 0.0.
     pool_kernel : int, optional
         Kernel size of the pooling layers, by default 2.
-    last_activation : Optional[Callable], optional
+    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__(
@@ -273,6 +355,7 @@ class UNet(nn.Module):
         pool_kernel: int = 2,
         final_activation: Union[SupportedActivation, str] = SupportedActivation.NONE,
         n2v2: bool = False,
+        independent_channels: bool = True,
         **kwargs: Any,
     ) -> None:
         """
@@ -296,11 +379,20 @@ class UNet(nn.Module):
             Dropout probability, by default 0.0.
         pool_kernel : int, optional
             Kernel size of the pooling layers, by default 2.
-        last_activation : Optional[Callable], optional
+        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,
@@ -310,6 +402,7 @@ class UNet(nn.Module):
             dropout=dropout,
             pool_kernel=pool_kernel,
             n2v2=n2v2,
+            groups=groups,
         )
 
         self.decoder = UnetDecoder(
@@ -319,11 +412,13 @@ class UNet(nn.Module):
             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,
+            in_channels=num_channels_init * groups,
             out_channels=num_classes,
             kernel_size=1,
+            groups=groups,
         )
         self.final_activation = get_activation(final_activation)
 

careamics/prediction/stitch_prediction.py

@@ -1,8 +1,5 @@
-"""
-Prediction convenience functions.
+"""Prediction utility functions."""
 
-These functions are used during prediction.
-"""
 from typing import List
 
 import numpy as np
@@ -20,7 +17,7 @@ def stitch_prediction(
     ----------
     tiles : List[torch.Tensor]
         Cropped tiles and their respective stitching coordinates.
-    stitching_coords : List
+    stitching_data : List
         List of information and coordinates obtained from
         `dataset.tiled_patching.extract_tiles`.
 

careamics/transforms/__init__.py

@@ -3,39 +3,18 @@
 __all__ = [
     "get_all_transforms",
     "N2VManipulate",
-    "NDFlip",
+    "XYFlip",
     "XYRandomRotate90",
     "ImageRestorationTTA",
     "Denormalize",
     "Normalize",
+    "Compose",
 ]
 
 
+from .compose import Compose, get_all_transforms
 from .n2v_manipulate import N2VManipulate
-from .nd_flip import NDFlip
 from .normalize import Denormalize, Normalize
 from .tta import ImageRestorationTTA
+from .xy_flip import XYFlip
 from .xy_random_rotate90 import XYRandomRotate90
-
-ALL_TRANSFORMS = {
-    "Normalize": Normalize,
-    "N2VManipulate": N2VManipulate,
-    "NDFlip": NDFlip,
-    "XYRandomRotate90": XYRandomRotate90,
-}
-
-
-def get_all_transforms() -> dict:
-    """Return all the transforms accepted by CAREamics.
-
-    Note that while CAREamics accepts any `Compose` transforms from Albumentations (see
-    https://albumentations.ai/), only a few transformations are explicitely supported
-    (see `SupportedTransform`).
-
-    Returns
-    -------
-    dict
-        A dictionary with all the transforms accepted by CAREamics, where the keys are
-        the transform names and the values are the transform classes.
-    """
-    return ALL_TRANSFORMS