careamics 0.0.19__py3-none-any.whl

careamics/lightning/dataset_ng/lightning_modules/unet_module.py

@@ -0,0 +1,221 @@
+"""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,
+    PN2VAlgorithm,
+)
+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, PN2VAlgorithm, 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,
+            additional_metadata={},
+        )
+        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/dataset_ng/prediction/__init__.py

@@ -0,0 +1,16 @@
+"""Prediction utilities for the NG Dataset."""
+
+__all__ = [
+    "combine_samples",
+    "convert_prediction",
+    "decollate_image_region_data",
+    "stitch_prediction",
+    "stitch_single_prediction",
+]
+
+from .convert_prediction import (
+    combine_samples,
+    convert_prediction,
+    decollate_image_region_data,
+)
+from .stitch_prediction import stitch_prediction, stitch_single_prediction

careamics/lightning/dataset_ng/prediction/convert_prediction.py

@@ -0,0 +1,198 @@
+"""Module containing functions to convert prediction outputs to desired form."""
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+from numpy.typing import NDArray
+
+from careamics.dataset_ng.dataset import ImageRegionData
+
+from .stitch_prediction import group_tiles_by_key, stitch_prediction
+
+if TYPE_CHECKING:
+    from torch import Tensor
+
+
+def _decollate_batch_dict(
+    batched_dict: "dict[str, list | Tensor]",
+    index: int,
+) -> dict[str, int | tuple[int, ...]]:
+    """
+    Decollate element `index` from a batched_dict.
+
+    This method is only compatible with integer elements.
+
+    Parameters
+    ----------
+    batched_dict : dict of {str: list or Tensor}
+        Batch dictionary where each value is a list of elements of length B or a
+        Tensor of shape (B,).
+    index : int
+        Index of the element to extract.
+
+    Returns
+    -------
+    dict of {str: int | tuple[int, ...]}
+        Dictionary of the `index` element in the collated batch.
+    """
+    item_dict = {
+        key: (
+            # cast to int otherwise we have Tensor scalars
+            # TODO for additional types (e.g. axes in additional_metadata), we will need
+            # to handle it differently
+            tuple(int(value[idx][index]) for idx in range(len(value)))
+            if isinstance(value, list)
+            else int(value[index])
+        )  # handles tensor (1D) vs list of 1D tensors (2D)
+        for key, value in batched_dict.items()
+    }
+
+    return item_dict
+
+
+def decollate_image_region_data(
+    batch: ImageRegionData,
+) -> list[ImageRegionData]:
+    """
+    Decollate a batch of `ImageRegionData` into a list of `ImageRegionData`.
+
+    Input batch has the following structure:
+    - data: (B, C, (Z), Y, X) numpy.ndarray
+    - source: sequence of str, length B
+    - data_shape: sequence of tuple of int, each tuple being of length B
+    - dtype: list of numpy.dtype, length B
+    - axes: list of str, length B
+    - region_spec: dict of {str: sequence}, each sequence being of length B
+    - additional_metadata: dict of {str: Any}, each sequence being of length B
+
+    Parameters
+    ----------
+    batch : ImageRegionData
+        Batch of `ImageRegionData`.
+
+    Returns
+    -------
+    list of ImageRegionData
+        List of `ImageRegionData`.
+    """
+    batch_size = batch.data.shape[0]
+    decollated: list[ImageRegionData] = []
+    for i in range(batch_size):
+        # unpack region spec irrespective of whether it is a PatchSpecs or TileSpecs
+        region_spec = _decollate_batch_dict(batch.region_spec, i)
+
+        # handle additional metadata
+        # currently only zarr chunks and shards may be stored there, as tuples.
+        # TODO if additional metadata becomes used for anything else, this function
+        # call may not be appropriate anymore.
+        additional_metadata = _decollate_batch_dict(batch.additional_metadata, i)
+
+        # data shape
+        assert isinstance(batch.data_shape, list)
+        data_shape = tuple(int(dim[i]) for dim in batch.data_shape)
+
+        image_region = ImageRegionData(
+            data=batch.data[i],  # discard batch dimension
+            source=batch.source[i],
+            dtype=batch.dtype[i],
+            data_shape=data_shape,
+            axes=batch.axes[i],
+            region_spec=region_spec,  # type: ignore
+            additional_metadata=additional_metadata,
+        )
+        decollated.append(image_region)
+
+    return decollated
+
+
+def combine_samples(
+    predictions: list[ImageRegionData],
+) -> tuple[list[NDArray], list[str]]:
+    """
+    Combine predictions by `data_idx`.
+
+    Images are first grouped by their `data_idx` found in their `region_spec`, then
+    sorted by ascending `sample_idx` before being stacked along the `S` dimension.
+
+    Parameters
+    ----------
+    predictions : list of ImageRegionData
+        List of `ImageRegionData`.
+
+    Returns
+    -------
+    list of numpy.ndarray
+        List of combined predictions, one per unique `data_idx`.
+    list of str
+        List of sources, one per unique `data_idx`.
+    """
+    # group predictions by data idx
+    grouped_prediction: dict[int, list[ImageRegionData]] = group_tiles_by_key(
+        predictions, key="data_idx"
+    )
+
+    # sort predictions by sample idx
+    combined_predictions: list[NDArray] = []
+    combined_sources: list[str] = []
+    for data_idx in sorted(grouped_prediction.keys()):
+        image_regions = grouped_prediction[data_idx]
+        combined_sources.append(image_regions[0].source)
+
+        # sort by sample idx
+        image_regions.sort(key=lambda x: x.region_spec["sample_idx"])
+
+        # remove singleton dims and stack along S axis
+        combined_data = np.stack([img.data.squeeze() for img in image_regions], axis=0)
+        combined_predictions.append(combined_data)
+
+    return combined_predictions, combined_sources
+
+
+def convert_prediction(
+    predictions: list[ImageRegionData],
+    tiled: bool,
+) -> tuple[list[NDArray], list[str]]:
+    """
+    Convert the Lightning trainer outputs to the desired form.
+
+    This method allows decollating batches and stitching back together tiled
+    predictions.
+
+    If the `source` of all predictions is "array" (see `InMemoryImageStack`), then the
+    returned sources list will be empty.
+
+    Parameters
+    ----------
+    predictions : list[ImageRegionData]
+        Output from `Trainer.predict`, list of batches.
+    tiled : bool
+        Whether the predictions are tiled.
+
+    Returns
+    -------
+    list of numpy.ndarray
+        List of arrays with the axes SC(Z)YX.
+    list of str
+        List of sources, one per output or empty if all equal to `array`.
+    """
+    # decollate batches
+    decollated_predictions: list[ImageRegionData] = []
+    for batch in predictions:
+        decollated_batch = decollate_image_region_data(batch)
+        decollated_predictions.extend(decollated_batch)
+
+    if not tiled and "total_tiles" in decollated_predictions[0].region_spec:
+        raise ValueError(
+            "Predictions contain `total_tiles` in region_spec but `tiled` is set to "
+            "False."
+        )
+
+    if tiled:
+        predictions_output, sources = stitch_prediction(decollated_predictions)
+    else:
+        predictions_output, sources = combine_samples(decollated_predictions)
+
+    if set(sources) == {"array"}:
+        sources = []
+
+    return predictions_output, sources

careamics/lightning/dataset_ng/prediction/stitch_prediction.py

@@ -0,0 +1,171 @@
+"""Tiled prediction stitching utilities."""
+
+import builtins
+from collections import defaultdict
+from typing import Literal
+
+import numpy as np
+from numpy.typing import NDArray
+
+from careamics.dataset_ng.dataset import ImageRegionData
+from careamics.dataset_ng.patching_strategies import TileSpecs
+
+
+def group_tiles_by_key(
+    tiles: list[ImageRegionData], key: Literal["data_idx", "sample_idx"]
+) -> dict[int, list[ImageRegionData]]:
+    """
+    Sort tiles by key.
+
+    Parameters
+    ----------
+    tiles : list of ImageRegionData
+        List of tiles to sort.
+    key : {'data_idx', 'sample_idx'}
+        Key to group tiles by.
+
+    Returns
+    -------
+    {int: list of ImageRegionData}
+        Dictionary mapping data indices to lists of tiles.
+    """
+    sorted_tiles: dict[int, list[ImageRegionData]] = defaultdict(list)
+    for tile in tiles:
+        key_value = tile.region_spec[key]
+        sorted_tiles[key_value].append(tile)
+    return sorted_tiles
+
+
+def stitch_prediction(
+    tiles: list[ImageRegionData],
+) -> tuple[list[NDArray], list[str]]:
+    """
+    Stitch tiles back together to form full images.
+
+    Tiles are of dimensions SC(Z)YX, where C is the number of channels and can be a
+    singleton dimension.
+
+    Parameters
+    ----------
+    tiles : list of ImageRegionData
+        Cropped tiles and their respective stitching coordinates. Can contain tiles
+        from multiple images.
+
+    Returns
+    -------
+    list of numpy.ndarray
+        Full images, may be a single image.
+    list of str
+        List of sources, one per output.
+    """
+    # sort tiles by data index
+    grouped_tiles: dict[int, list[ImageRegionData]] = group_tiles_by_key(
+        tiles, key="data_idx"
+    )
+
+    # stitch each image separately
+    image_predictions: list[NDArray] = []
+    image_sources: list[str] = []
+    for data_idx in sorted(grouped_tiles.keys()):
+        image_predictions.append(stitch_single_prediction(grouped_tiles[data_idx]))
+        image_sources.append(grouped_tiles[data_idx][0].source)
+
+    return image_predictions, image_sources
+
+
+def stitch_single_prediction(
+    tiles: list[ImageRegionData],
+) -> NDArray:
+    """
+    Stitch tiles back together to form a full image.
+
+    Tiles are of dimensions C(Z)YX, where C is the number of channels and can be a
+    singleton dimension.
+
+    Parameters
+    ----------
+    tiles : list of ImageRegionData
+        Cropped tiles and their respective stitching coordinates.
+
+    Returns
+    -------
+    numpy.ndarray
+        Full image, with dimensions SC(Z)YX.
+    """
+    data_shape = tiles[0].data_shape
+    predicted_image = np.zeros(data_shape, dtype=np.float32)
+
+    if "S" in tiles[0].axes:
+        tiles_by_sample = group_tiles_by_key(tiles, key="sample_idx")
+        for sample_idx in tiles_by_sample.keys():
+            sample_tiles = tiles_by_sample[sample_idx]
+            stitched_sample = stitch_single_sample(sample_tiles)
+
+            # compute sample slice
+            sample_slice = slice(
+                sample_idx,
+                sample_idx + 1,
+            )
+
+            # insert stitched sample into predicted image
+            predicted_image[sample_slice] = stitched_sample.astype(np.float32)
+    else:
+        # stitch as a single sample
+        # predicted_image has singleton sample dimension
+        predicted_image[0] = stitch_single_sample(tiles)
+
+    return predicted_image
+
+
+def stitch_single_sample(
+    tiles: list[ImageRegionData],
+) -> NDArray:
+    """
+    Stitch tiles back together to form a full sample.
+
+    Tiles are of dimensions C(Z)YX, where C is the number of channels and can be a
+    singleton dimension.
+
+    Parameters
+    ----------
+    tiles : list of ImageRegionData
+        Cropped tiles and their respective stitching coordinates.
+
+    Returns
+    -------
+    numpy.ndarray
+        Full sample, with dimensions C(Z)YX.
+    """
+    data_shape = tiles[0].data_shape  # SC(Z)YX
+    predicted_sample = np.zeros(data_shape[1:], dtype=np.float32)
+
+    for tile in tiles:
+        # compute crop coordinates and stitiching coordinates
+        tile_spec: TileSpecs = tile.region_spec  # type: ignore
+        crop_coords = tile_spec["crop_coords"]
+        crop_size = tile_spec["crop_size"]
+        stitch_coords = tile_spec["stitch_coords"]
+
+        crop_slices: tuple[builtins.ellipsis | slice, ...] = (
+            ...,
+            *[
+                slice(start, start + length)
+                for start, length in zip(crop_coords, crop_size, strict=True)
+            ],
+        )
+
+        stitch_slices: tuple[builtins.ellipsis | slice, ...] = (
+            ...,
+            *[
+                slice(start, start + length)
+                for start, length in zip(stitch_coords, crop_size, strict=True)
+            ],
+        )
+
+        # crop predited tile according to overlap coordinates
+        cropped_tile = tile.data[crop_slices]
+
+        # insert cropped tile into predicted image
+        predicted_sample[stitch_slices] = cropped_tile.astype(np.float32)
+
+    return predicted_sample