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

careamics/file_io/read/get_func.py

@@ -1,7 +1,7 @@
 """Module to get read functions."""
 
 from pathlib import Path
-from typing import Callable, Dict, Protocol, Union
+from typing import Callable, Protocol, Union
 
 from numpy.typing import NDArray
 
@@ -30,7 +30,7 @@ class ReadFunc(Protocol):
         """
 
 
-READ_FUNCS: Dict[SupportedData, ReadFunc] = {
+READ_FUNCS: dict[SupportedData, ReadFunc] = {
     SupportedData.TIFF: read_tiff,
 }
 

careamics/file_io/write/__init__.py

@@ -1,10 +1,10 @@
 """Functions relating to writing image files of different formats."""
 
 __all__ = [
+    "SupportedWriteType",
+    "WriteFunc",
     "get_write_func",
     "write_tiff",
-    "WriteFunc",
-    "SupportedWriteType",
 ]
 
 from .get_func import (

careamics/lightning/__init__.py

@@ -2,14 +2,14 @@
 
 __all__ = [
     "FCNModule",
+    "HyperParametersCallback",
+    "PredictDataModule",
+    "ProgressBarCallback",
+    "TrainDataModule",
     "VAEModule",
     "create_careamics_module",
-    "TrainDataModule",
-    "create_train_datamodule",
-    "PredictDataModule",
     "create_predict_datamodule",
-    "HyperParametersCallback",
-    "ProgressBarCallback",
+    "create_train_datamodule",
 ]
 
 from .callbacks import HyperParametersCallback, ProgressBarCallback

careamics/lightning/callbacks/__init__.py

@@ -2,8 +2,8 @@
 
 __all__ = [
     "HyperParametersCallback",
-    "ProgressBarCallback",
     "PredictionWriterCallback",
+    "ProgressBarCallback",
 ]
 
 from .hyperparameters_callback import HyperParametersCallback

careamics/lightning/callbacks/prediction_writer_callback/__init__.py

@@ -1,12 +1,12 @@
 """A package for the `PredictionWriterCallback` class and utilities."""
 
 __all__ = [
+    "CacheTiles",
     "PredictionWriterCallback",
-    "create_write_strategy",
-    "WriteStrategy",
     "WriteImage",
-    "CacheTiles",
+    "WriteStrategy",
     "WriteTilesZarr",
+    "create_write_strategy",
     "select_write_extension",
     "select_write_func",
 ]

careamics/lightning/callbacks/prediction_writer_callback/prediction_writer_callback.py

@@ -2,8 +2,9 @@
 
 from __future__ import annotations
 
+from collections.abc import Sequence
 from pathlib import Path
-from typing import Any, Optional, Sequence, Union
+from typing import Any, Optional, Union
 
 from pytorch_lightning import LightningModule, Trainer
 from pytorch_lightning.callbacks import BasePredictionWriter

careamics/lightning/callbacks/prediction_writer_callback/write_strategy.py

@@ -1,7 +1,8 @@
 """Module containing different strategies for writing predictions."""
 
+from collections.abc import Sequence
 from pathlib import Path
-from typing import Any, Optional, Protocol, Sequence, Union
+from typing import Any, Optional, Protocol, Union
 
 import numpy as np
 from numpy.typing import NDArray

careamics/lightning/callbacks/progress_bar_callback.py

@@ -1,7 +1,7 @@
 """Progressbar callback."""
 
 import sys
-from typing import Dict, Union
+from typing import Union
 
 from pytorch_lightning import LightningModule, Trainer
 from pytorch_lightning.callbacks import TQDMProgressBar
@@ -71,7 +71,7 @@ class ProgressBarCallback(TQDMProgressBar):
 
     def get_metrics(
         self, trainer: Trainer, pl_module: LightningModule
-    ) -> Dict[str, Union[int, str, float, Dict[str, float]]]:
+    ) -> dict[str, Union[int, str, float, dict[str, float]]]:
         """Override this to customize the metrics displayed in the progress bar.
 
         Parameters

careamics/lightning/lightning_module.py

@@ -6,7 +6,7 @@ import numpy as np
 import pytorch_lightning as L
 from torch import Tensor, nn
 
-from careamics.config import FCNAlgorithmConfig, VAEAlgorithmConfig
+from careamics.config import UNetBasedAlgorithm, VAEBasedAlgorithm
 from careamics.config.support import (
     SupportedAlgorithm,
     SupportedArchitecture,
@@ -14,8 +14,8 @@ from careamics.config.support import (
     SupportedOptimizer,
     SupportedScheduler,
 )
+from careamics.config.tile_information import TileInformation
 from careamics.losses import loss_factory
-from careamics.losses.loss_factory import LVAELossParameters
 from careamics.models.lvae.likelihoods import (
     GaussianLikelihood,
     NoiseModelLikelihood,
@@ -34,6 +34,7 @@ from careamics.utils.torch_utils import get_optimizer, get_scheduler
 NoiseModel = Union[GaussianMixtureNoiseModel, MultiChannelNoiseModel]
 
 
+# TODO rename to UNetModule
 class FCNModule(L.LightningModule):
     """
     CAREamics Lightning module.
@@ -60,7 +61,7 @@ class FCNModule(L.LightningModule):
         Learning rate scheduler name.
     """
 
-    def __init__(self, algorithm_config: Union[FCNAlgorithmConfig, dict]) -> None:
+    def __init__(self, algorithm_config: Union[UNetBasedAlgorithm, dict]) -> None:
         """Lightning module for CAREamics.
 
         This class encapsulates the a PyTorch model along with the training, validation,
@@ -74,7 +75,9 @@ class FCNModule(L.LightningModule):
         super().__init__()
         # if loading from a checkpoint, AlgorithmModel needs to be instantiated
         if isinstance(algorithm_config, dict):
-            algorithm_config = FCNAlgorithmConfig(**algorithm_config)
+            algorithm_config = UNetBasedAlgorithm(
+                **algorithm_config
+            )  # TODO this needs to be updated using the algorithm-specific class
 
         # create model and loss function
         self.model: nn.Module = model_factory(algorithm_config.model)
@@ -164,7 +167,17 @@ class FCNModule(L.LightningModule):
         Any
             Model output.
         """
-        if self._trainer.datamodule.tiled:
+        # TODO refactor when redoing datasets
+        # hacky way to determine if it is PredictDataModule, otherwise there is a
+        # circular import to solve with isinstance
+        from_prediction = hasattr(self._trainer.datamodule, "tiled")
+        is_tiled = (
+            len(batch) > 1
+            and isinstance(batch[1], list)
+            and isinstance(batch[1][0], TileInformation)
+        )
+
+        if is_tiled:
             x, *aux = batch
         else:
             x = batch
@@ -172,7 +185,10 @@ class FCNModule(L.LightningModule):
 
         # apply test-time augmentation if available
         # TODO: probably wont work with batch size > 1
-        if self._trainer.datamodule.prediction_config.tta_transforms:
+        if (
+            from_prediction
+            and self._trainer.datamodule.prediction_config.tta_transforms
+        ):
             tta = ImageRestorationTTA()
             augmented_batch = tta.forward(x)  # list of augmented tensors
             augmented_output = []
@@ -184,9 +200,18 @@ class FCNModule(L.LightningModule):
             output = self.model(x)
 
         # Denormalize the output
+        # TODO incompatible API between predict and train datasets
         denorm = Denormalize(
-            image_means=self._trainer.datamodule.predict_dataset.image_means,
-            image_stds=self._trainer.datamodule.predict_dataset.image_stds,
+            image_means=(
+                self._trainer.datamodule.predict_dataset.image_means
+                if from_prediction
+                else self._trainer.datamodule.train_dataset.image_stats.means
+            ),
+            image_stds=(
+                self._trainer.datamodule.predict_dataset.image_stds
+                if from_prediction
+                else self._trainer.datamodule.train_dataset.image_stats.stds
+            ),
         )
         denormalized_output = denorm(patch=output.cpu().numpy())
 
@@ -244,7 +269,7 @@ class VAEModule(L.LightningModule):
         Learning rate scheduler name.
     """
 
-    def __init__(self, algorithm_config: Union[VAEAlgorithmConfig, dict]) -> None:
+    def __init__(self, algorithm_config: Union[VAEBasedAlgorithm, dict]) -> None:
         """Lightning module for CAREamics.
 
         This class encapsulates the a PyTorch model along with the training, validation,
@@ -258,7 +283,7 @@ class VAEModule(L.LightningModule):
         super().__init__()
         # if loading from a checkpoint, AlgorithmModel needs to be instantiated
         self.algorithm_config = (
-            VAEAlgorithmConfig(**algorithm_config)
+            VAEBasedAlgorithm(**algorithm_config)
             if isinstance(algorithm_config, dict)
             else algorithm_config
         )
@@ -266,29 +291,27 @@ class VAEModule(L.LightningModule):
         # TODO: log algorithm config
         # self.save_hyperparameters(self.algorithm_config.model_dump())
 
-        # create model and loss function
+        # create model
         self.model: nn.Module = model_factory(self.algorithm_config.model)
-        self.noise_model: NoiseModel = noise_model_factory(
+
+        # create loss function
+        self.noise_model: Optional[NoiseModel] = noise_model_factory(
             self.algorithm_config.noise_model
         )
-        # TODO: here we can add some code to check whether the noise model is not None
-        # and `self.algorithm_config.noise_model_likelihood_model.noise_model` is,
-        # instead, None. In that case we could assign the noise model to the latter.
-        # This is particular useful when loading an algorithm config from file.
-        # Indeed, in that case the noise model in the nm likelihood is likely
-        # not available since excluded from serializaion.
-        self.noise_model_likelihood: NoiseModelLikelihood = likelihood_factory(
-            self.algorithm_config.noise_model_likelihood_model
+
+        self.noise_model_likelihood: Optional[NoiseModelLikelihood] = (
+            likelihood_factory(
+                config=self.algorithm_config.noise_model_likelihood,
+                noise_model=self.noise_model,
+            )
         )
-        self.gaussian_likelihood: GaussianLikelihood = likelihood_factory(
-            self.algorithm_config.gaussian_likelihood_model
+
+        self.gaussian_likelihood: Optional[GaussianLikelihood] = likelihood_factory(
+            self.algorithm_config.gaussian_likelihood
         )
-        self.loss_parameters = LVAELossParameters(
-            noise_model_likelihood=self.noise_model_likelihood,
-            gaussian_likelihood=self.gaussian_likelihood,
-            # TODO: musplit/denoisplit weights ?
-        )  # type: ignore
-        self.loss_func = loss_factory(self.algorithm_config.loss)
+
+        self.loss_parameters = self.algorithm_config.loss
+        self.loss_func = loss_factory(self.algorithm_config.loss.loss_type)
 
         # save optimizer and lr_scheduler names and parameters
         self.optimizer_name = self.algorithm_config.optimizer.name
@@ -344,11 +367,16 @@ class VAEModule(L.LightningModule):
         out = self.model(x)
 
         # Update loss parameters
-        # TODO rethink loss parameters
-        self.loss_parameters.current_epoch = self.current_epoch
+        self.loss_parameters.kl_params.current_epoch = self.current_epoch
 
         # Compute loss
-        loss = self.loss_func(out, target, self.loss_parameters)  # TODO ugly ?
+        loss = self.loss_func(
+            model_outputs=out,
+            targets=target,
+            config=self.loss_parameters,
+            gaussian_likelihood=self.gaussian_likelihood,
+            noise_model_likelihood=self.noise_model_likelihood,
+        )
 
         # Logging
         # TODO: implement a separate logging method?
@@ -376,7 +404,13 @@ class VAEModule(L.LightningModule):
         out = self.model(x)
 
         # Compute loss
-        loss = self.loss_func(out, target, self.loss_parameters)
+        loss = self.loss_func(
+            model_outputs=out,
+            targets=target,
+            config=self.loss_parameters,
+            gaussian_likelihood=self.gaussian_likelihood,
+            noise_model_likelihood=self.noise_model_likelihood,
+        )
 
         # Logging
         # Rename val_loss dict
@@ -625,9 +659,10 @@ def create_careamics_module(
     algorithm_configuration["model"] = model_configuration
 
     # call the parent init using an AlgorithmModel instance
+    # TODO broken by new configutations!
     algorithm_str = algorithm_configuration["algorithm"]
-    if algorithm_str in FCNAlgorithmConfig.get_compatible_algorithms():
-        return FCNModule(FCNAlgorithmConfig(**algorithm_configuration))
+    if algorithm_str in UNetBasedAlgorithm.get_compatible_algorithms():
+        return FCNModule(UNetBasedAlgorithm(**algorithm_configuration))
     else:
         raise NotImplementedError(
             f"Model {algorithm_str} is not implemented or unknown."

careamics/lightning/train_data_module.py

@@ -2,13 +2,14 @@
 
 from pathlib import Path
 from typing import Any, Callable, Literal, Optional, Union
+from warnings import warn
 
 import numpy as np
 import pytorch_lightning as L
 from numpy.typing import NDArray
-from torch.utils.data import DataLoader
+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 (
@@ -118,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,
@@ -218,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
@@ -446,6 +447,19 @@ class TrainDataModule(L.LightningDataModule):
         Any
             Training dataloader.
         """
+        # check because iterable dataset cannot be shuffled
+        if not isinstance(self.train_dataset, IterableDataset):
+            if ("shuffle" in self.dataloader_params) and (
+                not self.dataloader_params["shuffle"]
+            ):
+                warn(
+                    "Dataloader parameters include `shuffle=False`, this will be "
+                    "passed to the training dataloader and may result in bad results.",
+                    stacklevel=1,
+                )
+            else:
+                self.dataloader_params["shuffle"] = True
+
         return DataLoader(
             self.train_dataset, batch_size=self.batch_size, **self.dataloader_params
         )
@@ -488,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
@@ -504,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
@@ -613,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,
@@ -645,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/losses/loss_factory.py

@@ -7,7 +7,7 @@ This module contains a factory function for creating loss functions.
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Callable, Literal, Optional, Union
+from typing import Callable, Union
 
 from torch import Tensor as tensor
 
@@ -15,18 +15,6 @@ from ..config.support import SupportedLoss
 from .fcn.losses import mae_loss, mse_loss, n2v_loss
 from .lvae.losses import denoisplit_loss, denoisplit_musplit_loss, musplit_loss
 
-if TYPE_CHECKING:
-    from careamics.models.lvae.likelihoods import (
-        GaussianLikelihood,
-        NoiseModelLikelihood,
-    )
-    from careamics.models.lvae.noise_models import (
-        GaussianMixtureNoiseModel,
-        MultiChannelNoiseModel,
-    )
-
-    NoiseModel = Union[GaussianMixtureNoiseModel, MultiChannelNoiseModel]
-
 
 @dataclass
 class FCNLossParameters:
@@ -40,78 +28,6 @@ class FCNLossParameters:
     loss_weight: float
 
 
-@dataclass  # TODO why not pydantic?
-class LVAELossParameters:
-    """Dataclass for LVAE loss."""
-
-    # TODO: refactor in more modular blocks (otherwise it gets messy very easily)
-    # e.g., - weights, - kl_params, ...
-
-    noise_model_likelihood: Optional[NoiseModelLikelihood] = None
-    """Noise model likelihood instance."""
-    gaussian_likelihood: Optional[GaussianLikelihood] = None
-    """Gaussian likelihood instance."""
-    current_epoch: int = 0
-    """Current epoch in the training loop."""
-    reconstruction_weight: float = 1.0
-    """Weight for the reconstruction loss in the total net loss
-    (i.e., `net_loss = reconstruction_weight * rec_loss + kl_weight * kl_loss`)."""
-    musplit_weight: float = 0.1
-    """Weight for the muSplit loss (used in the muSplit-denoiSplit loss)."""
-    denoisplit_weight: float = 0.9
-    """Weight for the denoiSplit loss (used in the muSplit-deonoiSplit loss)."""
-    kl_type: Literal["kl", "kl_restricted", "kl_spatial", "kl_channelwise"] = "kl"
-    """Type of KL divergence used as KL loss."""
-    kl_weight: float = 1.0
-    """Weight for the KL loss in the total net loss.
-    (i.e., `net_loss = reconstruction_weight * rec_loss + kl_weight * kl_loss`)."""
-    kl_annealing: bool = False
-    """Whether to apply KL loss annealing."""
-    kl_start: int = -1
-    """Epoch at which KL loss annealing starts."""
-    kl_annealtime: int = 10
-    """Number of epochs for which KL loss annealing is applied."""
-    non_stochastic: bool = False
-    """Whether to sample latents and compute KL."""
-
-
-# TODO: really needed?
-# like it is now, it is difficult to use, we need a way to specify the
-# loss parameters in a more user-friendly way.
-def loss_parameters_factory(
-    type: SupportedLoss,
-) -> Union[FCNLossParameters, LVAELossParameters]:
-    """Return loss parameters.
-
-    Parameters
-    ----------
-    type : SupportedLoss
-        Requested loss.
-
-    Returns
-    -------
-    Union[FCNLossParameters, LVAELossParameters]
-        Loss parameters.
-
-    Raises
-    ------
-    NotImplementedError
-        If the loss is unknown.
-    """
-    if type in [SupportedLoss.N2V, SupportedLoss.MSE, SupportedLoss.MAE]:
-        return FCNLossParameters
-
-    elif type in [
-        SupportedLoss.MUSPLIT,
-        SupportedLoss.DENOISPLIT,
-        SupportedLoss.DENOISPLIT_MUSPLIT,
-    ]:
-        return LVAELossParameters  # it returns the class, not an instance
-
-    else:
-        raise NotImplementedError(f"Loss {type} is not yet supported.")
-
-
 def loss_factory(loss: Union[SupportedLoss, str]) -> Callable:
     """Return loss function.