careamics 0.0.1__py3-none-any.whl → 0.0.2__py3-none-any.whl

careamics/lightning/lightning_module.py

@@ -0,0 +1,276 @@
+"""CAREamics Lightning module."""
+
+from typing import Any, Optional, Union
+
+import pytorch_lightning as L
+from torch import Tensor, nn
+
+from careamics.config import AlgorithmConfig
+from careamics.config.support import (
+    SupportedAlgorithm,
+    SupportedArchitecture,
+    SupportedLoss,
+    SupportedOptimizer,
+    SupportedScheduler,
+)
+from careamics.losses import loss_factory
+from careamics.models.model_factory import model_factory
+from careamics.transforms import Denormalize, ImageRestorationTTA
+from careamics.utils.torch_utils import get_optimizer, get_scheduler
+
+
+class CAREamicsModule(L.LightningModule):
+    """
+    CAREamics Lightning module.
+
+    This class encapsulates the PyTorch model along with the training, validation,
+    and testing logic. It is configured using an `AlgorithmModel` Pydantic class.
+
+    Parameters
+    ----------
+    algorithm_config : AlgorithmModel or dict
+        Algorithm configuration.
+
+    Attributes
+    ----------
+    model : torch.nn.Module
+        PyTorch model.
+    loss_func : torch.nn.Module
+        Loss function.
+    optimizer_name : str
+        Optimizer name.
+    optimizer_params : dict
+        Optimizer parameters.
+    lr_scheduler_name : str
+        Learning rate scheduler name.
+    """
+
+    def __init__(self, algorithm_config: Union[AlgorithmConfig, dict]) -> None:
+        """Lightning module for CAREamics.
+
+        This class encapsulates the a PyTorch model along with the training, validation,
+        and testing logic. It is configured using an `AlgorithmModel` Pydantic class.
+
+        Parameters
+        ----------
+        algorithm_config : AlgorithmModel or dict
+            Algorithm configuration.
+        """
+        super().__init__()
+        # if loading from a checkpoint, AlgorithmModel needs to be instantiated
+        if isinstance(algorithm_config, dict):
+            algorithm_config = AlgorithmConfig(**algorithm_config)
+
+        # create model and loss function
+        self.model: nn.Module = model_factory(algorithm_config.model)
+        self.loss_func = loss_factory(algorithm_config.loss)
+
+        # save optimizer and lr_scheduler names and parameters
+        self.optimizer_name = algorithm_config.optimizer.name
+        self.optimizer_params = algorithm_config.optimizer.parameters
+        self.lr_scheduler_name = algorithm_config.lr_scheduler.name
+        self.lr_scheduler_params = algorithm_config.lr_scheduler.parameters
+
+    def forward(self, x: Any) -> Any:
+        """Forward pass.
+
+        Parameters
+        ----------
+        x : Any
+            Input tensor.
+
+        Returns
+        -------
+        Any
+            Output tensor.
+        """
+        return self.model(x)
+
+    def training_step(self, batch: Tensor, batch_idx: Any) -> Any:
+        """Training step.
+
+        Parameters
+        ----------
+        batch : torch.Tensor
+            Input batch.
+        batch_idx : Any
+            Batch index.
+
+        Returns
+        -------
+        Any
+            Loss value.
+        """
+        # TODO can N2V be simplified by returning mask*original_patch
+        x, *aux = batch
+        out = self.model(x)
+        loss = self.loss_func(out, *aux)
+        self.log(
+            "train_loss", loss, on_step=True, on_epoch=True, prog_bar=True, logger=True
+        )
+        return loss
+
+    def validation_step(self, batch: Tensor, batch_idx: Any) -> None:
+        """Validation step.
+
+        Parameters
+        ----------
+        batch : torch.Tensor
+            Input batch.
+        batch_idx : Any
+            Batch index.
+        """
+        x, *aux = batch
+        out = self.model(x)
+        val_loss = self.loss_func(out, *aux)
+
+        # log validation loss
+        self.log(
+            "val_loss",
+            val_loss,
+            on_step=False,
+            on_epoch=True,
+            prog_bar=True,
+            logger=True,
+        )
+
+    def predict_step(self, batch: Tensor, batch_idx: Any) -> Any:
+        """Prediction step.
+
+        Parameters
+        ----------
+        batch : torch.Tensor
+            Input batch.
+        batch_idx : Any
+            Batch index.
+
+        Returns
+        -------
+        Any
+            Model output.
+        """
+        if self._trainer.datamodule.tiled:
+            x, *aux = batch
+        else:
+            x = batch
+            aux = []
+
+        # apply test-time augmentation if available
+        # TODO: probably wont work with batch size > 1
+        if self._trainer.datamodule.prediction_config.tta_transforms:
+            tta = ImageRestorationTTA()
+            augmented_batch = tta.forward(x)  # list of augmented tensors
+            augmented_output = []
+            for augmented in augmented_batch:
+                augmented_pred = self.model(augmented)
+                augmented_output.append(augmented_pred)
+            output = tta.backward(augmented_output)
+        else:
+            output = self.model(x)
+
+        # Denormalize the output
+        denorm = Denormalize(
+            image_means=self._trainer.datamodule.predict_dataset.image_means,
+            image_stds=self._trainer.datamodule.predict_dataset.image_stds,
+        )
+        denormalized_output = denorm(patch=output.cpu().numpy())
+
+        if len(aux) > 0:  # aux can be tiling information
+            return denormalized_output, *aux
+        else:
+            return denormalized_output
+
+    def configure_optimizers(self) -> Any:
+        """Configure optimizers and learning rate schedulers.
+
+        Returns
+        -------
+        Any
+            Optimizer and learning rate scheduler.
+        """
+        # instantiate optimizer
+        optimizer_func = get_optimizer(self.optimizer_name)
+        optimizer = optimizer_func(self.model.parameters(), **self.optimizer_params)
+
+        # and scheduler
+        scheduler_func = get_scheduler(self.lr_scheduler_name)
+        scheduler = scheduler_func(optimizer, **self.lr_scheduler_params)
+
+        return {
+            "optimizer": optimizer,
+            "lr_scheduler": scheduler,
+            "monitor": "val_loss",  # otherwise triggers MisconfigurationException
+        }
+
+
+def create_careamics_module(
+    algorithm: Union[SupportedAlgorithm, str],
+    loss: Union[SupportedLoss, str],
+    architecture: Union[SupportedArchitecture, str],
+    model_parameters: Optional[dict] = None,
+    optimizer: Union[SupportedOptimizer, str] = "Adam",
+    optimizer_parameters: Optional[dict] = None,
+    lr_scheduler: Union[SupportedScheduler, str] = "ReduceLROnPlateau",
+    lr_scheduler_parameters: Optional[dict] = None,
+) -> CAREamicsModule:
+    """Create a CAREamics Lithgning module.
+
+    This function exposes parameters used to create an AlgorithmModel instance,
+    triggering parameters validation.
+
+    Parameters
+    ----------
+    algorithm : SupportedAlgorithm or str
+        Algorithm to use for training (see SupportedAlgorithm).
+    loss : SupportedLoss or str
+        Loss function to use for training (see SupportedLoss).
+    architecture : SupportedArchitecture or str
+        Model architecture to use for training (see SupportedArchitecture).
+    model_parameters : dict, optional
+        Model parameters to use for training, by default {}. Model parameters are
+        defined in the relevant `torch.nn.Module` class, or Pyddantic model (see
+        `careamics.config.architectures`).
+    optimizer : SupportedOptimizer or str, optional
+        Optimizer to use for training, by default "Adam" (see SupportedOptimizer).
+    optimizer_parameters : dict, optional
+        Optimizer parameters to use for training, as defined in `torch.optim`, by
+        default {}.
+    lr_scheduler : SupportedScheduler or str, optional
+        Learning rate scheduler to use for training, by default "ReduceLROnPlateau"
+        (see SupportedScheduler).
+    lr_scheduler_parameters : dict, optional
+        Learning rate scheduler parameters to use for training, as defined in
+        `torch.optim`, by default {}.
+
+    Returns
+    -------
+    CAREamicsModule
+        CAREamics Lightning module.
+    """
+    # create a AlgorithmModel compatible dictionary
+    if lr_scheduler_parameters is None:
+        lr_scheduler_parameters = {}
+    if optimizer_parameters is None:
+        optimizer_parameters = {}
+    if model_parameters is None:
+        model_parameters = {}
+    algorithm_configuration = {
+        "algorithm": algorithm,
+        "loss": loss,
+        "optimizer": {
+            "name": optimizer,
+            "parameters": optimizer_parameters,
+        },
+        "lr_scheduler": {
+            "name": lr_scheduler,
+            "parameters": lr_scheduler_parameters,
+        },
+    }
+    model_configuration = {"architecture": architecture}
+    model_configuration.update(model_parameters)
+
+    # add model parameters to algorithm configuration
+    algorithm_configuration["model"] = model_configuration
+
+    # call the parent init using an AlgorithmModel instance
+    return CAREamicsModule(AlgorithmConfig(**algorithm_configuration))

careamics/lightning/predict_data_module.py

@@ -0,0 +1,333 @@
+"""Prediction Lightning data modules."""
+
+from pathlib import Path
+from typing import Any, Callable, Literal, Optional, Union
+
+import numpy as np
+import pytorch_lightning as L
+from numpy.typing import NDArray
+from torch.utils.data import DataLoader
+
+from careamics.config import InferenceConfig
+from careamics.config.support import SupportedData
+from careamics.dataset import (
+    InMemoryPredDataset,
+    InMemoryTiledPredDataset,
+    IterablePredDataset,
+    IterableTiledPredDataset,
+)
+from careamics.dataset.dataset_utils import list_files
+from careamics.dataset.tiling.collate_tiles import collate_tiles
+from careamics.file_io.read import get_read_func
+from careamics.utils import get_logger
+
+PredictDatasetType = Union[
+    InMemoryPredDataset,
+    InMemoryTiledPredDataset,
+    IterablePredDataset,
+    IterableTiledPredDataset,
+]
+
+logger = get_logger(__name__)
+
+
+class PredictDataModule(L.LightningDataModule):
+    """
+    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.
+
+    To read custom data types, you can set `data_type` to `custom` in `data_config`
+    and provide a function that returns a numpy array from a path as
+    `read_source_func` parameter. The function will receive a Path object and
+    an axies string as arguments, the axes being derived from the `data_config`.
+
+    You can also provide a `fnmatch` and `Path.rglob` compatible expression (e.g.
+    "*.czi") to filter the files extension using `extension_filter`.
+
+    Parameters
+    ----------
+    pred_config : InferenceModel
+        Pydantic model for CAREamics prediction configuration.
+    pred_data : pathlib.Path or str or numpy.ndarray
+        Prediction data, can be a path to a folder, a file or a numpy array.
+    read_source_func : Callable, optional
+        Function to read custom types, by default None.
+    extension_filter : str, optional
+        Filter to filter file extensions for custom types, by default "".
+    dataloader_params : dict, optional
+        Dataloader parameters, by default {}.
+    """
+
+    def __init__(
+        self,
+        pred_config: InferenceConfig,
+        pred_data: Union[Path, str, NDArray],
+        read_source_func: Optional[Callable] = None,
+        extension_filter: str = "",
+        dataloader_params: Optional[dict] = None,
+    ) -> None:
+        """
+        Constructor.
+
+        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.
+
+        To read custom data types, you can set `data_type` to `custom` in `data_config`
+        and provide a function that returns a numpy array from a path as
+        `read_source_func` parameter. The function will receive a Path object and
+        an axies string as arguments, the axes being derived from the `data_config`.
+
+        You can also provide a `fnmatch` and `Path.rglob` compatible expression (e.g.
+        "*.czi") to filter the files extension using `extension_filter`.
+
+        Parameters
+        ----------
+        pred_config : InferenceModel
+            Pydantic model for CAREamics prediction configuration.
+        pred_data : pathlib.Path or str or numpy.ndarray
+            Prediction data, can be a path to a folder, a file or a numpy array.
+        read_source_func : Callable, optional
+            Function to read custom types, by default None.
+        extension_filter : str, optional
+            Filter to filter file extensions for custom types, by default "".
+        dataloader_params : dict, optional
+            Dataloader parameters, by default {}.
+
+        Raises
+        ------
+        ValueError
+            If the data type is `custom` and no `read_source_func` is provided.
+        ValueError
+            If the data type is `array` and the input is not a numpy array.
+        ValueError
+            If the data type is `tiff` and the input is neither a Path nor a str.
+        """
+        if dataloader_params is None:
+            dataloader_params = {}
+        if dataloader_params is None:
+            dataloader_params = {}
+        super().__init__()
+
+        # check that a read source function is provided for custom types
+        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` and an `extension_filer`."
+            )
+
+        # check correct input type
+        if (
+            isinstance(pred_data, np.ndarray)
+            and pred_config.data_type != SupportedData.ARRAY
+        ):
+            raise ValueError(
+                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 (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"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 = 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 = pred_config.tile_size
+        self.tile_overlap = pred_config.tile_overlap
+
+        # check if it is tiled
+        self.tiled = self.tile_size is not None and self.tile_overlap is not None
+
+        # read source function
+        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 pred_config.data_type != SupportedData.ARRAY:
+            self.read_source_func = get_read_func(pred_config.data_type)
+
+        self.extension_filter = extension_filter
+
+    def prepare_data(self) -> None:
+        """Hook used to prepare the data before calling `setup`."""
+        # if the data is a Path or a str
+        if not isinstance(self.pred_data, np.ndarray):
+            self.pred_files = list_files(
+                self.pred_data, self.data_type, self.extension_filter
+            )
+
+    def setup(self, stage: Optional[str] = None) -> None:
+        """
+        Hook called at the beginning of predict.
+
+        Parameters
+        ----------
+        stage : Optional[str], optional
+            Stage, by default None.
+        """
+        # if numpy array
+        if self.data_type == SupportedData.ARRAY:
+            if self.tiled:
+                self.predict_dataset: PredictDatasetType = InMemoryTiledPredDataset(
+                    prediction_config=self.prediction_config,
+                    inputs=self.pred_data,
+                )
+            else:
+                self.predict_dataset = InMemoryPredDataset(
+                    prediction_config=self.prediction_config,
+                    inputs=self.pred_data,
+                )
+        else:
+            if self.tiled:
+                self.predict_dataset = IterableTiledPredDataset(
+                    prediction_config=self.prediction_config,
+                    src_files=self.pred_files,
+                    read_source_func=self.read_source_func,
+                )
+            else:
+                self.predict_dataset = IterablePredDataset(
+                    prediction_config=self.prediction_config,
+                    src_files=self.pred_files,
+                    read_source_func=self.read_source_func,
+                )
+
+    def predict_dataloader(self) -> DataLoader:
+        """
+        Create a dataloader for prediction.
+
+        Returns
+        -------
+        DataLoader
+            Prediction dataloader.
+        """
+        return DataLoader(
+            self.predict_dataset,
+            batch_size=self.batch_size,
+            collate_fn=collate_tiles if self.tiled else None,
+            **self.dataloader_params,
+        )
+
+
+def create_predict_datamodule(
+    pred_data: Union[str, Path, NDArray],
+    data_type: Union[Literal["array", "tiff", "custom"], SupportedData],
+    axes: str,
+    image_means: list[float],
+    image_stds: list[float],
+    tile_size: Optional[tuple[int, ...]] = None,
+    tile_overlap: Optional[tuple[int, ...]] = None,
+    batch_size: int = 1,
+    tta_transforms: bool = True,
+    read_source_func: Optional[Callable] = None,
+    extension_filter: str = "",
+    dataloader_params: Optional[dict] = None,
+) -> PredictDataModule:
+    """Create a CAREamics prediction Lightning datamodule.
+
+    This function is used to explicitely pass the parameters usually contained in an
+    `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
+    and are coherent. This can be done by creating a `Configuration` object beforehand
+    and passing its parameters to the different Lightning modules.
+
+    The data module can be used with Path, str or numpy arrays. To use array data, set
+    `data_type` to `array` and pass a numpy array to `train_data`.
+
+    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
+    path. Additionally, pass a `fnmatch` and `Path.rglob` compatible expression
+    (e.g. "*.jpeg") to filter the files extension using `extension_filter`.
+
+    In `dataloader_params`, you can pass any parameter accepted by PyTorch
+    dataloaders, except for `batch_size`, which is set by the `batch_size`
+    parameter.
+
+    Parameters
+    ----------
+    pred_data : str or pathlib.Path or numpy.ndarray
+        Prediction data.
+    data_type : {"array", "tiff", "custom"}
+        Data type, see `SupportedData` for available options.
+    axes : str
+        Axes of the data, choosen among SCZYX.
+    image_means : list of float
+        Mean values for normalization, only used if Normalization is defined.
+    image_stds : list of float
+        Std values for normalization, only used if Normalization is defined.
+    tile_size : tuple of int, optional
+        Tile size, 2D or 3D tile size.
+    tile_overlap : tuple of int, optional
+        Tile overlap, 2D or 3D tile overlap.
+    batch_size : int
+        Batch size.
+    tta_transforms : bool, optional
+        Use test time augmentation, by default True.
+    read_source_func : Callable, optional
+        Function to read the source data, used if `data_type` is `custom`, by
+        default None.
+    extension_filter : str, optional
+        Filter for file extensions, used if `data_type` is `custom`, by default "".
+    dataloader_params : dict, optional
+        Pytorch dataloader parameters, by default {}.
+
+    Returns
+    -------
+    PredictDataModule
+        CAREamics prediction datamodule.
+
+    Notes
+    -----
+    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.
+    """
+    if dataloader_params is None:
+        dataloader_params = {}
+
+    prediction_dict: dict[str, Any] = {
+        "data_type": data_type,
+        "tile_size": tile_size,
+        "tile_overlap": tile_overlap,
+        "axes": axes,
+        "image_means": image_means,
+        "image_stds": image_stds,
+        "tta_transforms": tta_transforms,
+        "batch_size": batch_size,
+    }
+
+    # validate configuration
+    prediction_config = InferenceConfig(**prediction_dict)
+
+    # sanity check on the dataloader parameters
+    if "batch_size" in dataloader_params:
+        # remove it
+        del dataloader_params["batch_size"]
+
+    return PredictDataModule(
+        pred_config=prediction_config,
+        pred_data=pred_data,
+        read_source_func=read_source_func,
+        extension_filter=extension_filter,
+        dataloader_params=dataloader_params,
+    )