careamics 0.1.0rc7__py3-none-any.whl → 0.1.0rc8__py3-none-any.whl

careamics/config/configuration_model.py

@@ -9,7 +9,7 @@ from typing import Literal, Union
 
 import yaml
 from bioimageio.spec.generic.v0_3 import CiteEntry
-from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
+from pydantic import BaseModel, ConfigDict, field_validator, model_validator
 from typing_extensions import Self
 
 from .algorithm_model import AlgorithmConfig
@@ -147,20 +147,25 @@ class Configuration(BaseModel):
     )
 
     # version
-    version: Literal["0.1.0"] = Field(
-        default="0.1.0", description="Version of the CAREamics configuration."
-    )
+    version: Literal["0.1.0"] = "0.1.0"
+    """CAREamics configuration version."""
 
     # required parameters
-    experiment_name: str = Field(
-        ..., description="Name of the experiment, used to name logs and checkpoints."
-    )
+    experiment_name: str
+    """Name of the experiment, used to name logs and checkpoints."""
 
     # Sub-configurations
     algorithm_config: AlgorithmConfig
+    """Algorithm configuration, holding all parameters required to configure the
+    model."""
 
     data_config: DataConfig
+    """Data configuration, holding all parameters required to configure the training
+    data loader."""
+
     training_config: TrainingConfig
+    """Training configuration, holding all parameters required to configure the
+    training process."""
 
     @field_validator("experiment_name")
     @classmethod

careamics/config/data_model.py

@@ -55,8 +55,8 @@ class DataConfig(BaseModel):
     ...     axes="YX"
     ... )
 
-    To change the mean and std of the data:
-    >>> data.set_mean_and_std(image_means=[214.3], image_stds=[84.5])
+    To change the image_means and image_stds of the data:
+    >>> data.set_means_and_stds(image_means=[214.3], image_stds=[84.5])
 
     One can pass also a list of transformations, by keyword, using the
     SupportedTransform value:
@@ -80,22 +80,38 @@ class DataConfig(BaseModel):
     )
 
     # Dataset configuration
-    data_type: Literal["array", "tiff", "custom"]  # As defined in SupportedData
+    data_type: Literal["array", "tiff", "custom"]
+    """Type of input data, numpy.ndarray (array) or paths (tiff and custom), as defined
+    in SupportedData."""
+
+    axes: str
+    """Axes of the data, as defined in SupportedAxes."""
+
     patch_size: Union[list[int]] = Field(..., min_length=2, max_length=3)
+    """Patch size, as used during training."""
+
     batch_size: int = Field(default=1, ge=1, validate_default=True)
-    axes: str
+    """Batch size for training."""
 
     # Optional fields
     image_means: Optional[list[float]] = Field(
         default=None, min_length=0, max_length=32
     )
+    """Means of the data across channels, used for normalization."""
+
     image_stds: Optional[list[float]] = Field(default=None, min_length=0, max_length=32)
+    """Standard deviations of the data across channels, used for normalization."""
+
     target_means: Optional[list[float]] = Field(
         default=None, min_length=0, max_length=32
     )
+    """Means of the target data across channels, used for normalization."""
+
     target_stds: Optional[list[float]] = Field(
         default=None, min_length=0, max_length=32
     )
+    """Standard deviations of the target data across channels, used for
+    normalization."""
 
     transforms: list[TRANSFORMS_UNION] = Field(
         default=[
@@ -111,8 +127,11 @@ class DataConfig(BaseModel):
         ],
         validate_default=True,
     )
+    """List of transformations to apply to the data, available transforms are defined
+    in SupportedTransform. The default values are set for Noise2Void."""
 
     dataloader_params: Optional[dict] = None
+    """Dictionary of PyTorch dataloader parameters."""
 
     @field_validator("patch_size")
     @classmethod
@@ -346,7 +365,7 @@ class DataConfig(BaseModel):
         if self.has_n2v_manipulate():
             self.transforms.pop(-1)
 
-    def set_mean_and_std(
+    def set_means_and_stds(
         self,
         image_means: Union[NDArray, tuple, list, None],
         image_stds: Union[NDArray, tuple, list, None],
@@ -354,20 +373,20 @@ class DataConfig(BaseModel):
         target_stds: Optional[Union[NDArray, tuple, list, None]] = None,
     ) -> None:
         """
-        Set mean and standard deviation of the data.
+        Set mean and standard deviation of the data across channels.
 
         This method should be used instead setting the fields directly, as it would
         otherwise trigger a validation error.
 
         Parameters
         ----------
-        image_means : NDArray or tuple or list
+        image_means : numpy.ndarray ,tuple or list
             Mean values for normalization.
-        image_stds : NDArray or tuple or list
+        image_stds : numpy.ndarray, tuple or list
             Standard deviation values for normalization.
-        target_means : NDArray or tuple or list, optional
+        target_means : numpy.ndarray, tuple or list, optional
             Target mean values for normalization, by default ().
-        target_stds : NDArray or tuple or list, optional
+        target_stds : numpy.ndarray, tuple or list, optional
             Target standard deviation values for normalization, by default ().
         """
         # make sure we pass a list

careamics/config/inference_model.py

@@ -15,25 +15,35 @@ class InferenceConfig(BaseModel):
 
     model_config = ConfigDict(validate_assignment=True, arbitrary_types_allowed=True)
 
-    # Mandatory fields
     data_type: Literal["array", "tiff", "custom"]  # As defined in SupportedData
+    """Type of input data: numpy.ndarray (array) or path (tiff or custom)."""
+
     tile_size: Optional[Union[list[int]]] = Field(
         default=None, min_length=2, max_length=3
     )
+    """Tile size of prediction, only effective if `tile_overlap` is specified."""
+
     tile_overlap: Optional[Union[list[int]]] = Field(
         default=None, min_length=2, max_length=3
     )
+    """Overlap between tiles, only effective if `tile_size` is specified."""
 
     axes: str
+    """Data axes (TSCZYX) in the order of the input data."""
 
     image_means: list = Field(..., min_length=0, max_length=32)
+    """Mean values for each input channel."""
+
     image_stds: list = Field(..., min_length=0, max_length=32)
+    """Standard deviation values for each input channel."""
 
-    # only default TTAs are supported for now
+    # TODO only default TTAs are supported for now
     tta_transforms: bool = Field(default=True)
+    """Whether to apply test-time augmentation (all 90 degrees rotations and flips)."""
 
     # Dataloader parameters
     batch_size: int = Field(default=1, ge=1)
+    """Batch size for prediction."""
 
     @field_validator("tile_overlap")
     @classmethod

careamics/config/optimizer_models.py

@@ -45,6 +45,7 @@ class OptimizerModel(BaseModel):
 
     # Mandatory field
     name: Literal["Adam", "SGD"] = Field(default="Adam", validate_default=True)
+    """Name of the optimizer, supported optimizers are defined in SupportedOptimizer."""
 
     # Optional parameters, empty dict default value to allow filtering dictionary
     parameters: dict = Field(
@@ -53,6 +54,7 @@ class OptimizerModel(BaseModel):
         },
         validate_default=True,
     )
+    """Parameters of the optimizer, see PyTorch documentation for more details."""
 
     @field_validator("parameters")
     @classmethod
@@ -140,9 +142,13 @@ class LrSchedulerModel(BaseModel):
 
     # Mandatory field
     name: Literal["ReduceLROnPlateau", "StepLR"] = Field(default="ReduceLROnPlateau")
+    """Name of the learning rate scheduler, supported schedulers are defined in
+    SupportedScheduler."""
 
     # Optional parameters
     parameters: dict = Field(default={}, validate_default=True)
+    """Parameters of the learning rate scheduler, see PyTorch documentation for more
+    details."""
 
     @field_validator("parameters")
     @classmethod

careamics/config/support/supported_data.py

@@ -60,9 +60,9 @@ class SupportedData(str, BaseEnum):
         return super()._missing_(value)
 
     @classmethod
-    def get_extension(cls, data_type: Union[str, SupportedData]) -> str:
+    def get_extension_pattern(cls, data_type: Union[str, SupportedData]) -> str:
         """
-        Path.rglob and fnmatch compatible extension.
+        Get Path.rglob and fnmatch compatible extension.
 
         Parameters
         ----------
@@ -72,13 +72,38 @@ class SupportedData(str, BaseEnum):
         Returns
         -------
         str
-            Corresponding extension.
+            Corresponding extension pattern.
         """
         if data_type == cls.ARRAY:
-            raise NotImplementedError(f"Data {data_type} are not loaded from file.")
+            raise NotImplementedError(f"Data '{data_type}' is not loaded from a file.")
         elif data_type == cls.TIFF:
             return "*.tif*"
         elif data_type == cls.CUSTOM:
             return "*.*"
         else:
             raise ValueError(f"Data type {data_type} is not supported.")
+
+    @classmethod
+    def get_extension(cls, data_type: Union[str, SupportedData]) -> str:
+        """
+        Get file extension of corresponding data type.
+
+        Parameters
+        ----------
+        data_type : str or SupportedData
+            Data type.
+
+        Returns
+        -------
+        str
+            Corresponding extension.
+        """
+        if data_type == cls.ARRAY:
+            raise NotImplementedError(f"Data '{data_type}' is not loaded from a file.")
+        elif data_type == cls.TIFF:
+            return ".tiff"
+        elif data_type == cls.CUSTOM:
+            # TODO: improve this message
+            raise NotImplementedError("Custom extensions have to be passed elsewhere.")
+        else:
+            raise ValueError(f"Data type {data_type} is not supported.")

careamics/config/tile_information.py

@@ -19,10 +19,20 @@ class TileInformation(BaseModel):
     model_config = ConfigDict(validate_default=True)
 
     array_shape: tuple[int, ...]
+    """Shape of the original (untiled) array."""
+
     last_tile: bool = False
+    """Whether this tile is the last one of the array."""
+
     overlap_crop_coords: tuple[tuple[int, ...], ...]
+    """Inner coordinates of the tile where to crop the prediction in order to stitch
+    it back into the original image."""
+
     stitch_coords: tuple[tuple[int, ...], ...]
+    """Coordinates in the original image where to stitch the cropped tile back."""
+
     sample_id: int
+    """Sample ID of the tile."""
 
     @field_validator("array_shape")
     @classmethod

careamics/config/training_model.py

@@ -35,15 +35,19 @@ class TrainingConfig(BaseModel):
     )
 
     num_epochs: int = Field(default=20, ge=1)
+    """Number of epochs, greater than 0."""
 
     logger: Optional[Literal["wandb", "tensorboard"]] = None
+    """Logger to use during training. If None, no logger will be used. Available
+    loggers are defined in SupportedLogger."""
 
     checkpoint_callback: CheckpointModel = CheckpointModel()
+    """Checkpoint callback configuration."""
 
     early_stopping_callback: Optional[EarlyStoppingModel] = Field(
         default=None, validate_default=True
     )
-    # precision: Literal["64", "32", "16", "bf16"] = 32
+    """Early stopping callback configuration."""
 
     def __str__(self) -> str:
         """Pretty string reprensenting the configuration.

careamics/dataset/dataset_utils/__init__.py

@@ -6,9 +6,6 @@ __all__ = [
     "get_files_size",
     "list_files",
     "validate_source_target_files",
-    "read_tiff",
-    "get_read_func",
-    "read_zarr",
     "iterate_over_files",
     "WelfordStatistics",
 ]
@@ -19,7 +16,4 @@ from .dataset_utils import (
 )
 from .file_utils import get_files_size, list_files, validate_source_target_files
 from .iterate_over_files import iterate_over_files
-from .read_tiff import read_tiff
-from .read_utils import get_read_func
-from .read_zarr import read_zarr
 from .running_stats import WelfordStatistics, compute_normalization_stats

careamics/dataset/dataset_utils/file_utils.py

@@ -75,7 +75,7 @@ def list_files(
         raise FileNotFoundError(f"Data path {data_path} does not exist.")
 
     # get extension compatible with fnmatch and rglob search
-    extension = SupportedData.get_extension(data_type)
+    extension = SupportedData.get_extension_pattern(data_type)
 
     if data_type == SupportedData.CUSTOM and extension_filter != "":
         extension = extension_filter

careamics/dataset/dataset_utils/iterate_over_files.py

@@ -9,10 +9,10 @@ from numpy.typing import NDArray
 from torch.utils.data import get_worker_info
 
 from careamics.config import DataConfig, InferenceConfig
+from careamics.file_io.read import read_tiff
 from careamics.utils.logging import get_logger
 
 from .dataset_utils import reshape_array
-from .read_tiff import read_tiff
 
 logger = get_logger(__name__)
 

careamics/dataset/in_memory_dataset.py

@@ -9,14 +9,15 @@ from typing import Any, Callable, Optional, Union
 import numpy as np
 from torch.utils.data import Dataset
 
+from careamics.file_io.read import read_tiff
 from careamics.transforms import Compose
 
 from ..config import DataConfig
 from ..config.transformations import NormalizeModel
 from ..utils.logging import get_logger
-from .dataset_utils import read_tiff
 from .patching.patching import (
     PatchedOutput,
+    Stats,
     prepare_patches_supervised,
     prepare_patches_supervised_array,
     prepare_patches_unsupervised,
@@ -77,47 +78,50 @@ class InMemoryDataset(Dataset):
         # read function
         self.read_source_func = read_source_func
 
-        # Generate patches
+        # generate patches
         supervised = self.input_targets is not None
         patches_data = self._prepare_patches(supervised)
 
-        # Unpack the dataclass
+        # unpack the dataclass
         self.data = patches_data.patches
         self.data_targets = patches_data.targets
 
+        # set image statistics
         if self.data_config.image_means is None:
-            self.image_means = patches_data.image_stats.means
-            self.image_stds = patches_data.image_stats.stds
+            self.image_stats = patches_data.image_stats
             logger.info(
-                f"Computed dataset mean: {self.image_means}, std: {self.image_stds}"
+                f"Computed dataset mean: {self.image_stats.means}, "
+                f"std: {self.image_stats.stds}"
             )
         else:
-            self.image_means = self.data_config.image_means
-            self.image_stds = self.data_config.image_stds
+            self.image_stats = Stats(
+                self.data_config.image_means, self.data_config.image_stds
+            )
 
+        # set target statistics
         if self.data_config.target_means is None:
-            self.target_means = patches_data.target_stats.means
-            self.target_stds = patches_data.target_stats.stds
+            self.target_stats = patches_data.target_stats
         else:
-            self.target_means = self.data_config.target_means
-            self.target_stds = self.data_config.target_stds
+            self.target_stats = Stats(
+                self.data_config.target_means, self.data_config.target_stds
+            )
 
         # update mean and std in configuration
         # the object is mutable and should then be recorded in the CAREamist obj
-        self.data_config.set_mean_and_std(
-            image_means=self.image_means,
-            image_stds=self.image_stds,
-            target_means=self.target_means,
-            target_stds=self.target_stds,
+        self.data_config.set_means_and_stds(
+            image_means=self.image_stats.means,
+            image_stds=self.image_stats.stds,
+            target_means=self.target_stats.means,
+            target_stds=self.target_stats.stds,
         )
         # get transforms
         self.patch_transform = Compose(
             transform_list=[
                 NormalizeModel(
-                    image_means=self.image_means,
-                    image_stds=self.image_stds,
-                    target_means=self.target_means,
-                    target_stds=self.target_stds,
+                    image_means=self.image_stats.means,
+                    image_stds=self.image_stats.stds,
+                    target_means=self.target_stats.means,
+                    target_stds=self.target_stats.stds,
                 )
             ]
             + self.data_config.transforms,
@@ -223,6 +227,18 @@ class InMemoryDataset(Dataset):
                 "and no N2V manipulation (no N2V training)."
             )
 
+    def get_data_statistics(self) -> tuple[list[float], list[float]]:
+        """Return training data statistics.
+
+        This does not return the target data statistics, only those of the input.
+
+        Returns
+        -------
+        tuple of list of floats
+            Means and standard deviations across channels of the training data.
+        """
+        return self.image_stats.get_statistics()
+
     def split_dataset(
         self,
         percentage: float = 0.1,

careamics/dataset/iterable_dataset.py

@@ -12,15 +12,13 @@ from torch.utils.data import IterableDataset
 
 from careamics.config import DataConfig
 from careamics.config.transformations import NormalizeModel
+from careamics.file_io.read import read_tiff
 from careamics.transforms import Compose
 
 from ..utils.logging import get_logger
-from .dataset_utils import (
-    iterate_over_files,
-    read_tiff,
-)
+from .dataset_utils import iterate_over_files
 from .dataset_utils.running_stats import WelfordStatistics
-from .patching.patching import Stats, StatsOutput
+from .patching.patching import Stats
 from .patching.random_patching import extract_patches_random
 
 logger = get_logger(__name__)
@@ -78,31 +76,31 @@ class PathIterableDataset(IterableDataset):
         # only checking the image_mean because the DataConfig class ensures that
         # if image_mean is provided, image_std is also provided
         if not self.data_config.image_means:
-            self.data_stats = self._calculate_mean_and_std()
+            self.image_stats, self.target_stats = self._calculate_mean_and_std()
             logger.info(
-                f"Computed dataset mean: {self.data_stats.image_stats.means},"
-                f"std: {self.data_stats.image_stats.stds}"
+                f"Computed dataset mean: {self.image_stats.means},"
+                f"std: {self.image_stats.stds}"
             )
 
             # update the mean in the config
-            self.data_config.set_mean_and_std(
-                image_means=self.data_stats.image_stats.means,
-                image_stds=self.data_stats.image_stats.stds,
+            self.data_config.set_means_and_stds(
+                image_means=self.image_stats.means,
+                image_stds=self.image_stats.stds,
                 target_means=(
-                    list(self.data_stats.target_stats.means)
-                    if self.data_stats.target_stats.means is not None
+                    list(self.target_stats.means)
+                    if self.target_stats.means is not None
                     else None
                 ),
                 target_stds=(
-                    list(self.data_stats.target_stats.stds)
-                    if self.data_stats.target_stats.stds is not None
+                    list(self.target_stats.stds)
+                    if self.target_stats.stds is not None
                     else None
                 ),
             )
 
         else:
             # if mean and std are provided in the config, use them
-            self.data_stats = StatsOutput(
+            self.image_stats, self.target_stats = (
                 Stats(self.data_config.image_means, self.data_config.image_stds),
                 Stats(self.data_config.target_means, self.data_config.target_stds),
             )
@@ -111,23 +109,23 @@ class PathIterableDataset(IterableDataset):
         self.patch_transform = Compose(
             transform_list=[
                 NormalizeModel(
-                    image_means=self.data_stats.image_stats.means,
-                    image_stds=self.data_stats.image_stats.stds,
-                    target_means=self.data_stats.target_stats.means,
-                    target_stds=self.data_stats.target_stats.stds,
+                    image_means=self.image_stats.means,
+                    image_stds=self.image_stats.stds,
+                    target_means=self.target_stats.means,
+                    target_stds=self.target_stats.stds,
                 )
             ]
             + data_config.transforms
         )
 
-    def _calculate_mean_and_std(self) -> StatsOutput:
+    def _calculate_mean_and_std(self) -> tuple[Stats, Stats]:
         """
         Calculate mean and std of the dataset.
 
         Returns
         -------
-        PatchedOutput
-            Data class containing the image statistics.
+        tuple of Stats and optional Stats
+            Data classes containing the image and target statistics.
         """
         num_samples = 0
         image_stats = WelfordStatistics()
@@ -155,15 +153,12 @@ class PathIterableDataset(IterableDataset):
         if target is not None:
             target_means, target_stds = target_stats.finalize()
 
-        logger.info(f"Calculated mean and std for {num_samples} images")
-        logger.info(f"Mean: {image_means}, std: {image_stds}")
-        return StatsOutput(
-            Stats(image_means, image_stds),
-            Stats(
-                np.array(target_means) if target is not None else None,
-                np.array(target_stds) if target is not None else None,
-            ),
-        )
+            return (
+                Stats(image_means, image_stds),
+                Stats(np.array(target_means), np.array(target_stds)),
+            )
+        else:
+            return Stats(image_means, image_stds), Stats(None, None)
 
     def __iter__(
         self,
@@ -177,8 +172,7 @@ class PathIterableDataset(IterableDataset):
             Single patch.
         """
         assert (
-            self.data_stats.image_stats.means is not None
-            and self.data_stats.image_stats.stds is not None
+            self.image_stats.means is not None and self.image_stats.stds is not None
         ), "Mean and std must be provided"
 
         # iterate over files
@@ -201,6 +195,16 @@ class PathIterableDataset(IterableDataset):
                     target=patch_data[1],
                 )
 
+    def get_data_statistics(self) -> tuple[list[float], list[float]]:
+        """Return training data statistics.
+
+        Returns
+        -------
+        tuple of list of floats
+            Means and standard deviations across channels of the training data.
+        """
+        return self.image_stats.get_statistics()
+
     def get_number_of_files(self) -> int:
         """
         Return the number of files in the dataset.

careamics/dataset/iterable_pred_dataset.py

@@ -8,11 +8,12 @@ from typing import Any, Callable, Generator
 from numpy.typing import NDArray
 from torch.utils.data import IterableDataset
 
+from careamics.file_io.read import read_tiff
 from careamics.transforms import Compose
 
 from ..config import InferenceConfig
 from ..config.transformations import NormalizeModel
-from .dataset_utils import iterate_over_files, read_tiff
+from .dataset_utils import iterate_over_files
 
 
 class IterablePredDataset(IterableDataset):

careamics/dataset/iterable_tiled_pred_dataset.py

@@ -8,12 +8,13 @@ from typing import Any, Callable, Generator
 from numpy.typing import NDArray
 from torch.utils.data import IterableDataset
 
+from careamics.file_io.read import read_tiff
 from careamics.transforms import Compose
 
 from ..config import InferenceConfig
 from ..config.tile_information import TileInformation
 from ..config.transformations import NormalizeModel
-from .dataset_utils import iterate_over_files, read_tiff
+from .dataset_utils import iterate_over_files
 from .tiling import extract_tiles