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

careamics/__init__.py

@@ -9,16 +9,18 @@ except PackageNotFoundError:
 
 __all__ = [
     "CAREamist",
-    "CAREamicsModule",
+    "CAREamicsModuleWrapper",
+    "CAREamicsPredictData",
+    "CAREamicsTrainData",
     "Configuration",
     "load_configuration",
     "save_configuration",
-    "CAREamicsTrainDataModule",
-    "CAREamicsPredictDataModule",
+    "TrainingDataWrapper",
+    "PredictDataWrapper",
 ]
 
 from .careamist import CAREamist
 from .config import Configuration, load_configuration, save_configuration
-from .lightning_datamodule import CAREamicsTrainDataModule
-from .lightning_module import CAREamicsModule
-from .lightning_prediction_datamodule import CAREamicsPredictDataModule
+from .lightning_datamodule import CAREamicsTrainData, TrainingDataWrapper
+from .lightning_module import CAREamicsModuleWrapper
+from .lightning_prediction_datamodule import CAREamicsPredictData, PredictDataWrapper

careamics/careamist.py

@@ -20,9 +20,9 @@ from careamics.config import (
 )
 from careamics.config.inference_model import TRANSFORMS_UNION
 from careamics.config.support import SupportedAlgorithm, SupportedData, SupportedLogger
-from careamics.lightning_datamodule import CAREamicsWood
-from careamics.lightning_module import CAREamicsKiln
-from careamics.lightning_prediction_datamodule import CAREamicsClay
+from careamics.lightning_datamodule import CAREamicsTrainData
+from careamics.lightning_module import CAREamicsModule
+from careamics.lightning_prediction_datamodule import CAREamicsPredictData
 from careamics.lightning_prediction_loop import CAREamicsPredictionLoop
 from careamics.model_io import export_to_bmz, load_pretrained
 from careamics.utils import check_path_exists, get_logger
@@ -73,8 +73,7 @@ class CAREamist:
         source: Union[Path, str],
         work_dir: Optional[str] = None,
         experiment_name: str = "CAREamics",
-    ) -> None:
-        ...
+    ) -> None: ...
 
     @overload
     def __init__(  # numpydoc ignore=GL08
@@ -82,8 +81,7 @@ class CAREamist:
         source: Configuration,
         work_dir: Optional[str] = None,
         experiment_name: str = "CAREamics",
-    ) -> None:
-        ...
+    ) -> None: ...
 
     def __init__(
         self,
@@ -140,7 +138,7 @@ class CAREamist:
             self.cfg = source
 
             # instantiate model
-            self.model = CAREamicsKiln(
+            self.model = CAREamicsModule(
                 algorithm_config=self.cfg.algorithm_config,
             )
 
@@ -156,7 +154,7 @@ class CAREamist:
                 self.cfg = load_configuration(source)
 
                 # instantiate model
-                self.model = CAREamicsKiln(
+                self.model = CAREamicsModule(
                     algorithm_config=self.cfg.algorithm_config,
                 )
 
@@ -193,8 +191,8 @@ class CAREamist:
         self.trainer.predict_loop = CAREamicsPredictionLoop(self.trainer)
 
         # place holder for the datamodules
-        self.train_datamodule: Optional[CAREamicsWood] = None
-        self.pred_datamodule: Optional[CAREamicsClay] = None
+        self.train_datamodule: Optional[CAREamicsTrainData] = None
+        self.pred_datamodule: Optional[CAREamicsPredictData] = None
 
     def _define_callbacks(self) -> List[Callback]:
         """
@@ -227,7 +225,7 @@ class CAREamist:
     def train(
         self,
         *,
-        datamodule: Optional[CAREamicsWood] = None,
+        datamodule: Optional[CAREamicsTrainData] = None,
         train_source: Optional[Union[Path, str, np.ndarray]] = None,
         val_source: Optional[Union[Path, str, np.ndarray]] = None,
         train_target: Optional[Union[Path, str, np.ndarray]] = None,
@@ -360,7 +358,7 @@ class CAREamist:
                     f"instance (got {type(train_source)})."
                 )
 
-    def _train_on_datamodule(self, datamodule: CAREamicsWood) -> None:
+    def _train_on_datamodule(self, datamodule: CAREamicsTrainData) -> None:
         """
         Train the model on the provided datamodule.
 
@@ -402,7 +400,7 @@ class CAREamist:
             Minimum number of patches to use for validation, by default 5.
         """
         # create datamodule
-        datamodule = CAREamicsWood(
+        datamodule = CAREamicsTrainData(
             data_config=self.cfg.data_config,
             train_data=train_data,
             val_data=val_data,
@@ -458,7 +456,7 @@ class CAREamist:
             path_to_val_target = check_path_exists(path_to_val_target)
 
         # create datamodule
-        datamodule = CAREamicsWood(
+        datamodule = CAREamicsTrainData(
             data_config=self.cfg.data_config,
             train_data=path_to_train_data,
             val_data=path_to_val_data,
@@ -475,11 +473,10 @@ class CAREamist:
     @overload
     def predict(  # numpydoc ignore=GL08
         self,
-        source: CAREamicsClay,
+        source: CAREamicsPredictData,
         *,
         checkpoint: Optional[Literal["best", "last"]] = None,
-    ) -> Union[list, np.ndarray]:
-        ...
+    ) -> Union[list, np.ndarray]: ...
 
     @overload
     def predict(  # numpydoc ignore=GL08
@@ -497,8 +494,7 @@ class CAREamist:
         read_source_func: Optional[Callable] = None,
         extension_filter: str = "",
         checkpoint: Optional[Literal["best", "last"]] = None,
-    ) -> Union[list, np.ndarray]:
-        ...
+    ) -> Union[list, np.ndarray]: ...
 
     @overload
     def predict(  # numpydoc ignore=GL08
@@ -514,12 +510,11 @@ class CAREamist:
         tta_transforms: bool = True,
         dataloader_params: Optional[Dict] = None,
         checkpoint: Optional[Literal["best", "last"]] = None,
-    ) -> Union[list, np.ndarray]:
-        ...
+    ) -> Union[list, np.ndarray]: ...
 
     def predict(
         self,
-        source: Union[CAREamicsClay, Path, str, np.ndarray],
+        source: Union[CAREamicsPredictData, Path, str, np.ndarray],
         *,
         batch_size: int = 1,
         tile_size: Optional[Tuple[int, ...]] = None,
@@ -548,6 +543,12 @@ class CAREamist:
         Test-time augmentation (TTA) can be switched off using the `tta_transforms`
         parameter.
 
+        Note that if you are using a UNet model and tiling, the tile size must be
+        divisible in every dimension by 2**d, where d is the depth of the model. This
+        avoids artefacts arising from the broken shift invariance induced by the
+        pooling layers of the UNet. If your image has less dimensions, as it may
+        happen in the Z dimension, consider padding your image.
+
         Parameters
         ----------
         source : Union[CAREamicsClay, Path, str, np.ndarray]
@@ -587,7 +588,7 @@ class CAREamist:
         ValueError
             If the input is not a CAREamicsClay instance, a path or a numpy array.
         """
-        if isinstance(source, CAREamicsClay):
+        if isinstance(source, CAREamicsPredictData):
             # record datamodule
             self.pred_datamodule = source
 
@@ -602,7 +603,7 @@ class CAREamist:
                 )
             # create predict config, reuse training config if parameters missing
             prediction_config = create_inference_configuration(
-                training_configuration=self.cfg,
+                configuration=self.cfg,
                 tile_size=tile_size,
                 tile_overlap=tile_overlap,
                 data_type=data_type,
@@ -623,8 +624,8 @@ class CAREamist:
                 source_path = check_path_exists(source)
 
                 # create datamodule
-                datamodule = CAREamicsClay(
-                    prediction_config=prediction_config,
+                datamodule = CAREamicsPredictData(
+                    pred_config=prediction_config,
                     pred_data=source_path,
                     read_source_func=read_source_func,
                     extension_filter=extension_filter,
@@ -640,8 +641,8 @@ class CAREamist:
 
             elif isinstance(source, np.ndarray):
                 # create datamodule
-                datamodule = CAREamicsClay(
-                    prediction_config=prediction_config,
+                datamodule = CAREamicsPredictData(
+                    pred_config=prediction_config,
                     pred_data=source,
                     dataloader_params=dataloader_params,
                 )

careamics/config/__init__.py

@@ -1,16 +1,17 @@
 """Configuration module."""
 
-
 __all__ = [
-    "AlgorithmModel",
-    "DataModel",
+    "AlgorithmConfig",
+    "DataConfig",
     "Configuration",
     "CheckpointModel",
-    "InferenceModel",
+    "InferenceConfig",
     "load_configuration",
     "save_configuration",
-    "TrainingModel",
+    "TrainingConfig",
     "create_n2v_configuration",
+    "create_n2n_configuration",
+    "create_care_configuration",
     "register_model",
     "CustomModel",
     "create_inference_configuration",
@@ -18,11 +19,13 @@ __all__ = [
     "ConfigurationInformation",
 ]
 
-from .algorithm_model import AlgorithmModel
+from .algorithm_model import AlgorithmConfig
 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,
 )
 from .configuration_model import (
@@ -30,6 +33,6 @@ from .configuration_model import (
     load_configuration,
     save_configuration,
 )
-from .data_model import DataModel
-from .inference_model import InferenceModel
-from .training_model import TrainingModel
+from .data_model import DataConfig
+from .inference_model import InferenceConfig
+from .training_model import TrainingConfig

careamics/config/algorithm_model.py

@@ -10,7 +10,7 @@ from .architectures import CustomModel, UNetModel, VAEModel
 from .optimizer_models import LrSchedulerModel, OptimizerModel
 
 
-class AlgorithmModel(BaseModel):
+class AlgorithmConfig(BaseModel):
     """Algorithm configuration.
 
     This Pydantic model validates the parameters governing the components of the
@@ -45,7 +45,7 @@ class AlgorithmModel(BaseModel):
     Examples
     --------
     Minimum example:
-    >>> from careamics.config import AlgorithmModel
+    >>> from careamics.config import AlgorithmConfig
     >>> config_dict = {
     ...     "algorithm": "n2v",
     ...     "loss": "n2v",
@@ -53,11 +53,11 @@ class AlgorithmModel(BaseModel):
     ...         "architecture": "UNet",
     ...     }
     ... }
-    >>> config = AlgorithmModel(**config_dict)
+    >>> config = AlgorithmConfig(**config_dict)
 
     Using a custom model:
     >>> from torch import nn, ones
-    >>> from careamics.config import AlgorithmModel, register_model
+    >>> from careamics.config import AlgorithmConfig, register_model
     ...
     >>> @register_model(name="linear_model")
     ... class LinearModel(nn.Module):
@@ -80,7 +80,7 @@ class AlgorithmModel(BaseModel):
     ...         "out_features": 5,
     ...     }
     ... }
-    >>> config = AlgorithmModel(**config_dict)
+    >>> config = AlgorithmConfig(**config_dict)
     """
 
     # Pydantic class configuration

careamics/config/architectures/unet_model.py

@@ -39,6 +39,7 @@ class UNetModel(ArchitectureModel):
         "None", "Sigmoid", "Softmax", "Tanh", "ReLU", "LeakyReLU"
     ] = Field(default="None", validate_default=True)
     n2v2: bool = Field(default=False, validate_default=True)
+    independent_channels: bool = Field(default=True, validate_default=True)
 
     @field_validator("num_channels_init")
     @classmethod

careamics/config/callback_model.py

@@ -1,4 +1,5 @@
 """Checkpoint saving configuration."""
+
 from __future__ import annotations
 
 from datetime import timedelta

careamics/config/configuration_example.py

@@ -0,0 +1,87 @@
+from .algorithm_model import AlgorithmConfig
+from .architectures import UNetModel
+from .configuration_model import Configuration
+from .data_model import DataConfig
+from .optimizer_models import LrSchedulerModel, OptimizerModel
+from .support import (
+    SupportedActivation,
+    SupportedAlgorithm,
+    SupportedArchitecture,
+    SupportedData,
+    SupportedLogger,
+    SupportedLoss,
+    SupportedOptimizer,
+    SupportedPixelManipulation,
+    SupportedScheduler,
+    SupportedTransform,
+)
+from .training_model import TrainingConfig
+
+
+def full_configuration_example() -> Configuration:
+    """Returns a dictionnary representing a full configuration example.
+
+    Returns
+    -------
+    Configuration
+        Full configuration example.
+    """
+    experiment_name = "Full example"
+    algorithm_model = AlgorithmConfig(
+        algorithm=SupportedAlgorithm.N2V.value,
+        loss=SupportedLoss.N2V.value,
+        model=UNetModel(
+            architecture=SupportedArchitecture.UNET.value,
+            in_channels=1,
+            num_classes=1,
+            depth=2,
+            num_channels_init=32,
+            final_activation=SupportedActivation.NONE.value,
+            n2v2=True,
+        ),
+        optimizer=OptimizerModel(
+            name=SupportedOptimizer.ADAM.value, parameters={"lr": 0.0001}
+        ),
+        lr_scheduler=LrSchedulerModel(
+            name=SupportedScheduler.REDUCE_LR_ON_PLATEAU.value,
+        ),
+    )
+    data_model = DataConfig(
+        data_type=SupportedData.ARRAY.value,
+        patch_size=(256, 256),
+        batch_size=8,
+        axes="YX",
+        transforms=[
+            {
+                "name": SupportedTransform.NORMALIZE.value,
+            },
+            {
+                "name": SupportedTransform.NDFLIP.value,
+            },
+            {
+                "name": SupportedTransform.XY_RANDOM_ROTATE90.value,
+            },
+            {
+                "name": SupportedTransform.N2V_MANIPULATE.value,
+                "roi_size": 11,
+                "masked_pixel_percentage": 0.2,
+                "strategy": SupportedPixelManipulation.MEDIAN.value,
+            },
+        ],
+        mean=0.485,
+        std=0.229,
+        dataloader_params={
+            "num_workers": 4,
+        },
+    )
+    training_model = TrainingConfig(
+        num_epochs=30,
+        logger=SupportedLogger.WANDB.value,
+    )
+
+    return Configuration(
+        experiment_name=experiment_name,
+        algorithm_config=algorithm_model,
+        data_config=data_model,
+        training_config=training_model,
+    )