careamics 0.0.19__py3-none-any.whl

careamics/lightning/callbacks/prediction_writer_callback/prediction_writer_callback.py

@@ -0,0 +1,234 @@
+"""Module containing `PredictionWriterCallback` class."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from pathlib import Path
+from typing import Any, Union
+
+from pytorch_lightning import LightningModule, Trainer
+from pytorch_lightning.callbacks import BasePredictionWriter
+from torch.utils.data import DataLoader
+
+from careamics.dataset import (
+    IterablePredDataset,
+    IterableTiledPredDataset,
+)
+from careamics.file_io import SupportedWriteType, WriteFunc
+from careamics.utils import get_logger
+
+from .write_strategy import WriteStrategy
+from .write_strategy_factory import create_write_strategy
+
+logger = get_logger(__name__)
+
+ValidPredDatasets = Union[IterablePredDataset, IterableTiledPredDataset]
+
+
+class PredictionWriterCallback(BasePredictionWriter):
+    """
+    A PyTorch Lightning callback to save predictions.
+
+    Parameters
+    ----------
+    write_strategy : WriteStrategy
+        A strategy for writing predictions.
+    dirpath : Path or str, default="predictions"
+        The path to the directory where prediction outputs will be saved. If
+        `dirpath` is not absolute it is assumed to be relative to current working
+        directory.
+
+    Attributes
+    ----------
+    write_strategy : WriteStrategy
+        A strategy for writing predictions.
+    dirpath : pathlib.Path, default="predictions"
+        The path to the directory where prediction outputs will be saved. If
+        `dirpath` is not absolute it is assumed to be relative to current working
+        directory.
+    writing_predictions : bool
+        If writing predictions is turned on or off.
+    """
+
+    def __init__(
+        self,
+        write_strategy: WriteStrategy,
+        dirpath: Union[Path, str] = "predictions",
+    ):
+        """
+        A PyTorch Lightning callback to save predictions.
+
+        Parameters
+        ----------
+        write_strategy : WriteStrategy
+            A strategy for writing predictions.
+        dirpath : pathlib.Path or str, default="predictions"
+            The path to the directory where prediction outputs will be saved. If
+            `dirpath` is not absolute it is assumed to be relative to current working
+            directory.
+        """
+        super().__init__(write_interval="batch")
+
+        # Toggle for CAREamist to switch off saving if desired
+        self.writing_predictions: bool = True
+
+        self.write_strategy: WriteStrategy = write_strategy
+
+        # forward declaration
+        self.dirpath: Path
+        # attribute initialisation
+        self._init_dirpath(dirpath)
+
+    @classmethod
+    def from_write_func_params(
+        cls,
+        write_type: SupportedWriteType,
+        tiled: bool,
+        write_func: WriteFunc | None = None,
+        write_extension: str | None = None,
+        write_func_kwargs: dict[str, Any] | None = None,
+        dirpath: Union[Path, str] = "predictions",
+    ) -> PredictionWriterCallback:  # TODO: change type hint to self (find out how)
+        """
+        Initialize a `PredictionWriterCallback` from write function parameters.
+
+        This will automatically create a `WriteStrategy` to be passed to the
+        initialization of `PredictionWriterCallback`.
+
+        Parameters
+        ----------
+        write_type : {"tiff", "custom"}
+            The data type to save as, includes custom.
+        tiled : bool
+            Whether the prediction will be tiled or not.
+        write_func : WriteFunc, optional
+            If a known `write_type` is selected this argument is ignored. For a custom
+            `write_type` a function to save the data must be passed. See notes below.
+        write_extension : str, optional
+            If a known `write_type` is selected this argument is ignored. For a custom
+            `write_type` an extension to save the data with must be passed.
+        write_func_kwargs : dict of {{str: any}}, optional
+            Additional keyword arguments to be passed to the save function.
+        dirpath : pathlib.Path or str, default="predictions"
+            The path to the directory where prediction outputs will be saved. If
+            `dirpath` is not absolute it is assumed to be relative to current working
+            directory.
+
+        Returns
+        -------
+        PredictionWriterCallback
+            Callback for writing predictions.
+        """
+        write_strategy = create_write_strategy(
+            write_type=write_type,
+            tiled=tiled,
+            write_func=write_func,
+            write_extension=write_extension,
+            write_func_kwargs=write_func_kwargs,
+        )
+        return cls(write_strategy=write_strategy, dirpath=dirpath)
+
+    def _init_dirpath(self, dirpath):
+        """
+        Initialize directory path. Should only be called from `__init__`.
+
+        Parameters
+        ----------
+        dirpath : pathlib.Path
+            See `__init__` description.
+        """
+        dirpath = Path(dirpath)
+        if not dirpath.is_absolute():
+            dirpath = Path.cwd() / dirpath
+            logger.warning(
+                "Prediction output directory is not absolute, absolute path assumed to"
+                f"be '{dirpath}'"
+            )
+        self.dirpath = dirpath
+
+    def setup(self, trainer: Trainer, pl_module: LightningModule, stage: str) -> None:
+        """
+        Create the prediction output directory when predict begins.
+
+        Called when fit, validate, test, predict, or tune begins.
+
+        Parameters
+        ----------
+        trainer : Trainer
+            PyTorch Lightning trainer.
+        pl_module : LightningModule
+            PyTorch Lightning module.
+        stage : str
+            Stage of training e.g. 'predict', 'fit', 'validate'.
+        """
+        super().setup(trainer, pl_module, stage)
+        if stage == "predict":
+            # make prediction output directory
+            logger.info("Making prediction output directory.")
+            self.dirpath.mkdir(parents=True, exist_ok=True)
+
+    def write_on_batch_end(
+        self,
+        trainer: Trainer,
+        pl_module: LightningModule,
+        prediction: Any,  # TODO: change to expected type
+        batch_indices: Sequence[int] | None,
+        batch: Any,  # TODO: change to expected type
+        batch_idx: int,
+        dataloader_idx: int,
+    ) -> None:
+        """
+        Write predictions at the end of a batch.
+
+        The method of prediction is determined by the attribute `write_strategy`.
+
+        Parameters
+        ----------
+        trainer : Trainer
+            PyTorch Lightning trainer.
+        pl_module : LightningModule
+            PyTorch Lightning module.
+        prediction : Any
+            Prediction outputs of `batch`.
+        batch_indices : sequence of Any, optional
+            Batch indices.
+        batch : Any
+            Input batch.
+        batch_idx : int
+            Batch index.
+        dataloader_idx : int
+            Dataloader index.
+        """
+        # if writing prediction is turned off
+        if not self.writing_predictions:
+            return
+
+        dataloaders: Union[DataLoader, list[DataLoader]] = trainer.predict_dataloaders
+        dataloader: DataLoader = (
+            dataloaders[dataloader_idx]
+            if isinstance(dataloaders, list)
+            else dataloaders
+        )
+        dataset: ValidPredDatasets = dataloader.dataset
+        if not (
+            isinstance(dataset, IterablePredDataset)
+            or isinstance(dataset, IterableTiledPredDataset)
+        ):
+            # Note: Error will be raised before here from the source type
+            # This is for extra redundancy of errors.
+            raise TypeError(
+                "Prediction dataset has to be `IterableTiledPredDataset` or "
+                "`IterablePredDataset`. Cannot be `InMemoryPredDataset` because "
+                "filenames are taken from the original file."
+            )
+
+        self.write_strategy.write_batch(
+            trainer=trainer,
+            pl_module=pl_module,
+            prediction=prediction,
+            batch_indices=batch_indices,
+            batch=batch,
+            batch_idx=batch_idx,
+            dataloader_idx=dataloader_idx,
+            dirpath=self.dirpath,
+        )

careamics/lightning/callbacks/prediction_writer_callback/write_strategy.py

@@ -0,0 +1,399 @@
+"""Module containing different strategies for writing predictions."""
+
+from collections.abc import Sequence
+from pathlib import Path
+from typing import Any, Protocol, Union
+
+import numpy as np
+from numpy.typing import NDArray
+from pytorch_lightning import LightningModule, Trainer
+from torch.utils.data import DataLoader
+
+from careamics.config.data.tile_information import TileInformation
+from careamics.dataset import IterablePredDataset, IterableTiledPredDataset
+from careamics.file_io import WriteFunc
+from careamics.prediction_utils import stitch_prediction_single
+
+from .file_path_utils import create_write_file_path, get_sample_file_path
+
+
+class WriteStrategy(Protocol):
+    """Protocol for write strategy classes."""
+
+    def write_batch(
+        self,
+        trainer: Trainer,
+        pl_module: LightningModule,
+        prediction: Any,  # TODO: change to expected type
+        batch_indices: Sequence[int] | None,
+        batch: Any,  # TODO: change to expected type
+        batch_idx: int,
+        dataloader_idx: int,
+        dirpath: Path,
+    ) -> None:
+        """
+        WriteStrategy subclasses must contain this function to write a batch.
+
+        Parameters
+        ----------
+        trainer : Trainer
+            PyTorch Lightning Trainer.
+        pl_module : LightningModule
+            PyTorch Lightning LightningModule.
+        prediction : Any
+            Predictions on `batch`.
+        batch_indices : sequence of int
+            Indices identifying the samples in the batch.
+        batch : Any
+            Input batch.
+        batch_idx : int
+            Batch index.
+        dataloader_idx : int
+            Dataloader index.
+        dirpath : Path
+            Path to directory to save predictions to.
+        """
+
+
+class CacheTiles(WriteStrategy):
+    """
+    A write strategy that will cache tiles.
+
+    Tiles are cached until a whole image is predicted on. Then the stitched
+    prediction is saved.
+
+    Parameters
+    ----------
+    write_func : WriteFunc
+        Function used to save predictions.
+    write_extension : str
+        Extension added to prediction file paths.
+    write_func_kwargs : dict of {str: Any}
+        Extra kwargs to pass to `write_func`.
+
+    Attributes
+    ----------
+    write_func : WriteFunc
+        Function used to save predictions.
+    write_extension : str
+        Extension added to prediction file paths.
+    write_func_kwargs : dict of {str: Any}
+        Extra kwargs to pass to `write_func`.
+    tile_cache : list of numpy.ndarray
+        Tiles cached for stitching prediction.
+    tile_info_cache : list of TileInformation
+        Cached tile information for stitching prediction.
+    """
+
+    def __init__(
+        self,
+        write_func: WriteFunc,
+        write_extension: str,
+        write_func_kwargs: dict[str, Any],
+    ) -> None:
+        """
+        A write strategy that will cache tiles.
+
+        Tiles are cached until a whole image is predicted on. Then the stitched
+        prediction is saved.
+
+        Parameters
+        ----------
+        write_func : WriteFunc
+            Function used to save predictions.
+        write_extension : str
+            Extension added to prediction file paths.
+        write_func_kwargs : dict of {str: Any}
+            Extra kwargs to pass to `write_func`.
+        """
+        super().__init__()
+
+        self.write_func: WriteFunc = write_func
+        self.write_extension: str = write_extension
+        self.write_func_kwargs: dict[str, Any] = write_func_kwargs
+
+        # where tiles will be cached until a whole image has been predicted
+        self.tile_cache: list[NDArray] = []
+        self.tile_info_cache: list[TileInformation] = []
+
+    @property
+    def last_tiles(self) -> list[bool]:
+        """
+        List of bool to determine whether each tile in the cache is the last tile.
+
+        Returns
+        -------
+        list of bool
+            Whether each tile in the tile cache is the last tile.
+        """
+        return [tile_info.last_tile for tile_info in self.tile_info_cache]
+
+    def write_batch(
+        self,
+        trainer: Trainer,
+        pl_module: LightningModule,
+        prediction: tuple[NDArray, list[TileInformation]],
+        batch_indices: Sequence[int] | None,
+        batch: tuple[NDArray, list[TileInformation]],
+        batch_idx: int,
+        dataloader_idx: int,
+        dirpath: Path,
+    ) -> None:
+        """
+        Cache tiles until the last tile is predicted; save the stitched prediction.
+
+        Parameters
+        ----------
+        trainer : Trainer
+            PyTorch Lightning Trainer.
+        pl_module : LightningModule
+            PyTorch Lightning LightningModule.
+        prediction : Any
+            Predictions on `batch`.
+        batch_indices : sequence of int
+            Indices identifying the samples in the batch.
+        batch : Any
+            Input batch.
+        batch_idx : int
+            Batch index.
+        dataloader_idx : int
+            Dataloader index.
+        dirpath : Path
+            Path to directory to save predictions to.
+        """
+        dataloaders: Union[DataLoader, list[DataLoader]] = trainer.predict_dataloaders
+        dataloader: DataLoader = (
+            dataloaders[dataloader_idx]
+            if isinstance(dataloaders, list)
+            else dataloaders
+        )
+        dataset: IterableTiledPredDataset = dataloader.dataset
+        if not isinstance(dataset, IterableTiledPredDataset):
+            raise TypeError("Prediction dataset is not `IterableTiledPredDataset`.")
+
+        # cache tiles (batches are split into single samples)
+        self.tile_cache.extend(np.split(prediction[0], prediction[0].shape[0]))
+        self.tile_info_cache.extend(prediction[1])
+
+        # save stitched prediction
+        if self._has_last_tile():
+
+            # get image tiles and remove them from the cache
+            tiles, tile_infos = self._get_image_tiles()
+            self._clear_cache()
+
+            # stitch prediction
+            prediction_image = stitch_prediction_single(
+                tiles=tiles, tile_infos=tile_infos
+            )
+
+            # write prediction
+            sample_id = tile_infos[0].sample_id  # need this to select correct file name
+            input_file_path = get_sample_file_path(dataset=dataset, sample_id=sample_id)
+            file_path = create_write_file_path(
+                dirpath=dirpath,
+                file_path=input_file_path,
+                write_extension=self.write_extension,
+            )
+            self.write_func(
+                file_path=file_path, img=prediction_image[0], **self.write_func_kwargs
+            )
+
+    def _has_last_tile(self) -> bool:
+        """
+        Whether a last tile is contained in the cached tiles.
+
+        Returns
+        -------
+        bool
+            Whether a last tile is contained in the cached tiles.
+        """
+        return any(self.last_tiles)
+
+    def _clear_cache(self) -> None:
+        """Remove the tiles in the cache up to the first last tile."""
+        index = self._last_tile_index()
+        self.tile_cache = self.tile_cache[index + 1 :]
+        self.tile_info_cache = self.tile_info_cache[index + 1 :]
+
+    def _last_tile_index(self) -> int:
+        """
+        Find the index of the last tile in the tile cache.
+
+        Returns
+        -------
+        int
+            Index of last tile.
+
+        Raises
+        ------
+        ValueError
+            If there is no last tile in the tile cache.
+        """
+        last_tiles = self.last_tiles
+        if not any(last_tiles):
+            raise ValueError("No last tile in the tile cache.")
+        index = np.where(last_tiles)[0][0]
+        return index
+
+    def _get_image_tiles(self) -> tuple[list[NDArray], list[TileInformation]]:
+        """
+        Get the tiles corresponding to a single image.
+
+        Returns
+        -------
+        tuple of (list of numpy.ndarray, list of TileInformation)
+            Tiles and tile information to stitch together a full image.
+        """
+        index = self._last_tile_index()
+        tiles = self.tile_cache[: index + 1]
+        tile_infos = self.tile_info_cache[: index + 1]
+        return tiles, tile_infos
+
+
+class WriteTilesZarr(WriteStrategy):
+    """Strategy to write tiles to Zarr file."""
+
+    def write_batch(
+        self,
+        trainer: Trainer,
+        pl_module: LightningModule,
+        prediction: Any,
+        batch_indices: Sequence[int] | None,
+        batch: Any,
+        batch_idx: int,
+        dataloader_idx: int,
+        dirpath: Path,
+    ) -> None:
+        """
+        Write tiles to zarr file.
+
+        Parameters
+        ----------
+        trainer : Trainer
+            PyTorch Lightning Trainer.
+        pl_module : LightningModule
+            PyTorch Lightning LightningModule.
+        prediction : Any
+            Predictions on `batch`.
+        batch_indices : sequence of int
+            Indices identifying the samples in the batch.
+        batch : Any
+            Input batch.
+        batch_idx : int
+            Batch index.
+        dataloader_idx : int
+            Dataloader index.
+        dirpath : Path
+            Path to directory to save predictions to.
+
+        Raises
+        ------
+        NotImplementedError
+        """
+        raise NotImplementedError
+
+
+class WriteImage(WriteStrategy):
+    """
+    A strategy for writing image predictions (i.e. un-tiled predictions).
+
+    Parameters
+    ----------
+    write_func : WriteFunc
+        Function used to save predictions.
+    write_extension : str
+        Extension added to prediction file paths.
+    write_func_kwargs : dict of {str: Any}
+        Extra kwargs to pass to `write_func`.
+
+    Attributes
+    ----------
+    write_func : WriteFunc
+        Function used to save predictions.
+    write_extension : str
+        Extension added to prediction file paths.
+    write_func_kwargs : dict of {str: Any}
+        Extra kwargs to pass to `write_func`.
+    """
+
+    def __init__(
+        self,
+        write_func: WriteFunc,
+        write_extension: str,
+        write_func_kwargs: dict[str, Any],
+    ) -> None:
+        """
+        A strategy for writing image predictions (i.e. un-tiled predictions).
+
+        Parameters
+        ----------
+        write_func : WriteFunc
+            Function used to save predictions.
+        write_extension : str
+            Extension added to prediction file paths.
+        write_func_kwargs : dict of {str: Any}
+            Extra kwargs to pass to `write_func`.
+        """
+        super().__init__()
+
+        self.write_func: WriteFunc = write_func
+        self.write_extension: str = write_extension
+        self.write_func_kwargs: dict[str, Any] = write_func_kwargs
+
+    def write_batch(
+        self,
+        trainer: Trainer,
+        pl_module: LightningModule,
+        prediction: NDArray,
+        batch_indices: Sequence[int] | None,
+        batch: NDArray,
+        batch_idx: int,
+        dataloader_idx: int,
+        dirpath: Path,
+    ) -> None:
+        """
+        Save full images.
+
+        Parameters
+        ----------
+        trainer : Trainer
+            PyTorch Lightning Trainer.
+        pl_module : LightningModule
+            PyTorch Lightning LightningModule.
+        prediction : Any
+            Predictions on `batch`.
+        batch_indices : sequence of int
+            Indices identifying the samples in the batch.
+        batch : Any
+            Input batch.
+        batch_idx : int
+            Batch index.
+        dataloader_idx : int
+            Dataloader index.
+        dirpath : Path
+            Path to directory to save predictions to.
+
+        Raises
+        ------
+        TypeError
+            If trainer prediction dataset is not `IterablePredDataset`.
+        """
+        dls: Union[DataLoader, list[DataLoader]] = trainer.predict_dataloaders
+        dl: DataLoader = dls[dataloader_idx] if isinstance(dls, list) else dls
+        ds: IterablePredDataset = dl.dataset
+        if not isinstance(ds, IterablePredDataset):
+            raise TypeError("Prediction dataset is not `IterablePredDataset`.")
+
+        for i in range(prediction.shape[0]):
+            prediction_image = prediction[0]
+            sample_id = batch_idx * dl.batch_size + i
+            input_file_path = get_sample_file_path(dataset=ds, sample_id=sample_id)
+            file_path = create_write_file_path(
+                dirpath=dirpath,
+                file_path=input_file_path,
+                write_extension=self.write_extension,
+            )
+            self.write_func(
+                file_path=file_path, img=prediction_image, **self.write_func_kwargs
+            )