careamics 0.0.19__py3-none-any.whl

careamics/lightning/callbacks/prediction_writer_callback/write_strategy_factory.py

@@ -0,0 +1,215 @@
+"""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 .write_strategy import CacheTiles, WriteImage, WriteStrategy
+
+
+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", "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 WriteTilesZarr (when implemented)
+        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
+    `WriteTilesZarr` for writing tiles directly to disk.
+
+    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.
+    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 == SupportedData.ZARR:
+    #    create *args, **kwargs
+    #    return WriteTilesZarr(*args, **kwargs)
+    # else:
+    if write_type == "zarr":
+        raise NotImplementedError("Saving to zarr is not implemented yet.")
+    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 CacheTiles(
+            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

careamics/lightning/callbacks/progress_bar_callback.py

@@ -0,0 +1,90 @@
+"""Progressbar callback."""
+
+import sys
+from typing import Union
+
+from pytorch_lightning import LightningModule, Trainer
+from pytorch_lightning.callbacks import TQDMProgressBar
+from tqdm.auto import tqdm
+
+
+class ProgressBarCallback(TQDMProgressBar):
+    """Progress bar for training and validation steps."""
+
+    def init_train_tqdm(self) -> tqdm:
+        """Override this to customize the tqdm bar for training.
+
+        Returns
+        -------
+        tqdm
+            A tqdm bar.
+        """
+        bar = tqdm(
+            desc="Training",
+            position=(2 * self.process_position),
+            disable=self.is_disabled,
+            leave=True,
+            dynamic_ncols=True,
+            file=sys.stdout,
+            smoothing=0,
+        )
+        return bar
+
+    def init_validation_tqdm(self) -> tqdm:
+        """Override this to customize the tqdm bar for validation.
+
+        Returns
+        -------
+        tqdm
+            A tqdm bar.
+        """
+        # The main progress bar doesn't exist in `trainer.validate()`
+        has_main_bar = self.train_progress_bar is not None
+        bar = tqdm(
+            desc="Validating",
+            position=(2 * self.process_position + has_main_bar),
+            disable=self.is_disabled,
+            leave=False,
+            dynamic_ncols=True,
+            file=sys.stdout,
+        )
+        return bar
+
+    def init_test_tqdm(self) -> tqdm:
+        """Override this to customize the tqdm bar for testing.
+
+        Returns
+        -------
+        tqdm
+            A tqdm bar.
+        """
+        bar = tqdm(
+            desc="Testing",
+            position=(2 * self.process_position),
+            disable=self.is_disabled,
+            leave=True,
+            dynamic_ncols=False,
+            ncols=100,
+            file=sys.stdout,
+        )
+        return bar
+
+    def get_metrics(
+        self, trainer: Trainer, pl_module: LightningModule
+    ) -> dict[str, Union[int, str, float, dict[str, float]]]:
+        """Override this to customize the metrics displayed in the progress bar.
+
+        Parameters
+        ----------
+        trainer : Trainer
+            The trainer object.
+        pl_module : LightningModule
+            The LightningModule object, unused.
+
+        Returns
+        -------
+        dict
+            A dictionary with the metrics to display in the progress bar.
+        """
+        pbar_metrics = trainer.progress_bar_metrics
+        return {**pbar_metrics}

careamics/lightning/dataset_ng/__init__.py

@@ -0,0 +1 @@
+"""Next-Generation DataModules for Careamics."""

careamics/lightning/dataset_ng/callbacks/__init__.py

@@ -0,0 +1 @@
+"""NG Dataset compatible callbacks for PyTorch Lightning."""

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

@@ -0,0 +1,29 @@
+"""A package for the `PredictionWriterCallback` class and utilities."""
+
+__all__ = [
+    "CachedTiles",
+    "PredictionWriterCallback",
+    "WriteImage",
+    "WriteStrategy",
+    "WriteTilesZarr",
+    "create_write_file_path",
+    "create_write_strategy",
+    "decollate_image_region_data",
+    "select_write_extension",
+    "select_write_func",
+]
+
+from .cached_tiles_strategy import CachedTiles
+from .file_path_utils import create_write_file_path
+from .prediction_writer_callback import (
+    PredictionWriterCallback,
+    decollate_image_region_data,
+)
+from .write_image_strategy import WriteImage
+from .write_strategy import WriteStrategy
+from .write_strategy_factory import (
+    create_write_strategy,
+    select_write_extension,
+    select_write_func,
+)
+from .write_tiles_zarr_strategy import WriteTilesZarr

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

@@ -0,0 +1,164 @@
+"""A writing strategy that caches tiles until a whole image is predicted."""
+
+from collections import defaultdict
+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 (
+    stitch_single_prediction,
+)
+
+from .file_path_utils import create_write_file_path
+from .write_strategy import WriteStrategy
+
+
+class CachedTiles(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: dict[int, list[ImageRegionData]] = defaultdict(list)
+
+    def write_batch(
+        self,
+        dirpath: Path,
+        predictions: list[ImageRegionData],
+    ) -> None:
+        """
+        Cache tiles until the last tile is predicted, then save the stitched image.
+
+        Parameters
+        ----------
+        dirpath : Path
+            Path to directory to save predictions to.
+        predictions : list[ImageRegionData]
+            Decollated predictions.
+        """
+        assert predictions is not None
+
+        # cache tiles
+        for tile in predictions:
+            data_idx = tile.region_spec["data_idx"]
+            self.tile_cache[data_idx].append(tile)
+
+        self._write_images(dirpath)
+
+    def _get_full_images(self) -> list[int]:
+        """
+        Get data indices of full images contained in the cache.
+
+        Returns
+        -------
+        list of int
+            Data indices of full images contained in the cache.
+        """
+        full_images = []
+        for data_idx in self.tile_cache.keys():
+            exp_n_tiles = self.tile_cache[data_idx][0].region_spec["total_tiles"]
+
+            if len(self.tile_cache[data_idx]) == exp_n_tiles:
+                full_images.append(data_idx)
+            elif len(self.tile_cache[data_idx]) > exp_n_tiles:
+                raise ValueError(
+                    f"More tiles cached for data_idx {data_idx} than expected. "
+                    f"Expected {exp_n_tiles}, found "
+                    f"{len(self.tile_cache[data_idx])}."
+                )
+
+        return full_images
+
+    def _stitch_and_write_single(
+        self, dirpath: Path, tiles: list[ImageRegionData]
+    ) -> None:
+        """
+        Stitch and write a single image from tiles.
+
+        Parameters
+        ----------
+        dirpath : Path
+            Path to directory to save predictions to.
+        tiles : list[ImageRegionData]
+            Tiles to stitch and write.
+        """
+        # stitch prediction
+        prediction_image = stitch_single_prediction(tiles)
+
+        # write prediction
+        source: Path = Path(tiles[0].source)
+        file_path = create_write_file_path(
+            dirpath=dirpath,
+            file_path=source,
+            write_extension=self.write_extension,
+        )
+        self.write_func(
+            file_path=file_path, img=prediction_image, **self.write_func_kwargs
+        )
+
+    def _write_images(self, dirpath: Path) -> None:
+        """
+        Write full images from cached tiles.
+
+        Parameters
+        ----------
+        dirpath : Path
+            Path to directory to save predictions to.
+        """
+        full_images = self._get_full_images()
+        for data_idx in full_images:
+            tiles = self.tile_cache.pop(data_idx)
+            self._stitch_and_write_single(dirpath, tiles)

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

@@ -0,0 +1,33 @@
+"""Module containing file path utilities for `WriteStrategy` to use."""
+
+from pathlib import Path
+
+
+def create_write_file_path(
+    dirpath: Path, file_path: Path, write_extension: str
+) -> Path:
+    """
+    Create the file name for the output file.
+
+    Takes the original file path, changes the directory to `dirpath` and changes
+    the extension to `write_extension`.
+
+    Parameters
+    ----------
+    dirpath : pathlib.Path
+        The output directory to write file to.
+    file_path : pathlib.Path
+        The original file path.
+    write_extension : str
+        The extension that output files should have.
+
+    Returns
+    -------
+    Path
+        The output file path.
+    """
+    file_path = Path(file_path)  # as a guard against str input
+
+    file_name = Path(file_path.stem).with_suffix(write_extension)
+    file_path = dirpath / file_name
+    return file_path