careamics 0.0.11__py3-none-any.whl → 0.0.13__py3-none-any.whl

careamics/lightning/dataset_ng/lightning_modules/care_module.py

@@ -0,0 +1,97 @@
+"""CARE Lightning DataModule."""
+
+from collections.abc import Callable
+from typing import Any, Union
+
+from careamics.config.algorithms.care_algorithm_model import CAREAlgorithm
+from careamics.config.algorithms.n2n_algorithm_model import N2NAlgorithm
+from careamics.config.support import SupportedLoss
+from careamics.dataset_ng.dataset import ImageRegionData
+from careamics.losses import mae_loss, mse_loss
+from careamics.utils.logging import get_logger
+
+from .unet_module import UnetModule
+
+logger = get_logger(__name__)
+
+
+class CAREModule(UnetModule):
+    """CAREamics PyTorch Lightning module for CARE algorithm.
+
+    Parameters
+    ----------
+    algorithm_config : CAREAlgorithm or dict
+        Configuration for the CARE algorithm, either as a CAREAlgorithm instance or a
+        dictionary.
+    """
+
+    def __init__(self, algorithm_config: Union[CAREAlgorithm, dict]) -> None:
+        """Instantiate CARE DataModule.
+
+        Parameters
+        ----------
+        algorithm_config : CAREAlgorithm or dict
+            Configuration for the CARE algorithm, either as a CAREAlgorithm instance or
+            a dictionary.
+        """
+        super().__init__(algorithm_config)
+        assert isinstance(
+            algorithm_config, CAREAlgorithm | N2NAlgorithm
+        ), "algorithm_config must be a CAREAlgorithm or a N2NAlgorithm"
+        loss = algorithm_config.loss
+        if loss == SupportedLoss.MAE:
+            self.loss_func: Callable = mae_loss
+        elif loss == SupportedLoss.MSE:
+            self.loss_func = mse_loss
+        else:
+            raise ValueError(f"Unsupported loss for Care: {loss}")
+
+    def training_step(
+        self,
+        batch: tuple[ImageRegionData, ImageRegionData],
+        batch_idx: Any,
+    ) -> Any:
+        """Training step for CARE module.
+
+        Parameters
+        ----------
+        batch : (ImageRegionData, ImageRegionData)
+            A tuple containing the input data and the target data.
+        batch_idx : Any
+            The index of the current batch in the training loop.
+
+        Returns
+        -------
+        Any
+            The loss value computed for the current batch.
+        """
+        # TODO: add validation to determine if target is initialized
+        x, target = batch[0], batch[1]
+
+        prediction = self.model(x.data)
+        loss = self.loss_func(prediction, target.data)
+
+        self._log_training_stats(loss, batch_size=x.data.shape[0])
+
+        return loss
+
+    def validation_step(
+        self,
+        batch: tuple[ImageRegionData, ImageRegionData],
+        batch_idx: Any,
+    ) -> None:
+        """Validation step for CARE module.
+
+        Parameters
+        ----------
+        batch : (ImageRegionData, ImageRegionData)
+            A tuple containing the input data and the target data.
+        batch_idx : Any
+            The index of the current batch in the training loop.
+        """
+        x, target = batch[0], batch[1]
+
+        prediction = self.model(x.data)
+        val_loss = self.loss_func(prediction, target.data)
+        self.metrics(prediction, target.data)
+        self._log_validation_stats(val_loss, batch_size=x.data.shape[0])

careamics/lightning/dataset_ng/lightning_modules/n2v_module.py

@@ -0,0 +1,106 @@
+"""Noise2Void Lightning DataModule."""
+
+from typing import Any, Union
+
+from careamics.config import (
+    N2VAlgorithm,
+)
+from careamics.dataset_ng.dataset import ImageRegionData
+from careamics.losses import n2v_loss
+from careamics.transforms import N2VManipulateTorch
+from careamics.utils.logging import get_logger
+
+from .unet_module import UnetModule
+
+logger = get_logger(__name__)
+
+
+class N2VModule(UnetModule):
+    """CAREamics PyTorch Lightning module for N2V algorithm.
+
+    Parameters
+    ----------
+    algorithm_config : N2VAlgorithm or dict
+        Configuration for the N2V algorithm, either as an N2VAlgorithm instance or a
+        dictionary.
+    """
+
+    def __init__(self, algorithm_config: Union[N2VAlgorithm, dict]) -> None:
+        """Instantiate N2V DataModule.
+
+        Parameters
+        ----------
+        algorithm_config : N2VAlgorithm or dict
+            Configuration for the N2V algorithm, either as an N2VAlgorithm instance or a
+            dictionary.
+        """
+        super().__init__(algorithm_config)
+
+        assert isinstance(
+            algorithm_config, N2VAlgorithm
+        ), "algorithm_config must be a N2VAlgorithm"
+
+        self.n2v_manipulate = N2VManipulateTorch(
+            n2v_manipulate_config=algorithm_config.n2v_config
+        )
+        self.loss_func = n2v_loss
+
+    def _load_best_checkpoint(self) -> None:
+        """Load the best checkpoint for N2V model."""
+        logger.warning(
+            "Loading best checkpoint for N2V model. Note that for N2V, "
+            "the checkpoint with the best validation metrics may not necessarily "
+            "have the best denoising performance."
+        )
+        super()._load_best_checkpoint()
+
+    def training_step(
+        self,
+        batch: Union[tuple[ImageRegionData], tuple[ImageRegionData, ImageRegionData]],
+        batch_idx: Any,
+    ) -> Any:
+        """Training step for N2V model.
+
+        Parameters
+        ----------
+        batch : ImageRegionData or (ImageRegionData, ImageRegionData)
+            A tuple containing the input data and the target data.
+        batch_idx : Any
+            The index of the current batch in the training loop.
+
+        Returns
+        -------
+        Any
+            The loss value for the current training step.
+        """
+        x = batch[0]
+        x_masked, x_original, mask = self.n2v_manipulate(x.data)
+        prediction = self.model(x_masked)
+        loss = self.loss_func(prediction, x_original, mask)
+
+        self._log_training_stats(loss, batch_size=x.data.shape[0])
+
+        return loss
+
+    def validation_step(
+        self,
+        batch: Union[tuple[ImageRegionData], tuple[ImageRegionData, ImageRegionData]],
+        batch_idx: Any,
+    ) -> None:
+        """Validation step for N2V model.
+
+        Parameters
+        ----------
+        batch : ImageRegionData or (ImageRegionData, ImageRegionData)
+            A tuple containing the input data and the target data.
+        batch_idx : Any
+            The index of the current batch in the training loop.
+        """
+        x = batch[0]
+
+        x_masked, x_original, mask = self.n2v_manipulate(x.data)
+        prediction = self.model(x_masked)
+
+        val_loss = self.loss_func(prediction, x_original, mask)
+        self.metrics(prediction, x_original)
+        self._log_validation_stats(val_loss, batch_size=x.data.shape[0])

careamics/lightning/dataset_ng/lightning_modules/unet_module.py

@@ -0,0 +1,212 @@
+"""Generic UNet Lightning DataModule."""
+
+from typing import Any, Union
+
+import pytorch_lightning as L
+import torch
+from torch import nn
+from torchmetrics import MetricCollection
+from torchmetrics.image import PeakSignalNoiseRatio
+
+from careamics.config import algorithm_factory
+from careamics.config.algorithms import CAREAlgorithm, N2NAlgorithm, N2VAlgorithm
+from careamics.dataset_ng.dataset import ImageRegionData
+from careamics.models.unet import UNet
+from careamics.transforms import Denormalize
+from careamics.utils.logging import get_logger
+from careamics.utils.torch_utils import get_optimizer, get_scheduler
+
+logger = get_logger(__name__)
+
+
+class UnetModule(L.LightningModule):
+    """CAREamics PyTorch Lightning module for UNet based algorithms.
+
+    Parameters
+    ----------
+    algorithm_config : CAREAlgorithm, N2VAlgorithm, N2NAlgorithm, or dict
+        Configuration for the algorithm, either as an instance of a specific algorithm
+        class or a dictionary that can be converted to an algorithm instance.
+    """
+
+    def __init__(
+        self, algorithm_config: Union[CAREAlgorithm, N2VAlgorithm, N2NAlgorithm, dict]
+    ) -> None:
+        """Instantiate UNet DataModule.
+
+        Parameters
+        ----------
+        algorithm_config : CAREAlgorithm, N2VAlgorithm, N2NAlgorithm, or dict
+            Configuration for the algorithm, either as an instance of a specific
+            algorithm class or a dictionary that can be converted to an algorithm
+            instance.
+        """
+        super().__init__()
+
+        if isinstance(algorithm_config, dict):
+            algorithm_config = algorithm_factory(algorithm_config)
+
+        self.config = algorithm_config
+        self.model: nn.Module = UNet(**algorithm_config.model.model_dump())
+
+        self._best_checkpoint_loaded = False
+
+        # TODO: how to support metric evaluation better
+        self.metrics = MetricCollection(PeakSignalNoiseRatio())
+
+    def forward(self, x: Any) -> Any:
+        """Default forward method.
+
+        Parameters
+        ----------
+        x : Any
+            Input data.
+
+        Returns
+        -------
+        Any
+            Output from the model.
+        """
+        return self.model(x)
+
+    def _log_training_stats(self, loss: Any, batch_size: Any) -> None:
+        """Log training statistics.
+
+        Parameters
+        ----------
+        loss : Any
+            The loss value for the current training step.
+        batch_size : Any
+            The size of the batch used in the current training step.
+        """
+        self.log(
+            "train_loss",
+            loss,
+            on_step=True,
+            on_epoch=True,
+            prog_bar=True,
+            logger=True,
+            batch_size=batch_size,
+        )
+
+        optimizer = self.optimizers()
+        if isinstance(optimizer, list):
+            current_lr = optimizer[0].param_groups[0]["lr"]
+        else:
+            current_lr = optimizer.param_groups[0]["lr"]
+        self.log(
+            "learning_rate",
+            current_lr,
+            on_step=False,
+            on_epoch=True,
+            logger=True,
+            batch_size=batch_size,
+        )
+
+    def _log_validation_stats(self, loss: Any, batch_size: Any) -> None:
+        """Log validation statistics.
+
+        Parameters
+        ----------
+        loss : Any
+            The loss value for the current validation step.
+        batch_size : Any
+            The size of the batch used in the current validation step.
+        """
+        self.log(
+            "val_loss",
+            loss,
+            on_step=False,
+            on_epoch=True,
+            prog_bar=True,
+            logger=True,
+            batch_size=batch_size,
+        )
+        self.log_dict(self.metrics, on_step=False, on_epoch=True, batch_size=batch_size)
+
+    def _load_best_checkpoint(self) -> None:
+        """Load the best checkpoint from the trainer's checkpoint callback."""
+        if (
+            not hasattr(self.trainer, "checkpoint_callback")
+            or self.trainer.checkpoint_callback is None
+        ):
+            logger.warning("No checkpoint callback found, cannot load best checkpoint.")
+            return
+
+        best_model_path = self.trainer.checkpoint_callback.best_model_path
+        if best_model_path and best_model_path != "":
+            logger.info(f"Loading best checkpoint from: {best_model_path}")
+            model_state = torch.load(best_model_path, weights_only=True)["state_dict"]
+            self.load_state_dict(model_state)
+        else:
+            logger.warning("No best checkpoint found.")
+
+    def predict_step(
+        self,
+        batch: Union[tuple[ImageRegionData], tuple[ImageRegionData, ImageRegionData]],
+        batch_idx: Any,
+        load_best_checkpoint=False,
+    ) -> Any:
+        """Default predict step.
+
+        Parameters
+        ----------
+        batch : ImageRegionData or (ImageRegionData, ImageRegionData)
+            A tuple containing the input data and optionally the target data.
+        batch_idx : Any
+            The index of the current batch in the prediction loop.
+        load_best_checkpoint : bool, default=False
+            Whether to load the best checkpoint before making predictions.
+
+        Returns
+        -------
+        Any
+            The output batch containing the predictions.
+        """
+        if self._best_checkpoint_loaded is False and load_best_checkpoint:
+            self._load_best_checkpoint()
+            self._best_checkpoint_loaded = True
+
+        x = batch[0]
+        # TODO: add TTA
+        prediction = self.model(x.data).cpu().numpy()
+
+        means = self._trainer.datamodule.stats.means
+        stds = self._trainer.datamodule.stats.stds
+        denormalize = Denormalize(
+            image_means=means,
+            image_stds=stds,
+        )
+        denormalized_output = denormalize(prediction)
+
+        output_batch = ImageRegionData(
+            data=denormalized_output,
+            source=x.source,
+            data_shape=x.data_shape,
+            dtype=x.dtype,
+            axes=x.axes,
+            region_spec=x.region_spec,
+        )
+        return output_batch
+
+    def configure_optimizers(self) -> Any:
+        """Configure optimizers.
+
+        Returns
+        -------
+        Any
+            A dictionary containing the optimizer and learning rate scheduler.
+        """
+        optimizer_func = get_optimizer(self.config.optimizer.name)
+        optimizer = optimizer_func(
+            self.model.parameters(), **self.config.optimizer.parameters
+        )
+
+        scheduler_func = get_scheduler(self.config.lr_scheduler.name)
+        scheduler = scheduler_func(optimizer, **self.config.lr_scheduler.parameters)
+
+        return {
+            "optimizer": optimizer,
+            "lr_scheduler": scheduler,
+            "monitor": "val_loss",  # otherwise triggers MisconfigurationException
+        }

careamics/lightning/lightning_module.py

@@ -1,6 +1,7 @@
 """CAREamics Lightning module."""
 
-from typing import Any, Callable, Literal, Optional, Union
+from collections.abc import Callable
+from typing import Any, Literal, Optional, Union
 
 import numpy as np
 import pytorch_lightning as L
@@ -148,6 +149,9 @@ class FCNModule(L.LightningModule):
         self.log(
             "train_loss", loss, on_step=True, on_epoch=True, prog_bar=True, logger=True
         )
+        optimizer = self.optimizers()
+        current_lr = optimizer.param_groups[0]["lr"]
+        self.log("learning_rate", current_lr, on_step=False, on_epoch=True, logger=True)
         return loss
 
     def validation_step(self, batch: Tensor, batch_idx: Any) -> None:

careamics/lightning/predict_data_module.py

@@ -1,7 +1,8 @@
 """Prediction Lightning data modules."""
 
+from collections.abc import Callable
 from pathlib import Path
-from typing import Any, Callable, Literal, Optional, Union
+from typing import Any, Literal, Optional, Union
 
 import numpy as np
 import pytorch_lightning as L

careamics/lightning/train_data_module.py

@@ -1,7 +1,8 @@
 """Training and validation Lightning data modules."""
 
+from collections.abc import Callable
 from pathlib import Path
-from typing import Any, Callable, Literal, Optional, Union
+from typing import Any, Literal, Optional, Union
 
 import numpy as np
 import pytorch_lightning as L

careamics/losses/loss_factory.py

@@ -6,8 +6,9 @@ This module contains a factory function for creating loss functions.
 
 from __future__ import annotations
 
+from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Callable, Union
+from typing import Union
 
 from torch import Tensor as tensor
 

careamics/lvae_training/dataset/__init__.py

@@ -1,14 +1,19 @@
-from .multich_dataset import MultiChDloader
+from .config import DatasetConfig
 from .lc_dataset import LCMultiChDloader
+from .ms_dataset_ref import MultiChDloaderRef
+from .multich_dataset import MultiChDloader
+from .multicrop_dset import MultiCropDset
 from .multifile_dataset import MultiFileDset
-from .config import DatasetConfig
-from .types import DataType, DataSplitType, TilingMode
+from .types import DataSplitType, DataType, TilingMode
 
 __all__ = [
     "DatasetConfig",
     "MultiChDloader",
     "LCMultiChDloader",
     "MultiFileDset",
+    "MultiCropDset",
+    "MultiChDloaderRef",
+    "LCMultiChDloaderRef",
     "DataType",
     "DataSplitType",
     "TilingMode",

careamics/lvae_training/dataset/config.py

@@ -1,4 +1,4 @@
-from typing import Any, Optional
+from typing import Any, Optional, Union
 
 from pydantic import BaseModel, ConfigDict
 
@@ -43,7 +43,7 @@ class DatasetConfig(BaseModel):
     image_size: tuple  # TODO: revisit, new model_config uses tuple
     """Size of one patch of data"""
 
-    grid_size: Optional[int] = None
+    grid_size: Optional[Union[int, tuple[int, int, int]]] = None
     """Frame is divided into square grids of this size. A patch centered on a grid 
     having size `image_size` is returned. Grid size not used in training,
     used only during val / test, grid size controls the overlap of the patches"""
@@ -82,7 +82,7 @@ class DatasetConfig(BaseModel):
     # TODO: why is this not used?
     enable_rotation_aug: Optional[bool] = False
 
-    max_val: Optional[float] = None
+    max_val: Optional[Union[float, tuple]] = None
     """Maximum data in the dataset. Is calculated for train split, and should be 
     externally set for val and test splits."""