careamics 0.0.1__py3-none-any.whl → 0.0.2__py3-none-any.whl

careamics/file_io/write/get_func.py

@@ -0,0 +1,63 @@
+"""Module to get write functions."""
+
+from pathlib import Path
+from typing import Literal, Protocol
+
+from numpy.typing import NDArray
+
+from careamics.config.support import SupportedData
+
+from .tiff import write_tiff
+
+SupportedWriteType = Literal["tiff", "custom"]
+
+
+# This is very strict, arguments have to be called file_path & img
+# Alternative? - doesn't capture *args & **kwargs
+# WriteFunc = Callable[[Path, NDArray], None]
+class WriteFunc(Protocol):
+    """Protocol for type hinting write functions."""
+
+    def __call__(self, file_path: Path, img: NDArray, *args, **kwargs) -> None:
+        """
+        Type hinted callables must match this function signature (not including self).
+
+        Parameters
+        ----------
+        file_path : pathlib.Path
+            Path to file.
+        img : numpy.ndarray
+            Image data to save.
+        *args
+            Other positional arguments.
+        **kwargs
+            Other keyword arguments.
+        """
+
+
+WRITE_FUNCS: dict[SupportedData, WriteFunc] = {
+    SupportedData.TIFF: write_tiff,
+}
+
+
+def get_write_func(data_type: SupportedWriteType) -> WriteFunc:
+    """
+    Get the write function for the data type.
+
+    Parameters
+    ----------
+    data_type : {"tiff", "custom"}
+        Data type.
+
+    Returns
+    -------
+    callable
+        Write function.
+    """
+    # error raised here if not supported
+    data_type_ = SupportedData(data_type)  # new variable for mypy
+    # error if no write func.
+    if data_type_ not in WRITE_FUNCS:
+        raise NotImplementedError(f"No write function for data type '{data_type}'.")
+
+    return WRITE_FUNCS[data_type_]

careamics/file_io/write/tiff.py

@@ -0,0 +1,40 @@
+"""Write tiff function."""
+
+from fnmatch import fnmatch
+from pathlib import Path
+
+import tifffile
+from numpy.typing import NDArray
+
+from careamics.config.support import SupportedData
+
+
+def write_tiff(file_path: Path, img: NDArray, *args, **kwargs) -> None:
+    # TODO: add link to tiffile docs for args kwrgs?
+    """
+    Write tiff files.
+
+    Parameters
+    ----------
+    file_path : pathlib.Path
+        Path to file.
+    img : numpy.ndarray
+        Image data to save.
+    *args
+        Positional arguments passed to `tifffile.imwrite`.
+    **kwargs
+        Keyword arguments passed to `tifffile.imwrite`.
+
+    Raises
+    ------
+    ValueError
+        When the file extension of `file_path` does not match the Unix shell-style
+        pattern '*.tif*'.
+    """
+    if not fnmatch(
+        file_path.suffix, SupportedData.get_extension_pattern(SupportedData.TIFF)
+    ):
+        raise ValueError(
+            f"Unexpected extension '{file_path.suffix}' for save file type 'tiff'."
+        )
+    tifffile.imwrite(file_path, img, *args, **kwargs)

careamics/lightning/__init__.py

@@ -0,0 +1,17 @@
+"""CAREamics PyTorch Lightning modules."""
+
+__all__ = [
+    "CAREamicsModule",
+    "create_careamics_module",
+    "TrainDataModule",
+    "create_train_datamodule",
+    "PredictDataModule",
+    "create_predict_datamodule",
+    "HyperParametersCallback",
+    "ProgressBarCallback",
+]
+
+from .callbacks import HyperParametersCallback, ProgressBarCallback
+from .lightning_module import CAREamicsModule, create_careamics_module
+from .predict_data_module import PredictDataModule, create_predict_datamodule
+from .train_data_module import TrainDataModule, create_train_datamodule

careamics/lightning/callbacks/__init__.py

@@ -0,0 +1,11 @@
+"""Callbacks module."""
+
+__all__ = [
+    "HyperParametersCallback",
+    "ProgressBarCallback",
+    "PredictionWriterCallback",
+]
+
+from .hyperparameters_callback import HyperParametersCallback
+from .prediction_writer_callback import PredictionWriterCallback
+from .progress_bar_callback import ProgressBarCallback

careamics/lightning/callbacks/hyperparameters_callback.py

@@ -0,0 +1,49 @@
+"""Callback saving CAREamics configuration as hyperparameters in the model."""
+
+from pytorch_lightning import LightningModule, Trainer
+from pytorch_lightning.callbacks import Callback
+
+from careamics.config import Configuration
+
+
+class HyperParametersCallback(Callback):
+    """
+    Callback allowing saving CAREamics configuration as hyperparameters in the model.
+
+    This allows saving the configuration as dictionnary in the checkpoints, and
+    loading it subsequently in a CAREamist instance.
+
+    Parameters
+    ----------
+    config : Configuration
+        CAREamics configuration to be saved as hyperparameter in the model.
+
+    Attributes
+    ----------
+    config : Configuration
+        CAREamics configuration to be saved as hyperparameter in the model.
+    """
+
+    def __init__(self, config: Configuration) -> None:
+        """
+        Constructor.
+
+        Parameters
+        ----------
+        config : Configuration
+            CAREamics configuration to be saved as hyperparameter in the model.
+        """
+        self.config = config
+
+    def on_train_start(self, trainer: Trainer, pl_module: LightningModule) -> None:
+        """
+        Update the hyperparameters of the model with the configuration on train start.
+
+        Parameters
+        ----------
+        trainer : Trainer
+            PyTorch Lightning trainer, unused.
+        pl_module : LightningModule
+            PyTorch Lightning module.
+        """
+        pl_module.hparams.update(self.config.model_dump())

careamics/lightning/callbacks/prediction_writer_callback/__init__.py

@@ -0,0 +1,20 @@
+"""A package for the `PredictionWriterCallback` class and utilities."""
+
+__all__ = [
+    "PredictionWriterCallback",
+    "create_write_strategy",
+    "WriteStrategy",
+    "WriteImage",
+    "CacheTiles",
+    "WriteTilesZarr",
+    "select_write_extension",
+    "select_write_func",
+]
+
+from .prediction_writer_callback import PredictionWriterCallback
+from .write_strategy import CacheTiles, WriteImage, WriteStrategy, WriteTilesZarr
+from .write_strategy_factory import (
+    create_write_strategy,
+    select_write_extension,
+    select_write_func,
+)

careamics/lightning/callbacks/prediction_writer_callback/file_path_utils.py

@@ -0,0 +1,56 @@
+"""Module containing file path utilities for `WriteStrategy` to use."""
+
+from pathlib import Path
+from typing import Union
+
+from careamics.dataset import IterablePredDataset, IterableTiledPredDataset
+
+
+# TODO: move to datasets package ?
+def get_sample_file_path(
+    dataset: Union[IterableTiledPredDataset, IterablePredDataset], sample_id: int
+) -> Path:
+    """
+    Get the file path for a particular sample.
+
+    Parameters
+    ----------
+    dataset : IterableTiledPredDataset or IterablePredDataset
+        Dataset.
+    sample_id : int
+        Sample ID, the index of the file in the dataset `dataset`.
+
+    Returns
+    -------
+    Path
+        The file path corresponding to the sample with the ID `sample_id`.
+    """
+    return dataset.data_files[sample_id]
+
+
+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_name = Path(file_path.stem).with_suffix(write_extension)
+    file_path = dirpath / file_name
+    return file_path

careamics/lightning/callbacks/prediction_writer_callback/prediction_writer_callback.py

@@ -0,0 +1,233 @@
+"""Module containing `PredictionWriterCallback` class."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any, Optional, Sequence, 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: Optional[WriteFunc] = None,
+        write_extension: Optional[str] = None,
+        write_func_kwargs: Optional[dict[str, Any]] = 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: Optional[Sequence[int]],
+        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,
+        )