careamics 0.1.0rc2__py3-none-any.whl → 0.1.0rc4__py3-none-any.whl

careamics/lightning_module.py

@@ -0,0 +1,292 @@
+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 a PyTorch model along with the training, validation,
+    and testing logic. It is configured using an `AlgorithmModel` Pydantic class.
+
+    Attributes
+    ----------
+    model : nn.Module
+        PyTorch model.
+    loss_func : 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:
+        """
+        CAREamics Lightning module.
+
+        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 : Union[AlgorithmModel, 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 : Tensor
+            Input batch.
+        batch_idx : Any
+            Batch index.
+
+        Returns
+        -------
+        Any
+            Loss value.
+        """
+        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 : 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 : Tensor
+            Input batch.
+        batch_idx : Any
+            Batch index.
+
+        Returns
+        -------
+        Any
+            Model output.
+        """
+        x, *aux = batch
+
+        # 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(batch[0])  # 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(
+            mean=self._trainer.datamodule.predict_dataset.mean,
+            std=self._trainer.datamodule.predict_dataset.std,
+        )
+        denormalized_output = denorm(image=output)["image"]
+
+        if len(aux) > 0:
+            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
+        }
+
+
+class CAREamicsModuleWrapper(CAREamicsModule):
+    """Class defining the API for CAREamics Lightning layer.
+
+    This class exposes parameters used to create an AlgorithmModel instance, triggering
+    parameters validation.
+
+    Parameters
+    ----------
+    algorithm : Union[SupportedAlgorithm, str]
+        Algorithm to use for training (see SupportedAlgorithm).
+    loss : Union[SupportedLoss, str]
+        Loss function to use for training (see SupportedLoss).
+    architecture : Union[SupportedArchitecture, 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 : Union[SupportedOptimizer, 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 : Union[SupportedScheduler, 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 {}.
+    """
+
+    def __init__(
+        self,
+        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,
+    ) -> None:
+        """
+        Wrapper for the CAREamics model, exposing all algorithm configuration arguments.
+
+        Parameters
+        ----------
+        algorithm : Union[SupportedAlgorithm, str]
+            Algorithm to use for training (see SupportedAlgorithm).
+        loss : Union[SupportedLoss, str]
+            Loss function to use for training (see SupportedLoss).
+        architecture : Union[SupportedArchitecture, 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 : Union[SupportedOptimizer, 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 : Union[SupportedScheduler, 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 {}.
+        """
+        # 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
+        super().__init__(AlgorithmConfig(**algorithm_configuration))
+
+        # TODO add load_from_checkpoint wrapper

careamics/lightning_prediction_datamodule.py

@@ -0,0 +1,396 @@
+"""Prediction Lightning data modules."""
+from pathlib import Path
+from typing import Any, Callable, List, Literal, Optional, Tuple, Union
+
+import numpy as np
+import pytorch_lightning as L
+from albumentations import Compose
+from torch.utils.data import DataLoader
+from torch.utils.data.dataloader import default_collate
+
+from careamics.config import InferenceConfig
+from careamics.config.support import SupportedData
+from careamics.config.tile_information import TileInformation
+from careamics.dataset.dataset_utils import (
+    get_read_func,
+    list_files,
+)
+from careamics.dataset.in_memory_dataset import (
+    InMemoryPredictionDataset,
+)
+from careamics.dataset.iterable_dataset import (
+    IterablePredictionDataset,
+)
+from careamics.utils import get_logger
+
+PredictDatasetType = Union[InMemoryPredictionDataset, IterablePredictionDataset]
+
+logger = get_logger(__name__)
+
+
+def _collate_tiles(batch: List[Tuple[np.ndarray, TileInformation]]) -> Any:
+    """
+    Collate tiles received from CAREamics prediction dataloader.
+
+    CAREamics prediction dataloader returns tuples of arrays and TileInformation. In
+    case of non-tiled data, this function will return the arrays. In case of tiled data,
+    it will return the arrays, the last tile flag, the overlap crop coordinates and the
+    stitch coordinates.
+
+    Parameters
+    ----------
+    batch : Tuple[Tuple[np.ndarray, TileInformation], ...]
+        Batch of tiles.
+
+    Returns
+    -------
+    Any
+        Collated batch.
+    """
+    first_tile_info: TileInformation = batch[0][1]
+    # if not tiled, then return arrays
+    if not first_tile_info.tiled:
+        arrays, _ = zip(*batch)
+
+        return default_collate(arrays)
+    # else we explicit the last_tile flag and coordinates
+    else:
+        new_batch = [
+            (tile, t.last_tile, t.array_shape, t.overlap_crop_coords, t.stitch_coords)
+            for tile, t in batch
+        ]
+
+        return default_collate(new_batch)
+
+
+class CAREamicsPredictData(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 : Union[Path, str, np.ndarray]
+        Prediction data, can be a path to a folder, a file or a numpy array.
+    read_source_func : Optional[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, np.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 : Union[Path, str, np.ndarray]
+            Prediction data, can be a path to a folder, a file or a numpy array.
+        read_source_func : Optional[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
+
+        # 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:
+            # prediction dataset
+            self.predict_dataset: PredictDatasetType = InMemoryPredictionDataset(
+                prediction_config=self.prediction_config,
+                inputs=self.pred_data,
+            )
+        else:
+            self.predict_dataset = IterablePredictionDataset(
+                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,
+            **self.dataloader_params,
+        )  # TODO check workers are used
+
+
+class PredictDataWrapper(CAREamicsPredictData):
+    """
+    Wrapper around the CAREamics inference Lightning data module.
+
+    This class is used to explicitely pass the parameters usually contained in a
+    `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.
+
+    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`.
+
+    The default transformations applied to the images are defined in
+    `careamics.config.inference_model`. To use different transformations, pass a list
+    of transforms or an albumentation `Compose` as `transforms` parameter. See examples
+    for more details.
+
+    The `mean` and `std` parameters are only used if Normalization is defined either
+    in the default transformations or in the `transforms` parameter, but not with
+    a `Compose` object. If you pass a `Normalization` transform in a list as
+    `transforms`, then the mean and std parameters will be overwritten by those passed
+    to this method.
+
+    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 : Union[str, Path, np.ndarray]
+        Prediction data.
+    data_type : Union[Literal["array", "tiff", "custom"], SupportedData]
+        Data type, see `SupportedData` for available options.
+    mean : float
+        Mean value for normalization, only used if Normalization is defined in the
+        transforms.
+    std : float
+        Standard deviation value for normalization, only used if Normalization is
+        defined in the transform.
+    tile_size : Tuple[int, ...]
+        Tile size, 2D or 3D tile size.
+    tile_overlap : Tuple[int, ...]
+        Tile overlap, 2D or 3D tile overlap.
+    axes : str
+        Axes of the data, choosen amongst SCZYX.
+    batch_size : int
+        Batch size.
+    tta_transforms : bool, optional
+        Use test time augmentation, by default True.
+    transforms : Optional[Union[List[TRANSFORMS_UNION], Compose]], optional
+        List of transforms to apply to prediction patches. If None, default
+        transforms are applied.
+    read_source_func : Optional[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 {}.
+    """
+
+    def __init__(
+        self,
+        pred_data: Union[str, Path, np.ndarray],
+        data_type: Union[Literal["array", "tiff", "custom"], SupportedData],
+        mean: float,
+        std: float,
+        tile_size: Optional[Tuple[int, ...]] = None,
+        tile_overlap: Optional[Tuple[int, ...]] = None,
+        axes: str = "YX",
+        batch_size: int = 1,
+        tta_transforms: bool = True,
+        transforms: Optional[Union[List, Compose]] = None,
+        read_source_func: Optional[Callable] = None,
+        extension_filter: str = "",
+        dataloader_params: Optional[dict] = None,
+    ) -> None:
+        """
+        Constructor.
+
+        Parameters
+        ----------
+        pred_data : Union[str, Path, np.ndarray]
+            Prediction data.
+        data_type : Union[Literal["array", "tiff", "custom"], SupportedData]
+            Data type, see `SupportedData` for available options.
+        mean : float
+            Mean value for normalization, only used if Normalization is defined in the
+            transforms.
+        std : float
+            Standard deviation value for normalization, only used if Normalization is
+            defined in the transform.
+        tile_size : List[int]
+            Tile size, 2D or 3D tile size.
+        tile_overlap : List[int]
+            Tile overlap, 2D or 3D tile overlap.
+        axes : str
+            Axes of the data, choosen amongst SCZYX.
+        batch_size : int
+            Batch size.
+        tta_transforms : bool, optional
+            Use test time augmentation, by default True.
+        transforms : Optional[Union[List[TRANSFORMS_UNION], Compose]], optional
+            List of transforms to apply to prediction patches. If None, default
+            transforms are applied.
+        read_source_func : Optional[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 {}.
+        """
+        if dataloader_params is None:
+            dataloader_params = {}
+        prediction_dict = {
+            "data_type": data_type,
+            "tile_size": tile_size,
+            "tile_overlap": tile_overlap,
+            "axes": axes,
+            "mean": mean,
+            "std": std,
+            "tta": tta_transforms,
+            "batch_size": batch_size,
+        }
+
+        # if transforms are passed (otherwise it will use the default ones)
+        if transforms is not None:
+            prediction_dict["transforms"] = transforms
+
+        # validate configuration
+        self.prediction_config = InferenceConfig(**prediction_dict)
+
+        # sanity check on the dataloader parameters
+        if "batch_size" in dataloader_params:
+            # remove it
+            del dataloader_params["batch_size"]
+
+        super().__init__(
+            pred_config=self.prediction_config,
+            pred_data=pred_data,
+            read_source_func=read_source_func,
+            extension_filter=extension_filter,
+            dataloader_params=dataloader_params,
+        )