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

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,9 @@ class UnetEncoder(nn.Module):
         Dropout probability, by default 0.0.
     pool_kernel : int, optional
         Kernel size for the max pooling layers, by default 2.
+    groups: int, optional
+        Number of blocked connections from input channels to output
+        channels, by default 1.
     """
 
     def __init__(
@@ -45,6 +49,7 @@ class UnetEncoder(nn.Module):
         dropout: float = 0.0,
         pool_kernel: int = 2,
         n2v2: bool = False,
+        groups: int = 1,
     ) -> None:
         """
         Constructor.
@@ -65,6 +70,9 @@ class UnetEncoder(nn.Module):
             Dropout probability, by default 0.0.
         pool_kernel : int, optional
             Kernel size for the max pooling layers, by default 2.
+        groups: int, optional
+            Number of blocked connections from input channels to output
+            channels, by default 1.
         """
         super().__init__()
 
@@ -77,7 +85,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 +94,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 +140,9 @@ class UnetDecoder(nn.Module):
         Whether to use batch normalization, by default True.
     dropout : float, optional
         Dropout probability, by default 0.0.
+    groups: int, optional
+        Number of blocked connections from input channels to output
+        channels, by default 1.
     """
 
     def __init__(
@@ -141,6 +153,7 @@ class UnetDecoder(nn.Module):
         use_batch_norm: bool = True,
         dropout: float = 0.0,
         n2v2: bool = False,
+        groups: int = 1,
     ) -> None:
         """
         Constructor.
@@ -157,15 +170,19 @@ class UnetDecoder(nn.Module):
             Whether to use batch normalization, by default True.
         dropout : float, optional
             Dropout probability, by default 0.0.
+        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 +191,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 +232,70 @@ 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:
+        """
+        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
+        B: torch.Tensor
+        groups: int
+            The number of groups.
+
+        Returns
+        -------
+        torch.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):
     """
@@ -273,6 +338,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:
         """
@@ -298,9 +364,14 @@ class UNet(nn.Module):
             Kernel size of the pooling layers, by default 2.
         last_activation : Optional[Callable], optional
             Activation function to use for the last layer, by default None.
+        independent_channels : bool
+            Whether to train parallel independent networks for each channel, by
+            default True.
         """
         super().__init__()
 
+        groups = in_channels if independent_channels else 1
+
         self.encoder = UnetEncoder(
             conv_dims,
             in_channels=in_channels,
@@ -310,6 +381,7 @@ class UNet(nn.Module):
             dropout=dropout,
             pool_kernel=pool_kernel,
             n2v2=n2v2,
+            groups=groups,
         )
 
         self.decoder = UnetDecoder(
@@ -319,11 +391,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

@@ -3,6 +3,7 @@ Prediction convenience functions.
 
 These functions are used during prediction.
 """
+
 from typing import List
 
 import numpy as np

careamics/transforms/__init__.py

@@ -8,34 +8,13 @@ __all__ = [
     "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_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

careamics/transforms/compose.py

@@ -0,0 +1,98 @@
+"""A class chaining transforms together."""
+
+from typing import Callable, List, Optional, Tuple
+
+import numpy as np
+
+from careamics.config.data_model import TRANSFORMS_UNION
+
+from .n2v_manipulate import N2VManipulate
+from .nd_flip import NDFlip
+from .normalize import Normalize
+from .transform import Transform
+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.
+
+    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
+
+
+class Compose:
+    """A class chaining transforms together."""
+
+    def __init__(self, transform_list: List[TRANSFORMS_UNION]) -> None:
+        """Instantiate a Compose object.
+
+        Parameters
+        ----------
+        transform_list : List[TRANSFORMS_UNION]
+            A list of dictionaries where each dictionary contains the name of a
+            transform and its parameters.
+        """
+        # retrieve all available transforms
+        all_transforms = get_all_transforms()
+
+        # instantiate all transforms
+        transforms = [all_transforms[t.name](**t.model_dump()) for t in transform_list]
+
+        self._callable_transforms = self._chain_transforms(transforms)
+
+    def _chain_transforms(self, transforms: List[Transform]) -> Callable:
+        """Chain the transforms together.
+
+        Parameters
+        ----------
+        transforms : List[Transform]
+            A list of transforms to chain together.
+
+        Returns
+        -------
+        Callable
+            A callable that applies the transforms in order to the input data.
+        """
+
+        def _chain(
+            patch: np.ndarray, target: Optional[np.ndarray]
+        ) -> Tuple[np.ndarray, ...]:
+            params = (patch, target)
+
+            for t in transforms:
+                params = t(*params)
+
+            return params
+
+        return _chain
+
+    def __call__(
+        self, patch: np.ndarray, target: Optional[np.ndarray] = None
+    ) -> Tuple[np.ndarray, ...]:
+        """Apply the transforms to the input data.
+
+        Parameters
+        ----------
+        patch : np.ndarray
+            The input data.
+        target : Optional[np.ndarray], optional
+            Target data, by default None
+
+        Returns
+        -------
+        Tuple[np.ndarray, ...]
+            The output of the transformations.
+        """
+        return self._callable_transforms(patch, target)

careamics/transforms/n2v_manipulate.py

@@ -1,19 +1,19 @@
 from typing import Any, Literal, Optional, Tuple
 
 import numpy as np
-from albumentations import ImageOnlyTransform
 
 from careamics.config.support import SupportedPixelManipulation, SupportedStructAxis
+from careamics.transforms.transform import Transform
 
 from .pixel_manipulation import median_manipulate, uniform_manipulate
 from .struct_mask_parameters import StructMaskParameters
 
 
-class N2VManipulate(ImageOnlyTransform):
+class N2VManipulate(Transform):
     """
     Default augmentation for the N2V model.
 
-    This transform expects (Z)YXC dimensions.
+    This transform expects C(Z)YX dimensions.
 
     Parameters
     ----------
@@ -33,6 +33,7 @@ class N2VManipulate(ImageOnlyTransform):
         remove_center: bool = True,
         struct_mask_axis: Literal["horizontal", "vertical", "none"] = "none",
         struct_mask_span: int = 5,
+        seed: Optional[int] = None,  # TODO use in pixel manipulation
     ):
         """Constructor.
 
@@ -50,8 +51,9 @@ class N2VManipulate(ImageOnlyTransform):
             StructN2V mask axis, by default "none"
         struct_mask_span : int, optional
             StructN2V mask span, by default 5
+        seed : Optional[int], optional
+            Random seed, by default None
         """
-        super().__init__(p=1)
         self.masked_pixel_percentage = masked_pixel_percentage
         self.roi_size = roi_size
         self.strategy = strategy
@@ -65,23 +67,26 @@ class N2VManipulate(ImageOnlyTransform):
                 span=struct_mask_span,
             )
 
-    def apply(
-        self, patch: np.ndarray, **kwargs: Any
+        # numpy random generator
+        self.rng = np.random.default_rng(seed=seed)
+
+    def __call__(
+        self, patch: np.ndarray, *args: Any, **kwargs: Any
     ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
         """Apply the transform to the image.
 
         Parameters
         ----------
         image : np.ndarray
-            Image or image patch, 2D or 3D, shape (y, x, c) or (z, y, x, c).
+            Image or image patch, 2D or 3D, shape C(Z)YX.
         """
         masked = np.zeros_like(patch)
         mask = np.zeros_like(patch)
         if self.strategy == SupportedPixelManipulation.UNIFORM:
             # Iterate over the channels to apply manipulation separately
-            for c in range(patch.shape[-1]):
-                masked[..., c], mask[..., c] = uniform_manipulate(
-                    patch=patch[..., c],
+            for c in range(patch.shape[0]):
+                masked[c, ...], mask[c, ...] = uniform_manipulate(
+                    patch=patch[c, ...],
                     mask_pixel_percentage=self.masked_pixel_percentage,
                     subpatch_size=self.roi_size,
                     remove_center=self.remove_center,
@@ -89,9 +94,9 @@ class N2VManipulate(ImageOnlyTransform):
                 )
         elif self.strategy == SupportedPixelManipulation.MEDIAN:
             # Iterate over the channels to apply manipulation separately
-            for c in range(patch.shape[-1]):
-                masked[..., c], mask[..., c] = median_manipulate(
-                    patch=patch[..., c],
+            for c in range(patch.shape[0]):
+                masked[c, ...], mask[c, ...] = median_manipulate(
+                    patch=patch[c, ...],
                     mask_pixel_percentage=self.masked_pixel_percentage,
                     subpatch_size=self.roi_size,
                     struct_params=self.struct_mask,
@@ -101,13 +106,3 @@ class N2VManipulate(ImageOnlyTransform):
 
         # TODO why return patch?
         return masked, patch, mask
-
-    def get_transform_init_args_names(self) -> Tuple[str, ...]:
-        """Get the transform parameters.
-
-        Returns
-        -------
-        Tuple[str, ...]
-            Transform parameters.
-        """
-        return ("roi_size", "masked_pixel_percentage", "strategy", "struct_mask")

careamics/transforms/nd_flip.py

@@ -1,93 +1,67 @@
-from typing import Any, Dict, Tuple
+from typing import Optional, Tuple
 
 import numpy as np
-from albumentations import DualTransform
 
+from careamics.transforms.transform import Transform
 
-class NDFlip(DualTransform):
+
+class NDFlip(Transform):
     """Flip ND arrays on a single axis.
 
     This transform ignores singleton axes and randomly flips one of the other
-    axes, to the exception of the first and last axes (sample and channels).
+    last two axes.
 
-    This transform expects (Z)YXC dimensions.
+    This transform expects C(Z)YX dimensions.
     """
 
-    def __init__(self, p: float = 0.5, is_3D: bool = False, flip_z: bool = True):
+    def __init__(self, seed: Optional[int] = None):
         """Constructor.
 
         Parameters
         ----------
-        p : float, optional
-            Probability to apply the transform, by default 0.5
-        is_3D : bool, optional
-            Whether the data is 3D, by default False
-        flip_z : bool, optional
-            Whether to flip Z dimension, by default True
+        seed : Optional[int], optional
+            Random seed, by default None
         """
-        super().__init__(p=p)
-
-        self.is_3D = is_3D
-        self.flip_z = flip_z
-
         # "flippable" axes
-        if is_3D:
-            self.axis_indices = [0, 1, 2] if flip_z else [1, 2]
-        else:
-            self.axis_indices = [0, 1]
+        self.axis_indices = [-2, -1]
 
-    def get_params(self, **kwargs: Any) -> Dict[str, int]:
-        """Get the transform parameters.
-
-        Returns
-        -------
-        Dict[str, int]
-            Transform parameters.
-        """
-        return {"flip_axis": np.random.choice(self.axis_indices)}
+        # numpy random generator
+        self.rng = np.random.default_rng(seed=seed)
 
-    def apply(self, patch: np.ndarray, flip_axis: int, **kwargs: Any) -> np.ndarray:
-        """Apply the transform to the image.
+    def __call__(
+        self, patch: np.ndarray, target: Optional[np.ndarray] = None
+    ) -> Tuple[np.ndarray, Optional[np.ndarray]]:
+        """Apply the transform to the source patch and the target (optional).
 
         Parameters
         ----------
         patch : np.ndarray
-            Image or image patch, 2D or 3D, shape (y, x, c) or (z, y, x, c).
-        flip_axis : int
-            Axis along which to flip the patch.
+            Patch, 2D or 3D, shape C(Z)YX.
+        target : Optional[np.ndarray], optional
+            Target for the patch, by default None
+
+        Returns
+        -------
+        Tuple[np.ndarray, Optional[np.ndarray]]
+            Transformed patch and target.
         """
-        if len(patch.shape) == 3 and self.is_3D:
-            raise ValueError(
-                "Incompatible patch shape and dimensionality. ZYXC patch shape "
-                "expected, but got YXC shape."
-            )
+        # choose an axis to flip
+        axis = self.rng.choice(self.axis_indices)
+
+        patch_transformed = self._apply(patch, axis)
+        target_transformed = self._apply(target, axis) if target is not None else None
 
-        return np.ascontiguousarray(np.flip(patch, axis=flip_axis))
+        return patch_transformed, target_transformed
 
-    def apply_to_mask(
-        self, mask: np.ndarray, flip_axis: int, **kwargs: Any
-    ) -> np.ndarray:
-        """Apply the transform to the mask.
+    def _apply(self, patch: np.ndarray, axis: int) -> np.ndarray:
+        """Apply the transform to the image.
 
         Parameters
         ----------
-        mask : np.ndarray
-            Mask or mask patch, 2D or 3D, shape (y, x, c) or (z, y, x, c).
-        """
-        if len(mask.shape) == 3 and self.is_3D:
-            raise ValueError(
-                "Incompatible mask shape and dimensionality. ZYXC patch shape "
-                "expected, but got YXC shape."
-            )
-
-        return np.ascontiguousarray(np.flip(mask, axis=flip_axis))
-
-    def get_transform_init_args_names(self, **kwargs: Any) -> Tuple[str, ...]:
-        """Get the transform arguments names.
-
-        Returns
-        -------
-        Tuple[str, ...]
-            Transform arguments names.
+        patch : np.ndarray
+            Image or image patch, 2D or 3D, shape C(Z)YX.
+        axis : int
+            Axis to flip.
         """
-        return ("is_3D", "flip_z")
+        # TODO why ascontiguousarray?
+        return np.ascontiguousarray(np.flip(patch, axis=axis))