careamics 0.0.19__py3-none-any.whl

careamics/lightning/dataset_ng/callbacks/prediction_writer/prediction_writer_callback.py

@@ -0,0 +1,219 @@
+"""Module containing `PredictionWriterCallback` class."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from pathlib import Path
+from typing import Any
+
+from pytorch_lightning import LightningModule, Trainer
+from pytorch_lightning.callbacks import BasePredictionWriter
+
+from careamics.dataset_ng.dataset import ImageRegionData
+from careamics.file_io.write.get_func import SupportedWriteType, WriteFunc
+from careamics.lightning.dataset_ng.prediction import decollate_image_region_data
+from careamics.utils import get_logger
+
+from .write_strategy import WriteStrategy
+from .write_strategy_factory import create_write_strategy
+
+logger = get_logger(__name__)
+
+
+class PredictionWriterCallback(BasePredictionWriter):
+    """
+    PyTorch Lightning callback to save predictions.
+
+    A `WriteStrategy` must be provided at instantiation or later via
+    `set_writing_strategy`.
+
+    Parameters
+    ----------
+    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.
+    write_strategy : WriteStrategy or None, default=None
+        A strategy for writing predictions.
+
+    Attributes
+    ----------
+    writing_predictions : bool
+        If writing predictions is turned on or off.
+    dirpath : pathlib.Path, default=""
+        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.
+    write_strategy : WriteStrategy or None
+            A strategy for writing predictions.
+    """
+
+    def __init__(
+        self,
+        dirpath: Path | str = "",
+        write_strategy: WriteStrategy | None = None,
+    ):
+        """
+        Constructor.
+
+        A `WriteStrategy` must be provided at instantiation or later via
+        `set_writing_strategy`.
+
+        Parameters
+        ----------
+        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.
+        write_strategy : WriteStrategy or None, default=None
+            A strategy for writing predictions.
+        """
+        super().__init__(write_interval="batch")
+
+        self.writing_predictions = True  # flag to turn off predictions
+
+        # forward declaration
+        self.write_strategy: WriteStrategy
+        if write_strategy is not None:  # avoid `WriteStrategy | None` type
+            self.write_strategy = write_strategy
+
+        self.dirpath: Path
+
+        # if a dirpath is provided, initialize it
+        # in some cases (e.g. zarr), destination is provided by the zarr store path
+        if dirpath != "":
+            self._init_dirpath(dirpath)
+
+    def disable_writing(self, disable_writing: bool) -> None:
+        """Disable writing.
+
+        Parameters
+        ----------
+        disable_writing : bool
+            If writing predictions should be disabled.
+        """
+        self.writing_predictions = disable_writing
+
+    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":
+            if self.dirpath is not None:
+                # make prediction output directory
+                logger.info("Making prediction output directory.")
+                self.dirpath.mkdir(parents=True, exist_ok=True)
+
+    def set_writing_strategy(
+        self,
+        write_type: SupportedWriteType,
+        tiled: bool,
+        write_func: WriteFunc | None = None,
+        write_extension: str | None = None,
+        write_func_kwargs: dict[str, Any] | None = None,
+    ) -> None:
+        """
+        Set the writing strategy.
+
+        Must be called before writing predictions.
+
+        Parameters
+        ----------
+        write_type : SupportedWriteType
+            The type of writing to perform.
+        tiled : bool
+            Whether to write in tiled format.
+        write_func : WriteFunc or None, default=None
+            A custom writing function.
+        write_extension : str or None, default=None
+            The file extension to use when writing files.
+        write_func_kwargs : dict of str to Any, default=None
+            Additional keyword arguments to pass to `write_func`.
+        """
+        self.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,
+        )
+
+    def write_on_batch_end(
+        self,
+        trainer: Trainer,
+        pl_module: LightningModule,
+        prediction: ImageRegionData,
+        batch_indices: Sequence[int] | None,
+        batch: ImageRegionData,
+        batch_idx: int,
+        dataloader_idx: int,
+    ) -> None:
+        """
+        Write predictions at the end of a batch.
+
+        Writing method is determined by the attribute `write_strategy`.
+
+        Parameters
+        ----------
+        trainer : Trainer
+            PyTorch Lightning trainer.
+        pl_module : LightningModule
+            PyTorch Lightning module.
+        prediction : ImageRegionData
+            Prediction outputs of `batch`.
+        batch_indices : sequence of Any, optional
+            Batch indices.
+        batch : ImageRegionData
+            Input batch.
+        batch_idx : int
+            Batch index.
+        dataloader_idx : int
+            Dataloader index.
+        """
+        # if writing prediction is turned off
+        if not self.writing_predictions:
+            return
+
+        if self.write_strategy is not None:
+            assert prediction is not None
+            predictions = decollate_image_region_data(prediction)
+
+            self.write_strategy.write_batch(
+                dirpath=self.dirpath,
+                predictions=predictions,
+            )
+        else:
+            raise RuntimeError(
+                "No write strategy defined for `PredictionWriterCallback`, cannot write"
+                " predictions. Call `set_writing_strategy` to pass a write strategy."
+            )

careamics/lightning/dataset_ng/callbacks/prediction_writer/write_image_strategy.py

@@ -0,0 +1,91 @@
+"""A strategy writing whole images directly."""
+
+from pathlib import Path
+from typing import Any
+
+from careamics.dataset_ng.dataset import ImageRegionData
+from careamics.file_io import WriteFunc
+from careamics.lightning.dataset_ng.prediction import (
+    combine_samples,
+)
+
+from .file_path_utils import create_write_file_path
+from .write_strategy import WriteStrategy
+
+
+# TODO bug: batch is over samples for whole images, if one batch does not cover
+# all samples, it will write an incomplete image, then overwrite it whith the next
+# batch
+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,
+        dirpath: Path,
+        predictions: list[ImageRegionData],
+    ) -> None:
+        """
+        Save full images.
+
+        Parameters
+        ----------
+        dirpath : Path
+            Path to directory to save predictions to.
+        predictions : list[ImageRegionData]
+            Decollated predictions.
+        """
+        assert predictions is not None
+
+        image_lst, sources = combine_samples(predictions)
+
+        for i, image in enumerate(image_lst):
+            file_path = create_write_file_path(
+                dirpath=dirpath,
+                file_path=Path(sources[i]),
+                write_extension=self.write_extension,
+            )
+            self.write_func(file_path=file_path, img=image, **self.write_func_kwargs)

careamics/lightning/dataset_ng/callbacks/prediction_writer/write_strategy.py

@@ -0,0 +1,27 @@
+"""Module containing different strategies for writing predictions."""
+
+from pathlib import Path
+from typing import Protocol
+
+from careamics.dataset_ng.dataset import ImageRegionData
+
+
+class WriteStrategy(Protocol):
+    """Protocol for write strategy classes."""
+
+    def write_batch(
+        self,
+        dirpath: Path,
+        predictions: list[ImageRegionData],
+    ) -> None:
+        """
+        WriteStrategy subclasses must contain this function to write a batch.
+
+        Parameters
+        ----------
+        dirpath : Path
+            Path to directory to save predictions to.
+        predictions : list[ImageRegionData]
+            Decollated predictions.
+        """
+        ...

careamics/lightning/dataset_ng/callbacks/prediction_writer/write_strategy_factory.py

@@ -0,0 +1,214 @@
+"""Module containing convenience function to create `WriteStrategy`."""
+
+from typing import Any
+
+from careamics.config.support import SupportedData
+from careamics.file_io import SupportedWriteType, WriteFunc, get_write_func
+
+from .cached_tiles_strategy import CachedTiles
+from .write_image_strategy import WriteImage
+from .write_strategy import WriteStrategy
+from .write_tiles_zarr_strategy import WriteTilesZarr
+
+
+def create_write_strategy(
+    write_type: SupportedWriteType,
+    tiled: bool,
+    write_func: WriteFunc | None = None,
+    write_extension: str | None = None,
+    write_func_kwargs: dict[str, Any] | None = None,
+) -> WriteStrategy:
+    """
+    Create a write strategy from convenient parameters.
+
+    Parameters
+    ----------
+    write_type : {"tiff", "zarr", "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.
+
+    Returns
+    -------
+    WriteStrategy
+        A strategy for writing predicions.
+
+    Notes
+    -----
+    The `write_func` function signature must match that of the example below
+        ```
+        write_func(file_path: Path, img: NDArray, *args, **kwargs) -> None: ...
+        ```
+
+    The `write_func_kwargs` will be passed to the `write_func` doing the following:
+        ```
+        write_func(file_path=file_path, img=img, **kwargs)
+        ```
+    """
+    if write_func_kwargs is None:
+        write_func_kwargs = {}
+
+    write_strategy: WriteStrategy
+    if not tiled:
+        write_func = select_write_func(write_type=write_type, write_func=write_func)
+        write_extension = select_write_extension(
+            write_type=write_type, write_extension=write_extension
+        )
+        write_strategy = WriteImage(
+            write_func=write_func,
+            write_extension=write_extension,
+            write_func_kwargs=write_func_kwargs,
+        )
+    else:
+        # select CacheTiles or ZarrTiles
+        write_strategy = _create_tiled_write_strategy(
+            write_type=write_type,
+            write_func=write_func,
+            write_extension=write_extension,
+            write_func_kwargs=write_func_kwargs,
+        )
+
+    return write_strategy
+
+
+def _create_tiled_write_strategy(
+    write_type: SupportedWriteType,
+    write_func: WriteFunc | None,
+    write_extension: str | None,
+    write_func_kwargs: dict[str, Any],
+) -> WriteStrategy:
+    """
+    Create a tiled write strategy.
+
+    Either `CacheTiles` for caching tiles until a whole image is predicted or
+    `ZarrTiles` for writing tiles directly to disk.
+
+    Parameters
+    ----------
+    write_type : {"tiff", "zarr", "custom"}
+        The data type to save as, includes custom.
+    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}
+        Additional keyword arguments to be passed to the save function.
+
+    Returns
+    -------
+    WriteStrategy
+        A strategy for writing tiled predictions.
+
+    Raises
+    ------
+    NotImplementedError
+        if `write_type="zarr" is chosen.
+    """
+    if write_type == "zarr":
+        return WriteTilesZarr()
+    else:
+        write_func = select_write_func(write_type=write_type, write_func=write_func)
+        write_extension = select_write_extension(
+            write_type=write_type, write_extension=write_extension
+        )
+        return CachedTiles(
+            write_func=write_func,
+            write_extension=write_extension,
+            write_func_kwargs=write_func_kwargs,
+        )
+
+
+def select_write_func(
+    write_type: SupportedWriteType, write_func: WriteFunc | None = None
+) -> WriteFunc:
+    """
+    Return a function to write images.
+
+    If `write_type` is "custom" then `write_func`, otherwise the known write function
+    is selected.
+
+    Parameters
+    ----------
+    write_type : {"tiff", "custom"}
+        The data type to save as, includes custom.
+    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.
+
+    Returns
+    -------
+    WriteFunc
+        A function for writing images.
+
+    Raises
+    ------
+    ValueError
+        If `write_type="custom"` but `write_func` has not been given.
+
+    Notes
+    -----
+    The `write_func` function signature must match that of the example below
+        ```
+        write_func(file_path: Path, img: NDArray, *args, **kwargs) -> None: ...
+        ```
+    """
+    if write_type == SupportedData.CUSTOM:
+        if write_func is None:
+            raise ValueError(
+                "A save function must be provided for custom data types."
+                # TODO: link to how save functions should be implemented
+            )
+        else:
+            write_func = write_func
+    else:
+        write_func = get_write_func(write_type)
+    return write_func
+
+
+def select_write_extension(
+    write_type: SupportedWriteType, write_extension: str | None = None
+) -> str:
+    """
+    Return an extension to add to file paths.
+
+    If `write_type` is "custom" then `write_extension`, otherwise the known
+    write extension is selected.
+
+    Parameters
+    ----------
+    write_type : {"tiff", "custom"}
+        The data type to save as, includes custom.
+    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.
+
+    Returns
+    -------
+    str
+        The extension to be added to file paths.
+
+    Raises
+    ------
+    ValueError
+        If `self.save_type="custom"` but `save_extension` has not been given.
+    """
+    write_type_: SupportedData = SupportedData(write_type)  # new variable for mypy
+    if write_type_ == SupportedData.CUSTOM:
+        if write_extension is None:
+            raise ValueError("A save extension must be provided for custom data types.")
+        else:
+            write_extension = write_extension
+    else:
+        # kind of a weird pattern -> reason to move get_extension from SupportedData
+        write_extension = write_type_.get_extension(write_type_)
+    return write_extension