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

careamics/__init__.py

@@ -7,20 +7,7 @@ try:
 except PackageNotFoundError:
     __version__ = "uninstalled"
 
-__all__ = [
-    "CAREamist",
-    "CAREamicsModuleWrapper",
-    "CAREamicsPredictData",
-    "CAREamicsTrainData",
-    "Configuration",
-    "load_configuration",
-    "save_configuration",
-    "TrainingDataWrapper",
-    "PredictDataWrapper",
-]
+__all__ = ["CAREamist", "Configuration", "load_configuration", "save_configuration"]
 
 from .careamist import CAREamist
 from .config import Configuration, load_configuration, save_configuration
-from .lightning_datamodule import CAREamicsTrainData, TrainingDataWrapper
-from .lightning_module import CAREamicsModuleWrapper
-from .lightning_prediction_datamodule import CAREamicsPredictData, PredictDataWrapper

careamics/careamist.py

@@ -13,22 +13,29 @@ from pytorch_lightning.callbacks import (
 )
 from pytorch_lightning.loggers import TensorBoardLogger, WandbLogger
 
-from careamics.callbacks import ProgressBarCallback
 from careamics.config import (
     Configuration,
     load_configuration,
 )
-from careamics.config.support import SupportedAlgorithm, SupportedData, SupportedLogger
+from careamics.config.support import (
+    SupportedAlgorithm,
+    SupportedArchitecture,
+    SupportedData,
+    SupportedLogger,
+)
 from careamics.dataset.dataset_utils import reshape_array
-from careamics.lightning_datamodule import CAREamicsTrainData
-from careamics.lightning_module import CAREamicsModule
+from careamics.lightning import (
+    CAREamicsModule,
+    HyperParametersCallback,
+    PredictDataModule,
+    ProgressBarCallback,
+    TrainDataModule,
+    create_predict_datamodule,
+)
 from careamics.model_io import export_to_bmz, load_pretrained
-from careamics.prediction_utils import convert_outputs, create_pred_datamodule
+from careamics.prediction_utils import convert_outputs
 from careamics.utils import check_path_exists, get_logger
 
-from .callbacks import HyperParametersCallback
-from .lightning_prediction_datamodule import CAREamicsPredictData
-
 logger = get_logger(__name__)
 
 LOGGER_TYPES = Optional[Union[TensorBoardLogger, WandbLogger]]
@@ -61,9 +68,9 @@ class CAREamist:
         Experiment logger, "wandb" or "tensorboard".
     work_dir : pathlib.Path
         Working directory.
-    train_datamodule : CAREamicsTrainData
+    train_datamodule : TrainDataModule
         Training datamodule.
-    pred_datamodule : CAREamicsPredictData
+    pred_datamodule : PredictDataModule
         Prediction datamodule.
     """
 
@@ -193,8 +200,8 @@ class CAREamist:
         )
 
         # place holder for the datamodules
-        self.train_datamodule: Optional[CAREamicsTrainData] = None
-        self.pred_datamodule: Optional[CAREamicsPredictData] = None
+        self.train_datamodule: Optional[TrainDataModule] = None
+        self.pred_datamodule: Optional[PredictDataModule] = None
 
     def _define_callbacks(self, callbacks: Optional[list[Callback]] = None) -> None:
         """
@@ -246,7 +253,7 @@ class CAREamist:
     def train(
         self,
         *,
-        datamodule: Optional[CAREamicsTrainData] = None,
+        datamodule: Optional[TrainDataModule] = None,
         train_source: Optional[Union[Path, str, NDArray]] = None,
         val_source: Optional[Union[Path, str, NDArray]] = None,
         train_target: Optional[Union[Path, str, NDArray]] = None,
@@ -273,7 +280,7 @@ class CAREamist:
 
         Parameters
         ----------
-        datamodule : CAREamicsTrainData, optional
+        datamodule : TrainDataModule, optional
             Datamodule to train on, by default None.
         train_source : pathlib.Path or str or NDArray, optional
             Train source, if no datamodule is provided, by default None.
@@ -375,17 +382,17 @@ class CAREamist:
 
             else:
                 raise ValueError(
-                    f"Invalid input, expected a str, Path, array or CAREamicsTrainData "
+                    f"Invalid input, expected a str, Path, array or TrainDataModule "
                     f"instance (got {type(train_source)})."
                 )
 
-    def _train_on_datamodule(self, datamodule: CAREamicsTrainData) -> None:
+    def _train_on_datamodule(self, datamodule: TrainDataModule) -> None:
         """
         Train the model on the provided datamodule.
 
         Parameters
         ----------
-        datamodule : CAREamicsTrainData
+        datamodule : TrainDataModule
             Datamodule to train on.
         """
         # record datamodule
@@ -421,7 +428,7 @@ class CAREamist:
             Minimum number of patches to use for validation, by default 5.
         """
         # create datamodule
-        datamodule = CAREamicsTrainData(
+        datamodule = TrainDataModule(
             data_config=self.cfg.data_config,
             train_data=train_data,
             val_data=val_data,
@@ -477,7 +484,7 @@ class CAREamist:
             path_to_val_target = check_path_exists(path_to_val_target)
 
         # create datamodule
-        datamodule = CAREamicsTrainData(
+        datamodule = TrainDataModule(
             data_config=self.cfg.data_config,
             train_data=path_to_train_data,
             val_data=path_to_val_data,
@@ -493,10 +500,7 @@ class CAREamist:
 
     @overload
     def predict(  # numpydoc ignore=GL08
-        self,
-        source: CAREamicsPredictData,
-        *,
-        checkpoint: Optional[Literal["best", "last"]] = None,
+        self, source: PredictDataModule
     ) -> Union[list[NDArray], NDArray]: ...
 
     @overload
@@ -513,7 +517,6 @@ class CAREamist:
         dataloader_params: Optional[dict] = None,
         read_source_func: Optional[Callable] = None,
         extension_filter: str = "",
-        checkpoint: Optional[Literal["best", "last"]] = None,
     ) -> Union[list[NDArray], NDArray]: ...
 
     @overload
@@ -528,12 +531,11 @@ class CAREamist:
         data_type: Optional[Literal["array"]] = None,
         tta_transforms: bool = True,
         dataloader_params: Optional[dict] = None,
-        checkpoint: Optional[Literal["best", "last"]] = None,
     ) -> Union[list[NDArray], NDArray]: ...
 
     def predict(
         self,
-        source: Union[CAREamicsPredictData, Path, str, NDArray],
+        source: Union[PredictDataModule, Path, str, NDArray],
         *,
         batch_size: Optional[int] = None,
         tile_size: Optional[tuple[int, ...]] = None,
@@ -544,7 +546,6 @@ class CAREamist:
         dataloader_params: Optional[dict] = None,
         read_source_func: Optional[Callable] = None,
         extension_filter: str = "",
-        checkpoint: Optional[Literal["best", "last"]] = None,
         **kwargs: Any,
     ) -> Union[list[NDArray], NDArray]:
         """
@@ -590,8 +591,6 @@ class CAREamist:
             Function to read the source data.
         extension_filter : str, default=""
             Filter for the file extension.
-        checkpoint : {"best", "last"}, optional
-            Checkpoint to use for prediction.
         **kwargs : Any
             Unused.
 
@@ -599,31 +598,64 @@ class CAREamist:
         -------
         list of NDArray or NDArray
             Predictions made by the model.
-        """
-        # Reuse batch size if not provided explicitly
-        if batch_size is None:
-            batch_size = (
-                self.train_datamodule.batch_size
-                if self.train_datamodule
-                else self.cfg.data_config.batch_size
-            )
 
-        self.pred_datamodule = create_pred_datamodule(
-            source=source,
-            config=self.cfg,
-            batch_size=batch_size,
+        Raises
+        ------
+        ValueError
+            If mean and std are not provided in the configuration.
+        ValueError
+            If tile size is not divisible by 2**depth for UNet models.
+        ValueError
+            If tile overlap is not specified.
+        """
+        if (
+            self.cfg.data_config.image_means is None
+            or self.cfg.data_config.image_stds is None
+        ):
+            raise ValueError("Mean and std must be provided in the configuration.")
+
+        # tile size for UNets
+        if tile_size is not None:
+            model = self.cfg.algorithm_config.model
+
+            if model.architecture == SupportedArchitecture.UNET.value:
+                # tile size must be equal to k*2^n, where n is the number of pooling
+                # layers (equal to the depth) and k is an integer
+                depth = model.depth
+                tile_increment = 2**depth
+
+                for i, t in enumerate(tile_size):
+                    if t % tile_increment != 0:
+                        raise ValueError(
+                            f"Tile size must be divisible by {tile_increment} along "
+                            f"all axes (got {t} for axis {i}). If your image size is "
+                            f"smaller along one axis (e.g. Z), consider padding the "
+                            f"image."
+                        )
+
+            # tile overlaps must be specified
+            if tile_overlap is None:
+                raise ValueError("Tile overlap must be specified.")
+
+        # create the prediction
+        self.pred_datamodule = create_predict_datamodule(
+            pred_data=source,
+            data_type=data_type or self.cfg.data_config.data_type,
+            axes=axes or self.cfg.data_config.axes,
+            image_means=self.cfg.data_config.image_means,
+            image_stds=self.cfg.data_config.image_stds,
             tile_size=tile_size,
             tile_overlap=tile_overlap,
-            axes=axes,
-            data_type=data_type,
+            batch_size=batch_size or self.cfg.data_config.batch_size,
             tta_transforms=tta_transforms,
-            dataloader_params=dataloader_params,
             read_source_func=read_source_func,
             extension_filter=extension_filter,
+            dataloader_params=dataloader_params,
         )
 
+        # predict
         predictions = self.trainer.predict(
-            model=self.model, datamodule=self.pred_datamodule, ckpt_path=checkpoint
+            model=self.model, datamodule=self.pred_datamodule
         )
         return convert_outputs(predictions, self.pred_datamodule.tiled)
 
@@ -659,27 +691,16 @@ class CAREamist:
         data_description : str, optional
             Description of the data, by default None.
         """
-        input_patch = reshape_array(input_array, self.cfg.data_config.axes)
-
-        # axes need to be reformated for the export because reshaping was done in the
-        # datamodule
-        if "Z" in self.cfg.data_config.axes:
-            axes = "SCZYX"
-        else:
-            axes = "SCYX"
+        # TODO: add in docs that it is expected that input_array dimensions match
+        # those in data_config
 
-        # predict output, remove extra dimensions for the purpose of the prediction
         output_patch = self.predict(
-            input_patch,
+            input_array,
             data_type=SupportedData.ARRAY.value,
-            axes=axes,
             tta_transforms=False,
         )
-
-        if isinstance(output_patch, list):
-            output = np.concatenate(output_patch, axis=0)
-        else:
-            output = output_patch
+        output = np.concatenate(output_patch, axis=0)
+        input_array = reshape_array(input_array, self.cfg.data_config.axes)
 
         export_to_bmz(
             model=self.model,
@@ -688,7 +709,7 @@ class CAREamist:
             name=name,
             general_description=general_description,
             authors=authors,
-            input_array=input_patch,
+            input_array=input_array,
             output_array=output,
             channel_names=channel_names,
             data_description=data_description,

careamics/config/__init__.py

@@ -14,9 +14,7 @@ __all__ = [
     "create_care_configuration",
     "register_model",
     "CustomModel",
-    "create_inference_configuration",
     "clear_custom_models",
-    "ConfigurationInformation",
 ]
 
 from .algorithm_model import AlgorithmConfig
@@ -24,7 +22,6 @@ from .architectures import CustomModel, clear_custom_models, register_model
 from .callback_model import CheckpointModel
 from .configuration_factory import (
     create_care_configuration,
-    create_inference_configuration,
     create_n2n_configuration,
     create_n2v_configuration,
 )

careamics/config/algorithm_model.py

@@ -93,12 +93,20 @@ class AlgorithmConfig(BaseModel):
 
     # Mandatory fields
     algorithm: Literal["n2v", "care", "n2n", "custom"]  # defined in SupportedAlgorithm
+    """Name of the algorithm, as defined in SupportedAlgorithm."""
+
     loss: Literal["n2v", "mae", "mse"]
+    """Loss function to use, as defined in SupportedLoss."""
+
     model: Union[UNetModel, VAEModel, CustomModel] = Field(discriminator="architecture")
+    """Model architecture to use, defined in SupportedArchitecture."""
 
     # Optional fields
     optimizer: OptimizerModel = OptimizerModel()
+    """Optimizer to use, defined in SupportedOptimizer."""
+
     lr_scheduler: LrSchedulerModel = LrSchedulerModel()
+    """Learning rate scheduler to use, defined in SupportedScheduler."""
 
     @model_validator(mode="after")
     def algorithm_cross_validation(self: Self) -> Self:

careamics/config/architectures/architecture_model.py

@@ -13,6 +13,7 @@ class ArchitectureModel(BaseModel):
     """
 
     architecture: str
+    """Name of the architecture."""
 
     def model_dump(self, **kwargs: Any) -> Dict[str, Any]:
         """

careamics/config/architectures/custom_model.py

@@ -72,9 +72,11 @@ class CustomModel(ArchitectureModel):
 
     # discriminator used for choosing the pydantic model in Model
     architecture: Literal["Custom"]
+    """Name of the architecture."""
 
     # name of the custom model
     name: str
+    """Name of the custom model."""
 
     @field_validator("name")
     @classmethod

careamics/config/architectures/unet_model.py

@@ -29,19 +29,38 @@ class UNetModel(ArchitectureModel):
 
     # discriminator used for choosing the pydantic model in Model
     architecture: Literal["UNet"]
+    """Name of the architecture."""
 
     # parameters
     # validate_defaults allow ignoring default values in the dump if they were not set
     conv_dims: Literal[2, 3] = Field(default=2, validate_default=True)
+    """Dimensions (2D or 3D) of the convolutional layers."""
+
     num_classes: int = Field(default=1, ge=1, validate_default=True)
+    """Number of classes or channels in the model output."""
+
     in_channels: int = Field(default=1, ge=1, validate_default=True)
+    """Number of channels in the input to the model."""
+
     depth: int = Field(default=2, ge=1, le=10, validate_default=True)
+    """Number of levels in the UNet."""
+
     num_channels_init: int = Field(default=32, ge=8, le=1024, validate_default=True)
+    """Number of convolutional filters in the first layer of the UNet."""
+
     final_activation: Literal[
         "None", "Sigmoid", "Softmax", "Tanh", "ReLU", "LeakyReLU"
     ] = Field(default="None", validate_default=True)
+    """Final activation function."""
+
     n2v2: bool = Field(default=False, validate_default=True)
+    """Whether to use N2V2 architecture modifications, with blur pool layers and fewer
+    skip connections.
+    """
+
     independent_channels: bool = Field(default=True, validate_default=True)
+    """Whether information is processed independently in each channel, used to train
+    channels independently."""
 
     @field_validator("num_channels_init")
     @classmethod

careamics/config/architectures/vae_model.py

@@ -17,6 +17,7 @@ class VAEModel(ArchitectureModel):
     )
 
     architecture: Literal["VAE"]
+    """Name of the architecture."""
 
     def set_3D(self, is_3D: bool) -> None:
         """

careamics/config/callback_model.py

@@ -13,69 +13,111 @@ from pydantic import (
 
 
 class CheckpointModel(BaseModel):
-    """Checkpoint saving callback Pydantic model."""
+    """Checkpoint saving callback Pydantic model.
+
+    The parameters corresponds to those of
+    `pytorch_lightning.callbacks.ModelCheckpoint`.
+
+    See:
+    https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.ModelCheckpoint.html#modelcheckpoint
+    """
 
     model_config = ConfigDict(
         validate_assignment=True,
     )
 
     monitor: Literal["val_loss"] = Field(default="val_loss", validate_default=True)
+    """Quantity to monitor."""
+
     verbose: bool = Field(default=False, validate_default=True)
+    """Verbosity mode."""
+
     save_weights_only: bool = Field(default=False, validate_default=True)
+    """When `True`, only the model's weights will be saved (model.save_weights)."""
+
+    save_last: Optional[Literal[True, False, "link"]] = Field(
+        default=True, validate_default=True
+    )
+    """When `True`, saves a last.ckpt copy whenever a checkpoint file gets saved."""
+
+    save_top_k: int = Field(default=3, ge=1, le=10, validate_default=True)
+    """If `save_top_k == kz, the best k models according to the quantity monitored
+    will be saved. If `save_top_k == 0`, no models are saved. if `save_top_k == -1`,
+    all models are saved."""
+
     mode: Literal["min", "max"] = Field(default="min", validate_default=True)
+    """One of {min, max}. If `save_top_k != 0`, the decision to overwrite the current
+    save file is made based on either the maximization or the minimization of the
+    monitored quantity. For 'val_acc', this should be 'max', for 'val_loss' this should
+    be 'min', etc.
+    """
+
     auto_insert_metric_name: bool = Field(default=False, validate_default=True)
+    """When `True`, the checkpoints filenames will contain the metric name."""
+
     every_n_train_steps: Optional[int] = Field(
         default=None, ge=1, le=10, validate_default=True
     )
+    """Number of training steps between checkpoints."""
+
     train_time_interval: Optional[timedelta] = Field(
         default=None, validate_default=True
     )
+    """Checkpoints are monitored at the specified time interval."""
+
     every_n_epochs: Optional[int] = Field(
         default=None, ge=1, le=10, validate_default=True
     )
-    save_last: Optional[Literal[True, False, "link"]] = Field(
-        default=True, validate_default=True
-    )
-    save_top_k: int = Field(default=3, ge=1, le=10, validate_default=True)
+    """Number of epochs between checkpoints."""
 
 
 class EarlyStoppingModel(BaseModel):
-    """Early stopping callback Pydantic model."""
+    """Early stopping callback Pydantic model.
+
+    The parameters corresponds to those of
+    `pytorch_lightning.callbacks.ModelCheckpoint`.
+
+    See:
+    https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.EarlyStopping.html#lightning.pytorch.callbacks.EarlyStopping
+    """
 
     model_config = ConfigDict(
         validate_assignment=True,
     )
 
     monitor: Literal["val_loss"] = Field(default="val_loss", validate_default=True)
+    """Quantity to monitor."""
+
+    min_delta: float = Field(default=0.0, ge=0.0, le=1.0, validate_default=True)
+    """Minimum change in the monitored quantity to qualify as an improvement, i.e. an
+    absolute change of less than or equal to min_delta, will count as no improvement."""
+
     patience: int = Field(default=3, ge=1, le=10, validate_default=True)
+    """Number of checks with no improvement after which training will be stopped."""
+
+    verbose: bool = Field(default=False, validate_default=True)
+    """Verbosity mode."""
+
     mode: Literal["min", "max", "auto"] = Field(default="min", validate_default=True)
-    min_delta: float = Field(default=0.0, ge=0.0, le=1.0, validate_default=True)
+    """One of {min, max, auto}."""
+
     check_finite: bool = Field(default=True, validate_default=True)
-    stop_on_nan: bool = Field(default=True, validate_default=True)
-    verbose: bool = Field(default=False, validate_default=True)
-    restore_best_weights: bool = Field(default=True, validate_default=True)
-    auto_lr_find: bool = Field(default=False, validate_default=True)
-    auto_lr_find_patience: int = Field(default=3, ge=1, le=10, validate_default=True)
-    auto_lr_find_mode: Literal["min", "max", "auto"] = Field(
-        default="min", validate_default=True
-    )
-    auto_lr_find_direction: Literal["forward", "backward"] = Field(
-        default="backward", validate_default=True
-    )
-    auto_lr_find_max_lr: float = Field(
-        default=10.0, ge=0.0, le=1e6, validate_default=True
-    )
-    auto_lr_find_min_lr: float = Field(
-        default=1e-8, ge=0.0, le=1e6, validate_default=True
-    )
-    auto_lr_find_num_training: int = Field(
-        default=100, ge=1, le=1e6, validate_default=True
-    )
-    auto_lr_find_divergence_threshold: float = Field(
-        default=5.0, ge=0.0, le=1e6, validate_default=True
-    )
-    auto_lr_find_accumulate_grad_batches: int = Field(
-        default=1, ge=1, le=1e6, validate_default=True
+    """When `True`, stops training when the monitored quantity becomes `NaN` or
+    `inf`."""
+
+    stopping_threshold: Optional[float] = Field(default=None, validate_default=True)
+    """Stop training immediately once the monitored quantity reaches this threshold."""
+
+    divergence_threshold: Optional[float] = Field(default=None, validate_default=True)
+    """Stop training as soon as the monitored quantity becomes worse than this
+    threshold."""
+
+    check_on_train_epoch_end: Optional[bool] = Field(
+        default=False, validate_default=True
     )
-    auto_lr_find_stop_divergence: bool = Field(default=True, validate_default=True)
-    auto_lr_find_step_scale: float = Field(default=0.1, ge=0.0, le=10)
+    """Whether to run early stopping at the end of the training epoch. If this is
+    `False`, then the check runs at the end of the validation."""
+
+    log_rank_zero_only: bool = Field(default=False, validate_default=True)
+    """When set `True`, logs the status of the early stopping callback only for rank 0
+    process."""

careamics/config/configuration_factory.py

@@ -1,12 +1,11 @@
 """Convenience functions to create configurations for training and inference."""
 
-from typing import Any, Dict, List, Literal, Optional, Tuple
+from typing import Any, Dict, List, Literal, Optional
 
 from .algorithm_model import AlgorithmConfig
 from .architectures import UNetModel
 from .configuration_model import Configuration
 from .data_model import DataConfig
-from .inference_model import InferenceConfig
 from .support import (
     SupportedAlgorithm,
     SupportedArchitecture,
@@ -574,80 +573,3 @@ def create_n2v_configuration(
     )
 
     return configuration
-
-
-def create_inference_configuration(
-    configuration: Configuration,
-    tile_size: Optional[Tuple[int, ...]] = None,
-    tile_overlap: Optional[Tuple[int, ...]] = None,
-    data_type: Optional[Literal["array", "tiff", "custom"]] = None,
-    axes: Optional[str] = None,
-    tta_transforms: bool = True,
-    batch_size: Optional[int] = 1,
-) -> InferenceConfig:
-    """
-    Create a configuration for inference with N2V.
-
-    If not provided, `data_type` and `axes` are taken from the training
-    configuration.
-
-    Parameters
-    ----------
-    configuration : Configuration
-        Global configuration.
-    tile_size : Tuple[int, ...], optional
-        Size of the tiles.
-    tile_overlap : Tuple[int, ...], optional
-        Overlap of the tiles.
-    data_type : str, optional
-        Type of the data, by default "tiff".
-    axes : str, optional
-        Axes of the data, by default "YX".
-    tta_transforms : bool, optional
-        Whether to apply test-time augmentations, by default True.
-    batch_size : int, optional
-        Batch size, by default 1.
-
-    Returns
-    -------
-    InferenceConfiguration
-        Configuration used to configure CAREamicsPredictData.
-    """
-    if (
-        configuration.data_config.image_means is None
-        or configuration.data_config.image_stds is None
-    ):
-        raise ValueError("Mean and std must be provided in the configuration.")
-
-    # tile size for UNets
-    if tile_size is not None:
-        model = configuration.algorithm_config.model
-
-        if model.architecture == SupportedArchitecture.UNET.value:
-            # tile size must be equal to k*2^n, where n is the number of pooling layers
-            # (equal to the depth) and k is an integer
-            depth = model.depth
-            tile_increment = 2**depth
-
-            for i, t in enumerate(tile_size):
-                if t % tile_increment != 0:
-                    raise ValueError(
-                        f"Tile size must be divisible by {tile_increment} along all "
-                        f"axes (got {t} for axis {i}). If your image size is smaller "
-                        f"along one axis (e.g. Z), consider padding the image."
-                    )
-
-        # tile overlaps must be specified
-        if tile_overlap is None:
-            raise ValueError("Tile overlap must be specified.")
-
-    return InferenceConfig(
-        data_type=data_type or configuration.data_config.data_type,
-        tile_size=tile_size,
-        tile_overlap=tile_overlap,
-        axes=axes or configuration.data_config.axes,
-        image_means=configuration.data_config.image_means,
-        image_stds=configuration.data_config.image_stds,
-        tta_transforms=tta_transforms,
-        batch_size=batch_size,
-    )