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

careamics/lightning_datamodule.py

@@ -1,3 +1,4 @@
+"""Training and validation Lightning data modules."""
 from pathlib import Path
 from typing import Any, Callable, Dict, List, Literal, Optional, Union
 
@@ -6,7 +7,7 @@ import pytorch_lightning as L
 from albumentations import Compose
 from torch.utils.data import DataLoader
 
-from careamics.config import DataModel
+from careamics.config import DataConfig
 from careamics.config.data_model import TRANSFORMS_UNION
 from careamics.config.support import SupportedData
 from careamics.dataset.dataset_utils import (
@@ -28,9 +29,9 @@ DatasetType = Union[InMemoryDataset, PathIterableDataset]
 logger = get_logger(__name__)
 
 
-class CAREamicsWood(L.LightningDataModule):
+class CAREamicsTrainData(L.LightningDataModule):
     """
-    LightningDataModule for training and validation datasets.
+    CAREamics Ligthning training and validation data module.
 
     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
@@ -53,11 +54,70 @@ class CAREamicsWood(L.LightningDataModule):
 
     You can also provide a `fnmatch` and `Path.rglob` compatible expression (e.g.
     "*.czi") to filter the files extension using `extension_filter`.
+
+    Parameters
+    ----------
+    data_config : DataModel
+        Pydantic model for CAREamics data configuration.
+    train_data : Union[Path, str, np.ndarray]
+        Training data, can be a path to a folder, a file or a numpy array.
+    val_data : Optional[Union[Path, str, np.ndarray]], optional
+        Validation data, can be a path to a folder, a file or a numpy array, by
+        default None.
+    train_data_target : Optional[Union[Path, str, np.ndarray]], optional
+        Training target data, can be a path to a folder, a file or a numpy array, by
+        default None.
+    val_data_target : Optional[Union[Path, str, np.ndarray]], optional
+        Validation target data, can be a path to a folder, a file or a numpy array,
+        by default None.
+    read_source_func : Optional[Callable], optional
+        Function to read the source data, by default None. Only used for `custom`
+        data type (see DataModel).
+    extension_filter : str, optional
+        Filter for file extensions, by default "". Only used for `custom` data types
+        (see DataModel).
+    val_percentage : float, optional
+        Percentage of the training data to use for validation, by default 0.1. Only
+        used if `val_data` is None.
+    val_minimum_split : int, optional
+        Minimum number of patches or files to split from the training data for
+        validation, by default 5. Only used if `val_data` is None.
+    use_in_memory : bool, optional
+        Use in memory dataset if possible, by default True.
+
+    Attributes
+    ----------
+    data_config : DataModel
+        CAREamics data configuration.
+    data_type : SupportedData
+        Expected data type, one of "tiff", "array" or "custom".
+    batch_size : int
+        Batch size.
+    use_in_memory : bool
+        Whether to use in memory dataset if possible.
+    train_data : Union[Path, str, np.ndarray]
+        Training data.
+    val_data : Optional[Union[Path, str, np.ndarray]]
+        Validation data.
+    train_data_target : Optional[Union[Path, str, np.ndarray]]
+        Training target data.
+    val_data_target : Optional[Union[Path, str, np.ndarray]]
+        Validation target data.
+    val_percentage : float
+        Percentage of the training data to use for validation, if no validation data is
+        provided.
+    val_minimum_split : int
+        Minimum number of patches or files to split from the training data for
+        validation, if no validation data is provided.
+    read_source_func : Optional[Callable]
+        Function to read the source data, used if `data_type` is `custom`.
+    extension_filter : str
+        Filter for file extensions, used if `data_type` is `custom`.
     """
 
     def __init__(
         self,
-        data_config: DataModel,
+        data_config: DataConfig,
         train_data: Union[Path, str, np.ndarray],
         val_data: Optional[Union[Path, str, np.ndarray]] = None,
         train_data_target: Optional[Union[Path, str, np.ndarray]] = None,
@@ -98,6 +158,8 @@ class CAREamicsWood(L.LightningDataModule):
         val_minimum_split : int, optional
             Minimum number of patches or files to split from the training data for
             validation, by default 5. Only used if `val_data` is None.
+        use_in_memory : bool, optional
+            Use in memory dataset if possible, by default True.
 
         Raises
         ------
@@ -128,25 +190,30 @@ class CAREamicsWood(L.LightningDataModule):
         if data_config.data_type == SupportedData.CUSTOM and read_source_func is None:
             raise ValueError(
                 f"Data type {SupportedData.CUSTOM} is not allowed without "
-                f"specifying a `read_source_func`."
+                f"specifying a `read_source_func` and an `extension_filer`."
             )
 
-        # and that arrays are passed, if array type specified
-        elif data_config.data_type == SupportedData.ARRAY and not isinstance(
-            train_data, np.ndarray
+        # check correct input type
+        if (
+            isinstance(train_data, np.ndarray)
+            and data_config.data_type != SupportedData.ARRAY
         ):
             raise ValueError(
-                f"Expected array input (see configuration.data.data_type), but got "
-                f"{type(train_data)} instead."
+                f"Received a numpy array as input, but the data type was set to "
+                f"{data_config.data_type}. Set the data type in the configuration "
+                f"to {SupportedData.ARRAY} to train on numpy arrays."
             )
 
         # and that Path or str are passed, if tiff file type specified
-        elif data_config.data_type == SupportedData.TIFF and (
-            not isinstance(train_data, Path) and not isinstance(train_data, str)
+        elif (isinstance(train_data, Path) or isinstance(train_data, str)) and (
+            data_config.data_type != SupportedData.TIFF
+            and data_config.data_type != SupportedData.CUSTOM
         ):
             raise ValueError(
-                f"Expected Path or str input (see configuration.data.data_type), "
-                f"but got {type(train_data)} instead."
+                f"Received a path as input, but the data type was neither set to "
+                f"{SupportedData.TIFF} nor {SupportedData.CUSTOM}. Set the data type "
+                f"in the configuration to {SupportedData.TIFF} or "
+                f"{SupportedData.CUSTOM} to train on files."
             )
 
         # configuration
@@ -231,7 +298,15 @@ class CAREamicsWood(L.LightningDataModule):
                 validate_source_target_files(self.val_files, self.val_target_files)
 
     def setup(self, *args: Any, **kwargs: Any) -> None:
-        """Hook called at the beginning of fit, validate, or predict."""
+        """Hook called at the beginning of fit, validate, or predict.
+
+        Parameters
+        ----------
+        *args : Any
+            Unused.
+        **kwargs : Any
+            Unused.
+        """
         # if numpy array
         if self.data_type == SupportedData.ARRAY:
             # train dataset
@@ -353,9 +428,12 @@ class CAREamicsWood(L.LightningDataModule):
         )
 
 
-class CAREamicsTrainDataModule(CAREamicsWood):
+class TrainingDataWrapper(CAREamicsTrainData):
     """
-    LightningDataModule wrapper for training and validation datasets.
+    Wrapper around the CAREamics Lightning training data module.
+
+    This class is used to explicitely pass the parameters usually contained in a
+    `data_model` configuration.
 
     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
@@ -442,11 +520,11 @@ class CAREamicsTrainDataModule(CAREamicsWood):
 
     Examples
     --------
-    Create a CAREamicsTrainDataModule with default transforms with a numpy array:
+    Create a TrainingDataWrapper with default transforms with a numpy array:
     >>> import numpy as np
-    >>> from careamics import CAREamicsTrainDataModule
+    >>> from careamics import TrainingDataWrapper
     >>> my_array = np.arange(256).reshape(16, 16)
-    >>> data_module = CAREamicsTrainDataModule(
+    >>> data_module = TrainingDataWrapper(
     ...     train_data=my_array,
     ...     data_type="array",
     ...     patch_size=(8, 8),
@@ -457,12 +535,12 @@ class CAREamicsTrainDataModule(CAREamicsWood):
     For custom data types (those not supported by CAREamics), then one can pass a read
     function and a filter for the files extension:
     >>> import numpy as np
-    >>> from careamics import CAREamicsTrainDataModule
+    >>> from careamics import TrainingDataWrapper
     >>>
     >>> def read_npy(path):
     ...     return np.load(path)
     >>>
-    >>> data_module = CAREamicsTrainDataModule(
+    >>> data_module = TrainingDataWrapper(
     ...     train_data="path/to/data",
     ...     data_type="custom",
     ...     patch_size=(8, 8),
@@ -475,7 +553,7 @@ class CAREamicsTrainDataModule(CAREamicsWood):
     If you want to use a different set of transformations, you can pass a list of
     transforms:
     >>> import numpy as np
-    >>> from careamics import CAREamicsTrainDataModule
+    >>> from careamics import TrainingDataWrapper
     >>> from careamics.config.support import SupportedTransform
     >>> my_array = np.arange(256).reshape(16, 16)
     >>> my_transforms = [
@@ -488,7 +566,7 @@ class CAREamicsTrainDataModule(CAREamicsWood):
     ...         "name": SupportedTransform.N2V_MANIPULATE.value,
     ...     }
     ... ]
-    >>> data_module = CAREamicsTrainDataModule(
+    >>> data_module = TrainingDataWrapper(
     ...     train_data=my_array,
     ...     data_type="array",
     ...     patch_size=(8, 8),
@@ -628,7 +706,7 @@ class CAREamicsTrainDataModule(CAREamicsWood):
             data_dict["transforms"] = transforms
 
         # validate configuration
-        self.data_config = DataModel(**data_dict)
+        self.data_config = DataConfig(**data_dict)
 
         # N2V specific checks, N2V, structN2V, and transforms
         if (

careamics/lightning_module.py

@@ -3,7 +3,7 @@ from typing import Any, Optional, Union
 import pytorch_lightning as L
 from torch import Tensor, nn
 
-from careamics.config import AlgorithmModel
+from careamics.config import AlgorithmConfig
 from careamics.config.support import (
     SupportedAlgorithm,
     SupportedArchitecture,
@@ -17,7 +17,7 @@ from careamics.transforms import Denormalize, ImageRestorationTTA
 from careamics.utils.torch_utils import get_optimizer, get_scheduler
 
 
-class CAREamicsKiln(L.LightningModule):
+class CAREamicsModule(L.LightningModule):
     """
     CAREamics Lightning module.
 
@@ -38,7 +38,7 @@ class CAREamicsKiln(L.LightningModule):
         Learning rate scheduler name.
     """
 
-    def __init__(self, algorithm_config: Union[AlgorithmModel, dict]) -> None:
+    def __init__(self, algorithm_config: Union[AlgorithmConfig, dict]) -> None:
         """
         CAREamics Lightning module.
 
@@ -53,7 +53,7 @@ class CAREamicsKiln(L.LightningModule):
         super().__init__()
         # if loading from a checkpoint, AlgorithmModel needs to be instantiated
         if isinstance(algorithm_config, dict):
-            algorithm_config = AlgorithmModel(**algorithm_config)
+            algorithm_config = AlgorithmConfig(**algorithm_config)
 
         # create model and loss function
         self.model: nn.Module = model_factory(algorithm_config.model)
@@ -192,7 +192,7 @@ class CAREamicsKiln(L.LightningModule):
         }
 
 
-class CAREamicsModule(CAREamicsKiln):
+class CAREamicsModuleWrapper(CAREamicsModule):
     """Class defining the API for CAREamics Lightning layer.
 
     This class exposes parameters used to create an AlgorithmModel instance, triggering
@@ -287,6 +287,6 @@ class CAREamicsModule(CAREamicsKiln):
         algorithm_configuration["model"] = model_configuration
 
         # call the parent init using an AlgorithmModel instance
-        super().__init__(AlgorithmModel(**algorithm_configuration))
+        super().__init__(AlgorithmConfig(**algorithm_configuration))
 
         # TODO add load_from_checkpoint wrapper

careamics/lightning_prediction_datamodule.py

@@ -1,3 +1,4 @@
+"""Prediction Lightning data modules."""
 from pathlib import Path
 from typing import Any, Callable, List, Literal, Optional, Tuple, Union
 
@@ -7,7 +8,7 @@ from albumentations import Compose
 from torch.utils.data import DataLoader
 from torch.utils.data.dataloader import default_collate
 
-from careamics.config import InferenceModel
+from careamics.config import InferenceConfig
 from careamics.config.support import SupportedData
 from careamics.config.tile_information import TileInformation
 from careamics.dataset.dataset_utils import (
@@ -62,9 +63,9 @@ def _collate_tiles(batch: List[Tuple[np.ndarray, TileInformation]]) -> Any:
         return default_collate(new_batch)
 
 
-class CAREamicsClay(L.LightningDataModule):
+class CAREamicsPredictData(L.LightningDataModule):
     """
-    LightningDataModule for prediction dataset.
+    CAREamics Lightning prediction data module.
 
     The data module can be used with Path, str or numpy arrays. The data can be either
     a folder containing images or a single file.
@@ -79,7 +80,7 @@ class CAREamicsClay(L.LightningDataModule):
 
     Parameters
     ----------
-    prediction_config : InferenceModel
+    pred_config : InferenceModel
         Pydantic model for CAREamics prediction configuration.
     pred_data : Union[Path, str, np.ndarray]
         Prediction data, can be a path to a folder, a file or a numpy array.
@@ -93,7 +94,7 @@ class CAREamicsClay(L.LightningDataModule):
 
     def __init__(
         self,
-        prediction_config: InferenceModel,
+        pred_config: InferenceConfig,
         pred_data: Union[Path, str, np.ndarray],
         read_source_func: Optional[Callable] = None,
         extension_filter: str = "",
@@ -115,7 +116,7 @@ class CAREamicsClay(L.LightningDataModule):
 
         Parameters
         ----------
-        prediction_config : InferenceModel
+        pred_config : InferenceModel
             Pydantic model for CAREamics prediction configuration.
         pred_data : Union[Path, str, np.ndarray]
             Prediction data, can be a path to a folder, a file or a numpy array.
@@ -142,51 +143,53 @@ class CAREamicsClay(L.LightningDataModule):
         super().__init__()
 
         # check that a read source function is provided for custom types
-        if (
-            prediction_config.data_type == SupportedData.CUSTOM
-            and read_source_func is None
-        ):
+        if pred_config.data_type == SupportedData.CUSTOM and read_source_func is None:
             raise ValueError(
                 f"Data type {SupportedData.CUSTOM} is not allowed without "
-                f"specifying a `read_source_func`."
+                f"specifying a `read_source_func` and an `extension_filer`."
             )
 
-        # and that arrays are passed, if array type specified
-        elif prediction_config.data_type == SupportedData.ARRAY and not isinstance(
-            pred_data, np.ndarray
+        # check correct input type
+        if (
+            isinstance(pred_data, np.ndarray)
+            and pred_config.data_type != SupportedData.ARRAY
         ):
             raise ValueError(
-                f"Expected array input (see configuration.data.data_type), but got "
-                f"{type(pred_data)} instead."
+                f"Received a numpy array as input, but the data type was set to "
+                f"{pred_config.data_type}. Set the data type "
+                f"to {SupportedData.ARRAY} to predict on numpy arrays."
             )
 
         # and that Path or str are passed, if tiff file type specified
-        elif prediction_config.data_type == SupportedData.TIFF and not (
-            isinstance(pred_data, Path) or isinstance(pred_data, str)
+        elif (isinstance(pred_data, Path) or isinstance(pred_config, str)) and (
+            pred_config.data_type != SupportedData.TIFF
+            and pred_config.data_type != SupportedData.CUSTOM
         ):
             raise ValueError(
-                f"Expected Path or str input (see configuration.data.data_type), "
-                f"but got {type(pred_data)} instead."
+                f"Received a path as input, but the data type was neither set to "
+                f"{SupportedData.TIFF} nor {SupportedData.CUSTOM}. Set the data type "
+                f" to {SupportedData.TIFF} or "
+                f"{SupportedData.CUSTOM} to predict on files."
             )
 
         # configuration data
-        self.prediction_config = prediction_config
-        self.data_type = prediction_config.data_type
-        self.batch_size = prediction_config.batch_size
+        self.prediction_config = pred_config
+        self.data_type = pred_config.data_type
+        self.batch_size = pred_config.batch_size
         self.dataloader_params = dataloader_params
 
         self.pred_data = pred_data
-        self.tile_size = prediction_config.tile_size
-        self.tile_overlap = prediction_config.tile_overlap
+        self.tile_size = pred_config.tile_size
+        self.tile_overlap = pred_config.tile_overlap
 
         # read source function
-        if prediction_config.data_type == SupportedData.CUSTOM:
+        if pred_config.data_type == SupportedData.CUSTOM:
             # mypy check
             assert read_source_func is not None
 
             self.read_source_func: Callable = read_source_func
-        elif prediction_config.data_type != SupportedData.ARRAY:
-            self.read_source_func = get_read_func(prediction_config.data_type)
+        elif pred_config.data_type != SupportedData.ARRAY:
+            self.read_source_func = get_read_func(pred_config.data_type)
 
         self.extension_filter = extension_filter
 
@@ -238,9 +241,12 @@ class CAREamicsClay(L.LightningDataModule):
         )  # TODO check workers are used
 
 
-class CAREamicsPredictDataModule(CAREamicsClay):
+class PredictDataWrapper(CAREamicsPredictData):
     """
-    LightningDataModule wrapper of an inference dataset.
+    Wrapper around the CAREamics inference Lightning data module.
+
+    This class is used to explicitely pass the parameters usually contained in a
+    `inference_model` configuration.
 
     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
@@ -329,6 +335,12 @@ class CAREamicsPredictDataModule(CAREamicsClay):
             Prediction data.
         data_type : Union[Literal["array", "tiff", "custom"], SupportedData]
             Data type, see `SupportedData` for available options.
+        mean : float
+            Mean value for normalization, only used if Normalization is defined in the
+            transforms.
+        std : float
+            Standard deviation value for normalization, only used if Normalization is
+            defined in the transform.
         tile_size : List[int]
             Tile size, 2D or 3D tile size.
         tile_overlap : List[int]
@@ -339,12 +351,6 @@ class CAREamicsPredictDataModule(CAREamicsClay):
             Batch size.
         tta_transforms : bool, optional
             Use test time augmentation, by default True.
-        mean : Optional[float], optional
-            Mean value for normalization, only used if Normalization is defined, by
-            default None.
-        std : Optional[float], optional
-            Standard deviation value for normalization, only used if Normalization is
-            defined, by default None.
         transforms : Optional[Union[List[TRANSFORMS_UNION], Compose]], optional
             List of transforms to apply to prediction patches. If None, default
             transforms are applied.
@@ -374,7 +380,7 @@ class CAREamicsPredictDataModule(CAREamicsClay):
             prediction_dict["transforms"] = transforms
 
         # validate configuration
-        self.prediction_config = InferenceModel(**prediction_dict)
+        self.prediction_config = InferenceConfig(**prediction_dict)
 
         # sanity check on the dataloader parameters
         if "batch_size" in dataloader_params:
@@ -382,7 +388,7 @@ class CAREamicsPredictDataModule(CAREamicsClay):
             del dataloader_params["batch_size"]
 
         super().__init__(
-            prediction_config=self.prediction_config,
+            pred_config=self.prediction_config,
             pred_data=pred_data,
             read_source_func=read_source_func,
             extension_filter=extension_filter,

careamics/model_io/bioimage/model_description.py

@@ -26,14 +26,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 +100,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

@@ -11,7 +11,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 +21,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 +51,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 +73,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,
@@ -185,7 +185,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 +223,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-0.1.0rc3.dist-info → careamics-0.1.0rc4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: careamics
-Version: 0.1.0rc3
+Version: 0.1.0rc4
 Summary: Toolbox for running N2V and friends.
 Project-URL: homepage, https://careamics.github.io/
 Project-URL: repository, https://github.com/CAREamics/careamics