careamics 0.0.19__py3-none-any.whl

careamics/dataset_ng/patching_strategies/whole_sample.py

@@ -0,0 +1,61 @@
+from collections.abc import Sequence
+
+from .patching_strategy_protocol import PatchSpecs
+
+
+class WholeSamplePatchingStrategy:
+    # TODO: warn this strategy should only be used with batch size = 1
+    #   for the case of multiple image stacks with different dimensions
+
+    # TODO: docs
+    def __init__(self, data_shapes: Sequence[Sequence[int]]):
+        self.data_shapes = data_shapes
+
+        self.patch_specs: list[PatchSpecs] = self._initialize_patch_specs()
+
+    @property
+    def n_patches(self) -> int:
+        return len(self.patch_specs)
+
+    def get_patch_spec(self, index: int) -> PatchSpecs:
+        return self.patch_specs[index]
+
+    # Note: this is used by the FileIterSampler
+    def get_patch_indices(self, data_idx: int) -> Sequence[int]:
+        """
+        Get the patch indices will return patches for a specific `image_stack`.
+
+        The `image_stack` corresponds to the given `data_idx`.
+
+        Parameters
+        ----------
+        data_idx : int
+            An index that corresponds to a given `image_stack`.
+
+        Returns
+        -------
+        sequence of int
+            A sequence of patch indices, that when used to index the `CAREamicsDataset
+            will return a patch that comes from the `image_stack` corresponding to the
+            given `data_idx`.
+        """
+        return [
+            i
+            for i, patch_spec in enumerate(self.patch_specs)
+            if patch_spec["data_idx"] == data_idx
+        ]
+
+    def _initialize_patch_specs(self) -> list[PatchSpecs]:
+        patch_specs: list[PatchSpecs] = []
+        for data_idx, data_shape in enumerate(self.data_shapes):
+            spatial_shape = data_shape[2:]
+            for sample_idx in range(data_shape[0]):
+                patch_specs.append(
+                    {
+                        "data_idx": data_idx,
+                        "sample_idx": sample_idx,
+                        "coords": tuple(0 for _ in spatial_shape),
+                        "patch_size": spatial_shape,
+                    }
+                )
+        return patch_specs

careamics/file_io/__init__.py

@@ -0,0 +1,15 @@
+"""Functions relating reading and writing image files."""
+
+__all__ = [
+    "ReadFunc",
+    "SupportedWriteType",
+    "WriteFunc",
+    "get_read_func",
+    "get_write_func",
+    "read",
+    "write",
+]
+
+from . import read, write
+from .read import ReadFunc, get_read_func
+from .write import SupportedWriteType, WriteFunc, get_write_func

careamics/file_io/read/__init__.py

@@ -0,0 +1,11 @@
+"""Functions relating to reading image files of different formats."""
+
+__all__ = [
+    "ReadFunc",
+    "get_read_func",
+    "read_tiff",
+    "read_zarr",
+]
+
+from .get_func import ReadFunc, get_read_func
+from .tiff import read_tiff

careamics/file_io/read/get_func.py

@@ -0,0 +1,57 @@
+"""Module to get read functions."""
+
+from collections.abc import Callable
+from pathlib import Path
+from typing import Protocol, Union
+
+from numpy.typing import NDArray
+
+from careamics.config.support import SupportedData
+
+from .tiff import read_tiff
+
+
+# This is very strict, function signature has to match including arg names
+# See WriteFunc notes
+class ReadFunc(Protocol):
+    """Protocol for type hinting read functions."""
+
+    def __call__(self, file_path: Path, *args, **kwargs) -> NDArray:
+        """
+        Type hinted callables must match this function signature (not including self).
+
+        Parameters
+        ----------
+        file_path : pathlib.Path
+            Path to file.
+        *args
+            Other positional arguments.
+        **kwargs
+            Other keyword arguments.
+        """
+
+
+READ_FUNCS: dict[SupportedData, ReadFunc] = {
+    SupportedData.TIFF: read_tiff,
+}
+
+
+def get_read_func(data_type: Union[str, SupportedData]) -> Callable:
+    """
+    Get the read function for the data type.
+
+    Parameters
+    ----------
+    data_type : SupportedData
+        Data type.
+
+    Returns
+    -------
+    callable
+        Read function.
+    """
+    if data_type in READ_FUNCS:
+        data_type = SupportedData(data_type)  # mypy complaining about dict key type
+        return READ_FUNCS[data_type]
+    else:
+        raise NotImplementedError(f"Data type '{data_type}' is not supported.")

careamics/file_io/read/tiff.py

@@ -0,0 +1,58 @@
+"""Functions to read tiff images."""
+
+import logging
+from fnmatch import fnmatch
+from pathlib import Path
+
+import numpy as np
+import tifffile
+
+from careamics.config.support import SupportedData
+from careamics.utils.logging import get_logger
+
+logger = get_logger(__name__)
+
+
+def read_tiff(file_path: Path, *args: list, **kwargs: dict) -> np.ndarray:
+    """
+    Read a tiff file and return a numpy array.
+
+    Parameters
+    ----------
+    file_path : Path
+        Path to a file.
+    *args : list
+        Additional arguments.
+    **kwargs : dict
+        Additional keyword arguments.
+
+    Returns
+    -------
+    np.ndarray
+        Resulting array.
+
+    Raises
+    ------
+    ValueError
+        If the file failed to open.
+    OSError
+        If the file failed to open.
+    ValueError
+        If the file is not a valid tiff.
+    ValueError
+        If the data dimensions are incorrect.
+    ValueError
+        If the axes length is incorrect.
+    """
+    if fnmatch(
+        file_path.suffix, SupportedData.get_extension_pattern(SupportedData.TIFF)
+    ):
+        try:
+            array = tifffile.imread(file_path)
+        except (ValueError, OSError) as e:
+            logging.exception(f"Exception in file {file_path}: {e}, skipping it.")
+            raise e
+    else:
+        raise ValueError(f"File {file_path} is not a valid tiff.")
+
+    return array

careamics/file_io/write/__init__.py

@@ -0,0 +1,15 @@
+"""Functions relating to writing image files of different formats."""
+
+__all__ = [
+    "SupportedWriteType",
+    "WriteFunc",
+    "get_write_func",
+    "write_tiff",
+]
+
+from .get_func import (
+    SupportedWriteType,
+    WriteFunc,
+    get_write_func,
+)
+from .tiff import write_tiff

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", "zarr", "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,32 @@
+"""CAREamics PyTorch Lightning modules."""
+
+__all__ = [
+    "DataStatsCallback",
+    "FCNModule",
+    "HyperParametersCallback",
+    "MicroSplitDataModule",
+    "PredictDataModule",
+    "ProgressBarCallback",
+    "TrainDataModule",
+    "VAEModule",
+    "create_careamics_module",
+    "create_microsplit_predict_datamodule",
+    "create_microsplit_train_datamodule",
+    "create_predict_datamodule",
+    "create_train_datamodule",
+    "create_unet_based_module",
+    "create_vae_based_module",
+]
+
+from .callbacks import DataStatsCallback, HyperParametersCallback, ProgressBarCallback
+from .lightning_module import FCNModule, VAEModule, create_careamics_module
+from .microsplit_data_module import (
+    MicroSplitDataModule,
+    create_microsplit_predict_datamodule,
+    create_microsplit_train_datamodule,
+)
+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,13 @@
+"""Callbacks module."""
+
+__all__ = [
+    "DataStatsCallback",
+    "HyperParametersCallback",
+    "PredictionWriterCallback",
+    "ProgressBarCallback",
+]
+
+from .data_stats_callback import DataStatsCallback
+from .hyperparameters_callback import HyperParametersCallback
+from .prediction_writer_callback import PredictionWriterCallback
+from .progress_bar_callback import ProgressBarCallback

careamics/lightning/callbacks/data_stats_callback.py

@@ -0,0 +1,33 @@
+"""Data statistics callback."""
+
+import pytorch_lightning as L
+from pytorch_lightning.callbacks import Callback
+
+
+class DataStatsCallback(Callback):
+    """Callback to update model's data statistics from datamodule.
+
+    This callback ensures that the model has access to the data statistics (mean, std)
+    calculated by the datamodule before training starts.
+    """
+
+    def setup(self, trainer: L.Trainer, module: L.LightningModule, stage: str) -> None:
+        """Called when trainer is setting up.
+
+        Parameters
+        ----------
+        trainer : Lightning.Trainer
+            PyTorch Lightning trainer.
+        module : Lightning.LightningModule
+            Lightning module.
+        stage : str
+            Current stage (fit, validate, test, or predict).
+        """
+        if stage == "fit":
+            # Get data statistics from datamodule
+            (data_mean, data_std), _ = trainer.datamodule.get_data_stats()
+
+            # Set data statistics in the model's likelihood module
+            module.noise_model_likelihood.set_data_stats(
+                data_mean=data_mean["target"], data_std=data_std["target"]
+            )

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 dictionary 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__ = [
+    "CacheTiles",
+    "PredictionWriterCallback",
+    "WriteImage",
+    "WriteStrategy",
+    "WriteTilesZarr",
+    "create_write_strategy",
+    "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