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

careamics/lightning/lightning_module.py

@@ -1,11 +1,12 @@
 """CAREamics Lightning module."""
 
-from typing import Any, Optional, Union
+from typing import Any, Callable, Optional, Union
 
+import numpy as np
 import pytorch_lightning as L
 from torch import Tensor, nn
 
-from careamics.config import AlgorithmConfig
+from careamics.config import FCNAlgorithmConfig, VAEAlgorithmConfig
 from careamics.config.support import (
     SupportedAlgorithm,
     SupportedArchitecture,
@@ -14,12 +15,26 @@ from careamics.config.support import (
     SupportedScheduler,
 )
 from careamics.losses import loss_factory
+from careamics.losses.loss_factory import LVAELossParameters
+from careamics.models.lvae.likelihoods import (
+    GaussianLikelihood,
+    NoiseModelLikelihood,
+    likelihood_factory,
+)
+from careamics.models.lvae.noise_models import (
+    GaussianMixtureNoiseModel,
+    MultiChannelNoiseModel,
+    noise_model_factory,
+)
 from careamics.models.model_factory import model_factory
 from careamics.transforms import Denormalize, ImageRestorationTTA
+from careamics.utils.metrics import RunningPSNR, scale_invariant_psnr
 from careamics.utils.torch_utils import get_optimizer, get_scheduler
 
+NoiseModel = Union[GaussianMixtureNoiseModel, MultiChannelNoiseModel]
+
 
-class CAREamicsModule(L.LightningModule):
+class FCNModule(L.LightningModule):
     """
     CAREamics Lightning module.
 
@@ -45,7 +60,7 @@ class CAREamicsModule(L.LightningModule):
         Learning rate scheduler name.
     """
 
-    def __init__(self, algorithm_config: Union[AlgorithmConfig, dict]) -> None:
+    def __init__(self, algorithm_config: Union[FCNAlgorithmConfig, dict]) -> None:
         """Lightning module for CAREamics.
 
         This class encapsulates the a PyTorch model along with the training, validation,
@@ -59,7 +74,7 @@ class CAREamicsModule(L.LightningModule):
         super().__init__()
         # if loading from a checkpoint, AlgorithmModel needs to be instantiated
         if isinstance(algorithm_config, dict):
-            algorithm_config = AlgorithmConfig(**algorithm_config)
+            algorithm_config = FCNAlgorithmConfig(**algorithm_config)
 
         # create model and loss function
         self.model: nn.Module = model_factory(algorithm_config.model)
@@ -203,6 +218,343 @@ class CAREamicsModule(L.LightningModule):
         }
 
 
+class VAEModule(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.
+
+    Parameters
+    ----------
+    algorithm_config : Union[VAEAlgorithmConfig, dict]
+        Algorithm configuration.
+
+    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[VAEAlgorithmConfig, 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 : Union[AlgorithmModel, dict]
+            Algorithm configuration.
+        """
+        super().__init__()
+        # if loading from a checkpoint, AlgorithmModel needs to be instantiated
+        self.algorithm_config = (
+            VAEAlgorithmConfig(**algorithm_config)
+            if isinstance(algorithm_config, dict)
+            else algorithm_config
+        )
+
+        # TODO: log algorithm config
+        # self.save_hyperparameters(self.algorithm_config.model_dump())
+
+        # create model and loss function
+        self.model: nn.Module = model_factory(self.algorithm_config.model)
+        self.noise_model: 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.gaussian_likelihood: GaussianLikelihood = likelihood_factory(
+            self.algorithm_config.gaussian_likelihood_model
+        )
+        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)
+
+        # save optimizer and lr_scheduler names and parameters
+        self.optimizer_name = self.algorithm_config.optimizer.name
+        self.optimizer_params = self.algorithm_config.optimizer.parameters
+        self.lr_scheduler_name = self.algorithm_config.lr_scheduler.name
+        self.lr_scheduler_params = self.algorithm_config.lr_scheduler.parameters
+
+        # initialize running PSNR
+        self.running_psnr = [
+            RunningPSNR() for _ in range(self.algorithm_config.model.output_channels)
+        ]
+
+    def forward(self, x: Tensor) -> tuple[Tensor, dict[str, Any]]:
+        """Forward pass.
+
+        Parameters
+        ----------
+        x : Tensor
+            Input tensor of shape (B, (1 + n_LC), [Z], Y, X), where n_LC is the
+            number of lateral inputs.
+
+        Returns
+        -------
+        tuple[Tensor, dict[str, Any]]
+            A tuple with the output tensor and additional data from the top-down pass.
+        """
+        return self.model(x)  # TODO Different model can have more than one output
+
+    def training_step(
+        self, batch: tuple[Tensor, Tensor], batch_idx: Any
+    ) -> Optional[dict[str, Tensor]]:
+        """Training step.
+
+        Parameters
+        ----------
+        batch : tuple[Tensor, Tensor]
+            Input batch. It is a tuple with the input tensor and the target tensor.
+            The input tensor has shape (B, (1 + n_LC), [Z], Y, X), where n_LC is the
+            number of lateral inputs. The target tensor has shape (B, C, [Z], Y, X),
+            where C is the number of target channels (e.g., 1 in HDN, >1 in
+            muSplit/denoiSplit).
+        batch_idx : Any
+            Batch index.
+
+        Returns
+        -------
+        Any
+            Loss value.
+        """
+        x, target = batch
+
+        # Forward pass
+        out = self.model(x)
+
+        # Update loss parameters
+        # TODO rethink loss parameters
+        self.loss_parameters.current_epoch = self.current_epoch
+
+        # Compute loss
+        loss = self.loss_func(out, target, self.loss_parameters)  # TODO ugly ?
+
+        # Logging
+        # TODO: implement a separate logging method?
+        self.log_dict(loss, on_step=True, on_epoch=True)
+        # self.log("lr", self, on_epoch=True)
+        return loss
+
+    def validation_step(self, batch: tuple[Tensor, Tensor], batch_idx: Any) -> None:
+        """Validation step.
+
+        Parameters
+        ----------
+        batch : tuple[Tensor, Tensor]
+            Input batch. It is a tuple with the input tensor and the target tensor.
+            The input tensor has shape (B, (1 + n_LC), [Z], Y, X), where n_LC is the
+            number of lateral inputs. The target tensor has shape (B, C, [Z], Y, X),
+            where C is the number of target channels (e.g., 1 in HDN, >1 in
+            muSplit/denoiSplit).
+        batch_idx : Any
+            Batch index.
+        """
+        x, target = batch
+
+        # Forward pass
+        out = self.model(x)
+
+        # Compute loss
+        loss = self.loss_func(out, target, self.loss_parameters)
+
+        # Logging
+        # Rename val_loss dict
+        loss = {"_".join(["val", k]): v for k, v in loss.items()}
+        self.log_dict(loss, on_epoch=True, prog_bar=True)
+        curr_psnr = self.compute_val_psnr(out, target)
+        for i, psnr in enumerate(curr_psnr):
+            self.log(f"val_psnr_ch{i+1}_batch", psnr, on_epoch=True)
+
+    def on_validation_epoch_end(self) -> None:
+        """Validation epoch end."""
+        psnr_ = self.reduce_running_psnr()
+        if psnr_ is not None:
+            self.log("val_psnr", psnr_, on_epoch=True, prog_bar=True)
+        else:
+            self.log("val_psnr", 0.0, on_epoch=True, prog_bar=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.
+        """
+        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
+        }
+
+    # TODO: find a way to move the following methods to a separate module
+    # TODO: this same operation is done in many other places, like in loss_func
+    # should we refactor LadderVAE so that it already outputs
+    # tuple(`mean`, `logvar`, `td_data`)?
+    def get_reconstructed_tensor(
+        self, model_outputs: tuple[Tensor, dict[str, Any]]
+    ) -> Tensor:
+        """Get the reconstructed tensor from the LVAE model outputs.
+
+        Parameters
+        ----------
+        model_outputs : tuple[Tensor, dict[str, Any]]
+            Model outputs. It is a tuple with a tensor representing the predicted mean
+            and (optionally) logvar, and the top-down data dictionary.
+
+        Returns
+        -------
+        Tensor
+            Reconstructed tensor, i.e., the predicted mean.
+        """
+        predictions, _ = model_outputs
+        if self.model.predict_logvar is None:
+            return predictions
+        elif self.model.predict_logvar == "pixelwise":
+            return predictions.chunk(2, dim=1)[0]
+
+    def compute_val_psnr(
+        self,
+        model_output: tuple[Tensor, dict[str, Any]],
+        target: Tensor,
+        psnr_func: Callable = scale_invariant_psnr,
+    ) -> list[float]:
+        """Compute the PSNR for the current validation batch.
+
+        Parameters
+        ----------
+        model_output : tuple[Tensor, dict[str, Any]]
+            Model output, a tuple with the predicted mean and (optionally) logvar,
+            and the top-down data dictionary.
+        target : Tensor
+            Target tensor.
+        psnr_func : Callable, optional
+            PSNR function to use, by default `scale_invariant_psnr`.
+
+        Returns
+        -------
+        list[float]
+            PSNR for each channel in the current batch.
+        """
+        out_channels = target.shape[1]
+
+        # get the reconstructed image
+        recons_img = self.get_reconstructed_tensor(model_output)
+
+        # update running psnr
+        for i in range(out_channels):
+            self.running_psnr[i].update(rec=recons_img[:, i], tar=target[:, i])
+
+        # compute psnr for each channel in the current batch
+        # TODO: this doesn't need do be a method of this class
+        # and hence can be moved to a separate module
+        return [
+            psnr_func(
+                gt=target[:, i].clone().detach().cpu().numpy(),
+                pred=recons_img[:, i].clone().detach().cpu().numpy(),
+            )
+            for i in range(out_channels)
+        ]
+
+    def reduce_running_psnr(self) -> Optional[float]:
+        """Reduce the running PSNR statistics and reset the running PSNR.
+
+        Returns
+        -------
+        Optional[float]
+            Running PSNR averaged over the different output channels.
+        """
+        psnr_arr = []  # type: ignore
+        for i in range(len(self.running_psnr)):
+            psnr = self.running_psnr[i].get()
+            if psnr is None:
+                psnr_arr = None  # type: ignore
+                break
+            psnr_arr.append(psnr.cpu().numpy())
+            self.running_psnr[i].reset()
+            # TODO: this line forces it to be a method of this class
+            # alternative is returning also the reset `running_psnr`
+        if psnr_arr is not None:
+            psnr = np.mean(psnr_arr)
+        return psnr
+
+
+# TODO: make this LVAE compatible (?)
 def create_careamics_module(
     algorithm: Union[SupportedAlgorithm, str],
     loss: Union[SupportedLoss, str],
@@ -212,8 +564,8 @@ def create_careamics_module(
     optimizer_parameters: Optional[dict] = None,
     lr_scheduler: Union[SupportedScheduler, str] = "ReduceLROnPlateau",
     lr_scheduler_parameters: Optional[dict] = None,
-) -> CAREamicsModule:
-    """Create a CAREamics Lithgning module.
+) -> Union[FCNModule, VAEModule]:
+    """Create a CAREamics Lightning module.
 
     This function exposes parameters used to create an AlgorithmModel instance,
     triggering parameters validation.
@@ -254,7 +606,7 @@ def create_careamics_module(
         optimizer_parameters = {}
     if model_parameters is None:
         model_parameters = {}
-    algorithm_configuration = {
+    algorithm_configuration: dict[str, Any] = {
         "algorithm": algorithm,
         "loss": loss,
         "optimizer": {
@@ -273,4 +625,10 @@ def create_careamics_module(
     algorithm_configuration["model"] = model_configuration
 
     # call the parent init using an AlgorithmModel instance
-    return CAREamicsModule(AlgorithmConfig(**algorithm_configuration))
+    algorithm_str = algorithm_configuration["algorithm"]
+    if algorithm_str in FCNAlgorithmConfig.get_compatible_algorithms():
+        return FCNModule(FCNAlgorithmConfig(**algorithm_configuration))
+    else:
+        raise NotImplementedError(
+            f"Model {algorithm_str} is not implemented or unknown."
+        )

careamics/lightning/predict_data_module.py

@@ -240,7 +240,7 @@ def create_predict_datamodule(
 ) -> PredictDataModule:
     """Create a CAREamics prediction Lightning datamodule.
 
-    This function is used to explicitely pass the parameters usually contained in an
+    This function is used to explicitly pass the parameters usually contained in an
     `inference_model` configuration.
 
     Since the lightning datamodule has no access to the model, make sure that the
@@ -268,7 +268,7 @@ def create_predict_datamodule(
     data_type : {"array", "tiff", "custom"}
         Data type, see `SupportedData` for available options.
     axes : str
-        Axes of the data, choosen among SCZYX.
+        Axes of the data, chosen among SCZYX.
     image_means : list of float
         Mean values for normalization, only used if Normalization is defined.
     image_stds : list of float

careamics/lightning/train_data_module.py

@@ -9,8 +9,8 @@ from numpy.typing import NDArray
 from torch.utils.data import DataLoader
 
 from careamics.config import DataConfig
-from careamics.config.data_model import TRANSFORMS_UNION
 from careamics.config.support import SupportedData
+from careamics.config.transformations import TransformModel
 from careamics.dataset.dataset_utils import (
     get_files_size,
     list_files,
@@ -472,7 +472,7 @@ def create_train_datamodule(
     axes: str,
     batch_size: int,
     val_data: Optional[Union[str, Path, NDArray]] = None,
-    transforms: Optional[list[TRANSFORMS_UNION]] = None,
+    transforms: Optional[list[TransformModel]] = None,
     train_target_data: Optional[Union[str, Path, NDArray]] = None,
     val_target_data: Optional[Union[str, Path, NDArray]] = None,
     read_source_func: Optional[Callable] = None,
@@ -487,7 +487,7 @@ def create_train_datamodule(
 ) -> TrainDataModule:
     """Create a TrainDataModule.
 
-    This function is used to explicitely pass the parameters usually contained in a
+    This function is used to explicitly pass the parameters usually contained in a
     `data_model` configuration to a TrainDataModule.
 
     Since the lightning datamodule has no access to the model, make sure that the
@@ -537,7 +537,7 @@ def create_train_datamodule(
     patch_size : list of int
         Patch size, 2D or 3D patch size.
     axes : str
-        Axes of the data, choosen amongst SCZYX.
+        Axes of the data, chosen amongst SCZYX.
     batch_size : int
         Batch size.
     val_data : pathlib.Path or str or numpy.ndarray, optional

careamics/losses/__init__.py

@@ -1,5 +1,15 @@
 """Losses module."""
 
-__all__ = ["loss_factory"]
+__all__ = [
+    "loss_factory",
+    "mae_loss",
+    "mse_loss",
+    "n2v_loss",
+    "denoisplit_loss",
+    "musplit_loss",
+    "denoisplit_musplit_loss",
+]
 
+from .fcn.losses import mae_loss, mse_loss, n2v_loss
 from .loss_factory import loss_factory
+from .lvae.losses import denoisplit_loss, denoisplit_musplit_loss, musplit_loss

careamics/losses/fcn/__init__.py

@@ -0,0 +1 @@
+"""FCN losses."""

careamics/losses/{losses.py → fcn/losses.py}

@@ -53,7 +53,7 @@ def n2v_loss(
     errors = (original_patches - manipulated_patches) ** 2
     # Average over pixels and batch
     loss = torch.sum(errors * masks) / torch.sum(masks)
-    return loss
+    return loss  # TODO change output to dict ?
 
 
 def mae_loss(samples: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:

careamics/losses/loss_factory.py

@@ -4,14 +4,114 @@ Loss factory module.
 This module contains a factory function for creating loss functions.
 """
 
-from typing import Callable, Union
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Callable, Literal, Optional, Union
+
+from torch import Tensor as tensor
 
 from ..config.support import SupportedLoss
-from .losses import mae_loss, mse_loss, n2v_loss
+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:
+    """Dataclass for FCN loss."""
+
+    # TODO check
+    prediction: tensor
+    targets: tensor
+    mask: tensor
+    current_epoch: int
+    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.")
 
 
-# TODO add tests
-# TODO add custom?
 def loss_factory(loss: Union[SupportedLoss, str]) -> Callable:
     """Return loss function.
 
@@ -42,8 +142,14 @@ def loss_factory(loss: Union[SupportedLoss, str]) -> Callable:
     elif loss == SupportedLoss.MSE:
         return mse_loss
 
-    # elif loss_type == SupportedLoss.DICE:
-    #     return dice_loss
+    elif loss == SupportedLoss.MUSPLIT:
+        return musplit_loss
+
+    elif loss == SupportedLoss.DENOISPLIT:
+        return denoisplit_loss
+
+    elif loss == SupportedLoss.DENOISPLIT_MUSPLIT:
+        return denoisplit_musplit_loss
 
     else:
         raise NotImplementedError(f"Loss {loss} is not yet supported.")

careamics/losses/lvae/__init__.py

@@ -0,0 +1 @@
+"""LVAE losses."""