careamics 0.0.5__py3-none-any.whl → 0.0.6__py3-none-any.whl

careamics/lightning/train_data_module.py

@@ -9,7 +9,7 @@ import pytorch_lightning as L
 from numpy.typing import NDArray
 from torch.utils.data import DataLoader, IterableDataset
 
-from careamics.config import DataConfig
+from careamics.config.data import DataConfig, GeneralDataConfig, N2VDataConfig
 from careamics.config.support import SupportedData
 from careamics.config.transformations import TransformModel
 from careamics.dataset.dataset_utils import (
@@ -119,7 +119,7 @@ class TrainDataModule(L.LightningDataModule):
 
     def __init__(
         self,
-        data_config: DataConfig,
+        data_config: GeneralDataConfig,
         train_data: Union[Path, str, NDArray],
         val_data: Optional[Union[Path, str, NDArray]] = None,
         train_data_target: Optional[Union[Path, str, NDArray]] = None,
@@ -219,7 +219,7 @@ class TrainDataModule(L.LightningDataModule):
             )
 
         # configuration
-        self.data_config: DataConfig = data_config
+        self.data_config: GeneralDataConfig = data_config
         self.data_type: str = data_config.data_type
         self.batch_size: int = data_config.batch_size
         self.use_in_memory: bool = use_in_memory
@@ -502,12 +502,23 @@ def create_train_datamodule(
     """Create a TrainDataModule.
 
     This function is used to explicitly pass the parameters usually contained in a
-    `data_model` configuration to a TrainDataModule.
+    `GenericDataConfig` to a TrainDataModule.
 
     Since the lightning datamodule has no access to the model, make sure that the
     parameters passed to the datamodule are consistent with the model's requirements and
     are coherent.
 
+    By default, the train DataModule will be set for Noise2Void if no target data is
+    provided. That means that it will add a `N2VManipulateModel` transformation to the
+    list of augmentations. The default augmentations are XY flip, XY rotation, and N2V
+    pixel manipulation. If you pass a training target data, the default behaviour is to
+    train a supervised model. It will use the default XY flip and rotation
+    augmentations.
+
+    To use a different set of transformations, you can pass a list of transforms to
+    `transforms`. Note that if you intend to use Noise2Void, you should add
+    `N2VManipulateModel` as the last transform in the list of transformations.
+
     The data module can be used with Path, str or numpy arrays. In the case of
     numpy arrays, it loads and computes all the patches in memory. For Path and str
     inputs, it calculates the total file size and estimate whether it can fit in
@@ -518,11 +529,6 @@ def create_train_datamodule(
     To use array data, set `data_type` to `array` and pass a numpy array to
     `train_data`.
 
-    In particular, N2V requires a specific transformation (N2V manipulates), which is
-    not compatible with supervised training. The default transformations applied to the
-    training patches are defined in `careamics.config.data_model`. To use different
-    transformations, pass a list of transforms. See examples for more details.
-
     By default, CAREamics only supports types defined in
     `careamics.config.support.SupportedData`. To read custom data types, you can set
     `data_type` to `custom` and provide a function that returns a numpy array from a
@@ -627,12 +633,12 @@ def create_train_datamodule(
     transforms:
     >>> import numpy as np
     >>> from careamics.lightning import create_train_datamodule
+    >>> from careamics.config.transformations import XYFlipModel, N2VManipulateModel
     >>> from careamics.config.support import SupportedTransform
     >>> my_array = np.arange(256).reshape(16, 16)
     >>> my_transforms = [
-    ...     {
-    ...         "name": SupportedTransform.XY_FLIP.value,
-    ...     }
+    ...     XYFlipModel(flip_y=False),
+    ...     N2VManipulateModel()
     ... ]
     >>> data_module = create_train_datamodule(
     ...     train_data=my_array,
@@ -659,21 +665,15 @@ def create_train_datamodule(
     if transforms is not None:
         data_dict["transforms"] = transforms
 
-    # validate configuration
-    data_config = DataConfig(**data_dict)
+    # TODO not compatible with HDN, consider adding an argument for n2v/hdn
+    if train_target_data is None:
+        data_config: GeneralDataConfig = N2VDataConfig(**data_dict)
+        assert isinstance(data_config, N2VDataConfig)
 
-    # N2V specific checks, N2V, structN2V, and transforms
-    if data_config.has_n2v_manipulate():
-        # there is not target, n2v2 and structN2V can be changed
-        if train_target_data is None:
-            data_config.set_N2V2(use_n2v2)
-            data_config.set_structN2V_mask(struct_n2v_axis, struct_n2v_span)
-        else:
-            raise ValueError(
-                "Cannot have both supervised training (target data) and "
-                "N2V manipulation in the transforms. Pass a list of transforms "
-                "that is compatible with your supervised training."
-            )
+        data_config.set_n2v2(use_n2v2)
+        data_config.set_structN2V_mask(struct_n2v_axis, struct_n2v_span)
+    else:
+        data_config = DataConfig(**data_dict)
 
     # sanity check on the dataloader parameters
     if "batch_size" in dataloader_params:

careamics/losses/__init__.py

@@ -1,13 +1,13 @@
 """Losses module."""
 
 __all__ = [
+    "denoisplit_loss",
+    "denoisplit_musplit_loss",
     "loss_factory",
     "mae_loss",
     "mse_loss",
-    "n2v_loss",
-    "denoisplit_loss",
     "musplit_loss",
-    "denoisplit_musplit_loss",
+    "n2v_loss",
 ]
 
 from .fcn.losses import mae_loss, mse_loss, n2v_loss

careamics/model_io/__init__.py

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

careamics/model_io/bioimage/__init__.py

@@ -1,10 +1,10 @@
 """Bioimage Model Zoo format functions."""
 
 __all__ = [
+    "create_env_text",
     "create_model_description",
     "extract_model_path",
     "get_unzip_path",
-    "create_env_text",
 ]
 
 from .bioimage_utils import create_env_text, get_unzip_path

careamics/model_io/bioimage/_readme_factory.py

@@ -55,7 +55,7 @@ def readme_factory(
         readme.touch()
 
         # algorithm pretty name
-        algorithm_flavour = config.get_algorithm_flavour()
+        algorithm_flavour = config.get_algorithm_friendly_name()
         algorithm_pretty_name = algorithm_flavour + " - CAREamics"
 
         description = [f"# {algorithm_pretty_name}\n\n"]

careamics/model_io/bioimage/model_description.py

@@ -1,7 +1,7 @@
 """Module use to build BMZ model description."""
 
 from pathlib import Path
-from typing import List, Optional, Tuple, Union
+from typing import Optional, Union
 
 import numpy as np
 from bioimageio.spec._internal.io import resolve_and_extract
@@ -28,17 +28,17 @@ from bioimageio.spec.model.v0_5 import (
     WeightsDescr,
 )
 
-from careamics.config import Configuration, DataConfig
+from careamics.config import Configuration, GeneralDataConfig
 
 from ._readme_factory import readme_factory
 
 
 def _create_axes(
     array: np.ndarray,
-    data_config: DataConfig,
-    channel_names: Optional[List[str]] = None,
+    data_config: GeneralDataConfig,
+    channel_names: Optional[list[str]] = None,
     is_input: bool = True,
-) -> List[AxisBase]:
+) -> list[AxisBase]:
     """Create axes description.
 
     Array shape is expected to be SC(Z)YX.
@@ -49,15 +49,15 @@ def _create_axes(
         Array.
     data_config : DataModel
         CAREamics data configuration.
-    channel_names : Optional[List[str]], optional
+    channel_names : Optional[list[str]], optional
         Channel names, by default None.
     is_input : bool, optional
         Whether the axes are input axes, by default True.
 
     Returns
     -------
-    List[AxisBase]
-        List of axes description.
+    list[AxisBase]
+        list of axes description.
 
     Raises
     ------
@@ -102,11 +102,11 @@ def _create_axes(
 def _create_inputs_ouputs(
     input_array: np.ndarray,
     output_array: np.ndarray,
-    data_config: DataConfig,
+    data_config: GeneralDataConfig,
     input_path: Union[Path, str],
     output_path: Union[Path, str],
-    channel_names: Optional[List[str]] = None,
-) -> Tuple[InputTensorDescr, OutputTensorDescr]:
+    channel_names: Optional[list[str]] = None,
+) -> tuple[InputTensorDescr, OutputTensorDescr]:
     """Create input and output tensor description.
 
     Input and output paths must point to a `.npy` file.
@@ -123,12 +123,12 @@ def _create_inputs_ouputs(
         Path to input .npy file.
     output_path : Union[Path, str]
         Path to output .npy file.
-    channel_names : Optional[List[str]], optional
+    channel_names : Optional[list[str]], optional
         Channel names, by default None.
 
     Returns
     -------
-    Tuple[InputTensorDescr, OutputTensorDescr]
+    tuple[InputTensorDescr, OutputTensorDescr]
         Input and output tensor descriptions.
     """
     input_axes = _create_axes(input_array, data_config, channel_names)
@@ -188,7 +188,7 @@ def create_model_description(
     name: str,
     general_description: str,
     data_description: str,
-    authors: List[Author],
+    authors: list[Author],
     inputs: Union[Path, str],
     outputs: Union[Path, str],
     weights_path: Union[Path, str],
@@ -197,7 +197,7 @@ def create_model_description(
     config_path: Union[Path, str],
     env_path: Union[Path, str],
     covers: list[Union[Path, str]],
-    channel_names: Optional[List[str]] = None,
+    channel_names: Optional[list[str]] = None,
     model_version: str = "0.1.0",
 ) -> ModelDescr:
     """Create model description.
@@ -212,7 +212,7 @@ def create_model_description(
         General description of the model.
     data_description : str
         Description of the data the model was trained on.
-    authors : List[Author]
+    authors : list[Author]
         Authors of the model.
     inputs : Union[Path, str]
         Path to input .npy file.
@@ -230,7 +230,7 @@ def create_model_description(
         Path to environment file.
     covers : list of pathlib.Path or str
         Paths to cover images.
-    channel_names : Optional[List[str]], optional
+    channel_names : Optional[list[str]], optional
         Channel names, by default None.
     model_version : str, default "0.1.0"
         Model version.

careamics/model_io/bmz_io.py

@@ -2,7 +2,7 @@
 
 import tempfile
 from pathlib import Path
-from typing import List, Optional, Tuple, Union
+from typing import Optional, Union
 
 import numpy as np
 import pkg_resources
@@ -87,11 +87,11 @@ def export_to_bmz(
     model_name: str,
     general_description: str,
     data_description: str,
-    authors: List[dict],
+    authors: list[dict],
     input_array: np.ndarray,
     output_array: np.ndarray,
     covers: Optional[list[Union[Path, str]]] = None,
-    channel_names: Optional[List[str]] = None,
+    channel_names: Optional[list[str]] = None,
     model_version: str = "0.1.0",
 ) -> None:
     """Export the model to BioImage Model Zoo format.
@@ -115,7 +115,7 @@ def export_to_bmz(
         General description of the model.
     data_description : str
         Description of the data the model was trained on.
-    authors : List[dict]
+    authors : list[dict]
         Authors of the model.
     input_array : np.ndarray
         Input array, should not have been normalized.
@@ -123,24 +123,13 @@ def export_to_bmz(
         Output array, should have been denormalized.
     covers : list of pathlib.Path or str, default=None
         Paths to the cover images.
-    channel_names : Optional[List[str]], optional
+    channel_names : Optional[list[str]], optional
         Channel names, by default None.
     model_version : str, default="0.1.0"
         Model version.
-
-    Raises
-    ------
-    ValueError
-        If the model is a Custom model.
     """
     path_to_archive = Path(path_to_archive)
 
-    # method is not compatible with Custom models
-    if config.algorithm_config.model.architecture == SupportedArchitecture.CUSTOM:
-        raise ValueError(
-            "Exporting Custom models to BioImage Model Zoo format is not supported."
-        )
-
     if path_to_archive.suffix != ".zip":
         raise ValueError(
             f"Path to archive must point to a zip file, got {path_to_archive}."
@@ -212,7 +201,7 @@ def export_to_bmz(
 
 def load_from_bmz(
     path: Union[Path, str, HttpUrl]
-) -> Tuple[Union[FCNModule, VAEModule], Configuration]:
+) -> tuple[Union[FCNModule, VAEModule], Configuration]:
     """Load a model from a BioImage Model Zoo archive.
 
     Parameters

careamics/model_io/model_io_utils.py

@@ -1,11 +1,11 @@
 """Utility functions to load pretrained models."""
 
 from pathlib import Path
-from typing import Tuple, Union
+from typing import Union
 
 import torch
 
-from careamics.config import Configuration
+from careamics.config import Configuration, configuration_factory
 from careamics.lightning.lightning_module import FCNModule, VAEModule
 from careamics.model_io.bmz_io import load_from_bmz
 from careamics.utils import check_path_exists
@@ -13,7 +13,7 @@ from careamics.utils import check_path_exists
 
 def load_pretrained(
     path: Union[Path, str]
-) -> Tuple[Union[FCNModule, VAEModule], Configuration]:
+) -> tuple[Union[FCNModule, VAEModule], Configuration]:
     """
     Load a pretrained model from a checkpoint or a BioImage Model Zoo model.
 
@@ -26,8 +26,8 @@ def load_pretrained(
 
     Returns
     -------
-    Tuple[CAREamicsKiln, Configuration]
-        Tuple of CAREamics model and its configuration.
+    tuple[CAREamicsKiln, Configuration]
+        tuple of CAREamics model and its configuration.
 
     Raises
     ------
@@ -48,7 +48,7 @@ def load_pretrained(
 
 def _load_checkpoint(
     path: Union[Path, str]
-) -> Tuple[Union[FCNModule, VAEModule], Configuration]:
+) -> tuple[Union[FCNModule, VAEModule], Configuration]:
     """
     Load a model from a checkpoint and return both model and configuration.
 
@@ -59,8 +59,8 @@ def _load_checkpoint(
 
     Returns
     -------
-    Tuple[CAREamicsKiln, Configuration]
-        Tuple of CAREamics model and its configuration.
+    tuple[CAREamicsKiln, Configuration]
+        tuple of CAREamics model and its configuration.
 
     Raises
     ------
@@ -92,4 +92,4 @@ def _load_checkpoint(
             f"{cfg_dict['algorithm_config']['model']['architecture']}"
         )
 
-    return model, Configuration(**cfg_dict)
+    return model, configuration_factory(cfg_dict)

careamics/models/layers.py

@@ -4,7 +4,7 @@ Layer module.
 This submodule contains layers used in the CAREamics models.
 """
 
-from typing import List, Optional, Tuple, Union
+from typing import Optional, Union
 
 import torch
 import torch.nn as nn
@@ -157,22 +157,22 @@ class Conv_Block(nn.Module):
 
 
 def _unpack_kernel_size(
-    kernel_size: Union[Tuple[int, ...], int], dim: int
-) -> Tuple[int, ...]:
+    kernel_size: Union[tuple[int, ...], int], dim: int
+) -> tuple[int, ...]:
     """Unpack kernel_size to a tuple of ints.
 
     Inspired by Kornia implementation. TODO: link
 
     Parameters
     ----------
-    kernel_size : Union[Tuple[int, ...], int]
+    kernel_size : Union[tuple[int, ...], int]
         Kernel size.
     dim : int
         Number of dimensions.
 
     Returns
     -------
-    Tuple[int, ...]
+    tuple[int, ...]
         Kernel size tuple.
     """
     if isinstance(kernel_size, int):
@@ -183,20 +183,20 @@ def _unpack_kernel_size(
 
 
 def _compute_zero_padding(
-    kernel_size: Union[Tuple[int, ...], int], dim: int
-) -> Tuple[int, ...]:
+    kernel_size: Union[tuple[int, ...], int], dim: int
+) -> tuple[int, ...]:
     """Utility function that computes zero padding tuple.
 
     Parameters
     ----------
-    kernel_size : Union[Tuple[int, ...], int]
+    kernel_size : Union[tuple[int, ...], int]
         Kernel size.
     dim : int
         Number of dimensions.
 
     Returns
     -------
-    Tuple[int, ...]
+    tuple[int, ...]
         Zero padding tuple.
     """
     kernel_dims = _unpack_kernel_size(kernel_size, dim)
@@ -245,8 +245,8 @@ def get_pascal_kernel_1d(
     >>> get_pascal_kernel_1d(6)
     tensor([ 1.,  5., 10., 10.,  5.,  1.])
     """
-    pre: List[float] = []
-    cur: List[float] = []
+    pre: list[float] = []
+    cur: list[float] = []
     for i in range(kernel_size):
         cur = [1.0] * (i + 1)
 
@@ -266,7 +266,7 @@ def get_pascal_kernel_1d(
 
 
 def _get_pascal_kernel_nd(
-    kernel_size: Union[Tuple[int, int], int],
+    kernel_size: Union[tuple[int, int], int],
     norm: bool = True,
     dim: int = 2,
     *,
@@ -282,7 +282,7 @@ def _get_pascal_kernel_nd(
 
     Parameters
     ----------
-    kernel_size : Union[Tuple[int, int], int]
+    kernel_size : Union[tuple[int, int], int]
         Kernel size for the pascal kernel.
     norm : bool
         Normalize the kernel, by default True.
@@ -420,7 +420,7 @@ class MaxBlurPool(nn.Module):
     ----------
     dim : int
         Toggles between 2D and 3D.
-    kernel_size : Union[Tuple[int, int], int]
+    kernel_size : Union[tuple[int, int], int]
         Kernel size for max pooling.
     stride : int
         Stride for pooling.
@@ -433,7 +433,7 @@ class MaxBlurPool(nn.Module):
     def __init__(
         self,
         dim: int,
-        kernel_size: Union[Tuple[int, int], int],
+        kernel_size: Union[tuple[int, int], int],
         stride: int = 2,
         max_pool_size: int = 2,
         ceil_mode: bool = False,
@@ -444,7 +444,7 @@ class MaxBlurPool(nn.Module):
         ----------
         dim : int
             Dimension of the convolution.
-        kernel_size : Union[Tuple[int, int], int]
+        kernel_size : Union[tuple[int, int], int]
             Kernel size for max pooling.
         stride : int, optional
             Stride, by default 2.

careamics/models/lvae/lvae.py

@@ -12,8 +12,6 @@ import numpy as np
 import torch
 import torch.nn as nn
 
-from careamics.config.architectures import register_model
-
 from ..activation import get_activation
 from .layers import (
     BottomUpDeterministicResBlock,
@@ -25,7 +23,6 @@ from .layers import (
 from .utils import Interpolate, ModelType, crop_img_tensor
 
 
-@register_model("LVAE")
 class LadderVAE(nn.Module):
     """
     Constructor.

careamics/models/model_factory.py

@@ -1,8 +1,4 @@
-"""
-Model factory.
-
-Model creation factory functions.
-"""
+"""Model creation factory functions."""
 
 from __future__ import annotations
 
@@ -10,10 +6,6 @@ from typing import TYPE_CHECKING, Union
 
 import torch
 
-from careamics.config.architectures import (
-    CustomModel,
-    get_custom_model,
-)
 from careamics.config.support import SupportedArchitecture
 from careamics.models.lvae import LadderVAE as LVAE
 from careamics.models.unet import UNet
@@ -21,7 +13,6 @@ from careamics.utils import get_logger
 
 if TYPE_CHECKING:
     from careamics.config.architectures import (
-        CustomModel,
         LVAEModel,
         UNetModel,
     )
@@ -31,7 +22,7 @@ logger = get_logger(__name__)
 
 
 def model_factory(
-    model_configuration: Union[UNetModel, LVAEModel, CustomModel],
+    model_configuration: Union[UNetModel, LVAEModel],
 ) -> torch.nn.Module:
     """
     Deep learning model factory.
@@ -57,10 +48,6 @@ def model_factory(
         return UNet(**model_configuration.model_dump())
     elif model_configuration.architecture == SupportedArchitecture.LVAE:
         return LVAE(**model_configuration.model_dump())
-    elif model_configuration.architecture == SupportedArchitecture.CUSTOM:
-        assert isinstance(model_configuration, CustomModel)
-        model = get_custom_model(model_configuration.name)
-        return model(**model_configuration.model_dump())
     else:
         raise NotImplementedError(
             f"Model {model_configuration.architecture} is not implemented or unknown."

careamics/models/unet.py

@@ -4,7 +4,7 @@ UNet model.
 A UNet encoder, decoder and complete model.
 """
 
-from typing import Any, List, Tuple, Union
+from typing import Any, Union
 
 import torch
 import torch.nn as nn
@@ -104,7 +104,7 @@ class UnetEncoder(nn.Module):
             encoder_blocks.append(self.pooling)
         self.encoder_blocks = nn.ModuleList(encoder_blocks)
 
-    def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
+    def forward(self, x: torch.Tensor) -> list[torch.Tensor]:
         """
         Forward pass.
 
@@ -115,7 +115,7 @@ class UnetEncoder(nn.Module):
 
         Returns
         -------
-        List[torch.Tensor]
+        list[torch.Tensor]
             Output of each encoder block (skip connections) and final output of the
             encoder.
         """
@@ -202,7 +202,7 @@ class UnetDecoder(nn.Module):
             groups=self.groups,
         )
 
-        decoder_blocks: List[nn.Module] = []
+        decoder_blocks: list[nn.Module] = []
         for n in range(depth):
             decoder_blocks.append(upsampling)
             in_channels = (num_channels_init * 2 ** (depth - n)) * groups
@@ -230,7 +230,7 @@ class UnetDecoder(nn.Module):
 
         Parameters
         ----------
-        *features :  List[torch.Tensor]
+        *features :  list[torch.Tensor]
             List containing the output of each encoder block(skip connections) and final
             output of the encoder.
 
@@ -240,7 +240,7 @@ class UnetDecoder(nn.Module):
             Output of the decoder.
         """
         x: torch.Tensor = features[0]
-        skip_connections: Tuple[torch.Tensor, ...] = features[-1:0:-1]
+        skip_connections: tuple[torch.Tensor, ...] = features[-1:0:-1]
 
         x = self.bottleneck(x)
 
@@ -289,10 +289,10 @@ class UnetDecoder(nn.Module):
         m = A.shape[1] // groups
         n = B.shape[1] // groups
 
-        A_groups: List[torch.Tensor] = [
+        A_groups: list[torch.Tensor] = [
             A[:, i * m : (i + 1) * m] for i in range(groups)
         ]
-        B_groups: List[torch.Tensor] = [
+        B_groups: list[torch.Tensor] = [
             B[:, i * n : (i + 1) * n] for i in range(groups)
         ]
 

careamics/prediction_utils/__init__.py

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