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

careamics/model_io/__init__.py

@@ -1,6 +1,5 @@
 """Model I/O utilities."""
 
-
 __all__ = ["load_pretrained", "export_to_bmz"]
 
 

careamics/model_io/bioimage/_readme_factory.py

@@ -1,4 +1,5 @@
 """Functions used to create a README.md file for BMZ export."""
+
 from pathlib import Path
 from typing import Optional
 
@@ -117,4 +118,4 @@ def readme_factory(
 
         readme.write_text("".join(description))
 
-    return readme
+        return readme.absolute()

careamics/model_io/bioimage/bioimage_utils.py

@@ -1,4 +1,5 @@
 """Bioimage.io utils."""
+
 from pathlib import Path
 from typing import Union
 

careamics/model_io/bioimage/model_description.py

@@ -1,4 +1,5 @@
 """Module use to build BMZ model description."""
+
 from pathlib import Path
 from typing import List, Optional, Tuple, Union
 
@@ -26,14 +27,14 @@ from bioimageio.spec.model.v0_5 import (
     WeightsDescr,
 )
 
-from careamics.config import Configuration, DataModel
+from careamics.config import Configuration, DataConfig
 
 from ._readme_factory import readme_factory
 
 
 def _create_axes(
     array: np.ndarray,
-    data_config: DataModel,
+    data_config: DataConfig,
     channel_names: Optional[List[str]] = None,
     is_input: bool = True,
 ) -> List[AxisBase]:
@@ -100,7 +101,7 @@ def _create_axes(
 def _create_inputs_ouputs(
     input_array: np.ndarray,
     output_array: np.ndarray,
-    data_config: DataModel,
+    data_config: DataConfig,
     input_path: Union[Path, str],
     output_path: Union[Path, str],
     channel_names: Optional[List[str]] = None,

careamics/model_io/bmz_io.py

@@ -1,4 +1,5 @@
 """Function to export to the BioImage Model Zoo format."""
+
 import tempfile
 from pathlib import Path
 from typing import List, Optional, Tuple, Union
@@ -11,7 +12,7 @@ from torch import __version__, load, save
 
 from careamics.config import Configuration, load_configuration, save_configuration
 from careamics.config.support import SupportedArchitecture
-from careamics.lightning_module import CAREamicsKiln
+from careamics.lightning_module import CAREamicsModule
 
 from .bioimage import (
     create_env_text,
@@ -21,7 +22,7 @@ from .bioimage import (
 )
 
 
-def _export_state_dict(model: CAREamicsKiln, path: Union[Path, str]) -> Path:
+def _export_state_dict(model: CAREamicsModule, path: Union[Path, str]) -> Path:
     """
     Export the model state dictionary to a file.
 
@@ -51,7 +52,7 @@ def _export_state_dict(model: CAREamicsKiln, path: Union[Path, str]) -> Path:
     return path
 
 
-def _load_state_dict(model: CAREamicsKiln, path: Union[Path, str]) -> None:
+def _load_state_dict(model: CAREamicsModule, path: Union[Path, str]) -> None:
     """
     Load a model from a state dictionary.
 
@@ -73,7 +74,7 @@ def _load_state_dict(model: CAREamicsKiln, path: Union[Path, str]) -> None:
 
 # TODO break down in subfunctions
 def export_to_bmz(
-    model: CAREamicsKiln,
+    model: CAREamicsModule,
     config: Configuration,
     path: Union[Path, str],
     name: str,
@@ -177,7 +178,7 @@ def export_to_bmz(
         )
 
         # test model description
-        summary: ValidationSummary = test_model(model_description)
+        summary: ValidationSummary = test_model(model_description, decimal=0)
         if summary.status == "failed":
             raise ValueError(f"Model description test failed: {summary}")
 
@@ -185,7 +186,7 @@ def export_to_bmz(
         save_bioimageio_package(model_description, output_path=path)
 
 
-def load_from_bmz(path: Union[Path, str]) -> Tuple[CAREamicsKiln, Configuration]:
+def load_from_bmz(path: Union[Path, str]) -> Tuple[CAREamicsModule, Configuration]:
     """Load a model from a BioImage Model Zoo archive.
 
     Parameters
@@ -223,7 +224,7 @@ def load_from_bmz(path: Union[Path, str]) -> Tuple[CAREamicsKiln, Configuration]
     config = load_configuration(config_path)
 
     # create careamics lightning module
-    model = CAREamicsKiln(algorithm_config=config.algorithm_config)
+    model = CAREamicsModule(algorithm_config=config.algorithm_config)
 
     # load model state dictionary
     _load_state_dict(model, weights_path)

careamics/model_io/model_io_utils.py

@@ -6,12 +6,12 @@ from typing import Tuple, Union
 from torch import load
 
 from careamics.config import Configuration
-from careamics.lightning_module import CAREamicsKiln
+from careamics.lightning_module import CAREamicsModule
 from careamics.model_io.bmz_io import load_from_bmz
 from careamics.utils import check_path_exists
 
 
-def load_pretrained(path: Union[Path, str]) -> Tuple[CAREamicsKiln, Configuration]:
+def load_pretrained(path: Union[Path, str]) -> Tuple[CAREamicsModule, Configuration]:
     """
     Load a pretrained model from a checkpoint or a BioImage Model Zoo model.
 
@@ -44,7 +44,7 @@ def load_pretrained(path: Union[Path, str]) -> Tuple[CAREamicsKiln, Configuratio
         )
 
 
-def _load_checkpoint(path: Union[Path, str]) -> Tuple[CAREamicsKiln, Configuration]:
+def _load_checkpoint(path: Union[Path, str]) -> Tuple[CAREamicsModule, Configuration]:
     """
     Load a model from a checkpoint and return both model and configuration.
 
@@ -75,6 +75,6 @@ def _load_checkpoint(path: Union[Path, str]) -> Tuple[CAREamicsKiln, Configurati
             f"checkpoint: {checkpoint.keys()}"
         ) from e
 
-    model = CAREamicsKiln.load_from_checkpoint(path)
+    model = CAREamicsModule.load_from_checkpoint(path)
 
     return model, Configuration(**cfg_dict)

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

careamics/models/model_factory.py

@@ -3,6 +3,7 @@ Model factory.
 
 Model creation factory functions.
 """
+
 from typing import Union
 
 import torch

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")