careamics 0.1.0rc6__py3-none-any.whl → 0.1.0rc8__py3-none-any.whl

careamics/dataset/iterable_tiled_pred_dataset.py

@@ -0,0 +1,140 @@
+"""Iterable tiled prediction dataset used to load data file by file."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any, Callable, Generator
+
+from numpy.typing import NDArray
+from torch.utils.data import IterableDataset
+
+from careamics.file_io.read import read_tiff
+from careamics.transforms import Compose
+
+from ..config import InferenceConfig
+from ..config.tile_information import TileInformation
+from ..config.transformations import NormalizeModel
+from .dataset_utils import iterate_over_files
+from .tiling import extract_tiles
+
+
+class IterableTiledPredDataset(IterableDataset):
+    """Tiled prediction dataset.
+
+    Parameters
+    ----------
+    prediction_config : InferenceConfig
+        Inference configuration.
+    src_files : list of pathlib.Path
+        List of data files.
+    read_source_func : Callable, optional
+        Read source function for custom types, by default read_tiff.
+    **kwargs : Any
+        Additional keyword arguments, unused.
+
+    Attributes
+    ----------
+    data_path : str or pathlib.Path
+        Path to the data, must be a directory.
+    axes : str
+        Description of axes in format STCZYX.
+    mean : float, optional
+        Expected mean of the dataset, by default None.
+    std : float, optional
+        Expected standard deviation of the dataset, by default None.
+    patch_transform : Callable, optional
+        Patch transform callable, by default None.
+    """
+
+    def __init__(
+        self,
+        prediction_config: InferenceConfig,
+        src_files: list[Path],
+        read_source_func: Callable = read_tiff,
+        **kwargs: Any,
+    ) -> None:
+        """Constructor.
+
+        Parameters
+        ----------
+        prediction_config : InferenceConfig
+            Inference configuration.
+        src_files : List[Path]
+            List of data files.
+        read_source_func : Callable, optional
+            Read source function for custom types, by default read_tiff.
+        **kwargs : Any
+            Additional keyword arguments, unused.
+
+        Raises
+        ------
+        ValueError
+            If mean and std are not provided in the inference configuration.
+        """
+        if (
+            prediction_config.tile_size is None
+            or prediction_config.tile_overlap is None
+        ):
+            raise ValueError(
+                "Tile size and overlap must be provided for tiled prediction."
+            )
+
+        self.prediction_config = prediction_config
+        self.data_files = src_files
+        self.axes = prediction_config.axes
+        self.tile_size = prediction_config.tile_size
+        self.tile_overlap = prediction_config.tile_overlap
+        self.read_source_func = read_source_func
+
+        # check mean and std and create normalize transform
+        if (
+            self.prediction_config.image_means is None
+            or self.prediction_config.image_stds is None
+        ):
+            raise ValueError("Mean and std must be provided for prediction.")
+        else:
+            self.image_means = self.prediction_config.image_means
+            self.image_stds = self.prediction_config.image_stds
+
+            # instantiate normalize transform
+            self.patch_transform = Compose(
+                transform_list=[
+                    NormalizeModel(
+                        image_means=self.image_means,
+                        image_stds=self.image_stds,
+                    )
+                ],
+            )
+
+    def __iter__(
+        self,
+    ) -> Generator[tuple[NDArray, TileInformation], None, None]:
+        """
+        Iterate over data source and yield single patch.
+
+        Yields
+        ------
+        Generator of NDArray and TileInformation tuple
+            Generator of single tiles.
+        """
+        assert (
+            self.image_means is not None and self.image_stds is not None
+        ), "Mean and std must be provided"
+
+        for sample, _ in iterate_over_files(
+            self.prediction_config,
+            self.data_files,
+            read_source_func=self.read_source_func,
+        ):
+            # generate patches, return a generator of single tiles
+            patch_gen = extract_tiles(
+                arr=sample,
+                tile_size=self.tile_size,
+                overlaps=self.tile_overlap,
+            )
+
+            # apply transform to patches
+            for patch_array, tile_info in patch_gen:
+                transformed_patch, _ = self.patch_transform(patch=patch_array)
+
+                yield transformed_patch, tile_info

careamics/dataset/patching/patching.py

@@ -1,37 +1,83 @@
 """Patching functions."""
 
+from dataclasses import dataclass
 from pathlib import Path
-from typing import Callable, List, Tuple, Union
+from typing import Callable, Union
 
 import numpy as np
+from numpy.typing import NDArray
 
 from ...utils.logging import get_logger
 from ..dataset_utils import reshape_array
+from ..dataset_utils.running_stats import compute_normalization_stats
 from .sequential_patching import extract_patches_sequential
 
 logger = get_logger(__name__)
 
 
+@dataclass
+class Stats:
+    """Dataclass to store statistics."""
+
+    means: Union[NDArray, tuple, list, None]
+    """Mean of the data across channels."""
+
+    stds: Union[NDArray, tuple, list, None]
+    """Standard deviation of the data across channels."""
+
+    def get_statistics(self) -> tuple[list[float], list[float]]:
+        """Return the means and standard deviations.
+
+        Returns
+        -------
+        tuple of two lists of floats
+            Means and standard deviations.
+        """
+        if self.means is None or self.stds is None:
+            return [], []
+
+        return list(self.means), list(self.stds)
+
+
+@dataclass
+class PatchedOutput:
+    """Dataclass to store patches and statistics."""
+
+    patches: Union[NDArray]
+    """Image patches."""
+
+    targets: Union[NDArray, None]
+    """Target patches."""
+
+    image_stats: Stats
+    """Statistics of the image patches."""
+
+    target_stats: Stats
+    """Statistics of the target patches."""
+
+
 # called by in memory dataset
 def prepare_patches_supervised(
-    train_files: List[Path],
-    target_files: List[Path],
+    train_files: list[Path],
+    target_files: list[Path],
     axes: str,
-    patch_size: Union[List[int], Tuple[int, ...]],
+    patch_size: Union[list[int], tuple[int, ...]],
     read_source_func: Callable,
-) -> Tuple[np.ndarray, np.ndarray, float, float]:
+) -> PatchedOutput:
     """
     Iterate over data source and create an array of patches and corresponding targets.
 
+    The lists of Paths should be pre-sorted.
+
     Parameters
     ----------
-    train_files : List[Path]
+    train_files : list of pathlib.Path
         List of paths to training data.
-    target_files : List[Path]
+    target_files : list of pathlib.Path
         List of paths to target data.
     axes : str
         Axes of the data.
-    patch_size : Union[List[int], Tuple[int]]
+    patch_size : list or tuple of int
         Size of the patches.
     read_source_func : Callable
         Function to read the data.
@@ -41,9 +87,6 @@ def prepare_patches_supervised(
     np.ndarray
         Array of patches.
     """
-    train_files.sort()
-    target_files.sort()
-
     means, stds, num_samples = 0, 0, 0
     all_patches, all_targets = [], []
     for train_filename, target_filename in zip(train_files, target_files):
@@ -83,46 +126,47 @@ def prepare_patches_supervised(
             f"{target_files}."
         )
 
-    result_mean, result_std = means / num_samples, stds / num_samples
+    image_means, image_stds = compute_normalization_stats(np.concatenate(all_patches))
+    target_means, target_stds = compute_normalization_stats(np.concatenate(all_targets))
 
     patch_array: np.ndarray = np.concatenate(all_patches, axis=0)
     target_array: np.ndarray = np.concatenate(all_targets, axis=0)
     logger.info(f"Extracted {patch_array.shape[0]} patches from input array.")
 
-    return (
+    return PatchedOutput(
         patch_array,
         target_array,
-        result_mean,
-        result_std,
+        Stats(image_means, image_stds),
+        Stats(target_means, target_stds),
     )
 
 
 # called by in_memory_dataset
 def prepare_patches_unsupervised(
-    train_files: List[Path],
+    train_files: list[Path],
     axes: str,
-    patch_size: Union[List[int], Tuple[int]],
+    patch_size: Union[list[int], tuple[int]],
     read_source_func: Callable,
-) -> Tuple[np.ndarray, None, float, float]:
+) -> PatchedOutput:
     """Iterate over data source and create an array of patches.
 
     This method returns the mean and standard deviation of the image.
 
     Parameters
     ----------
-    train_files : List[Path]
+    train_files : list of pathlib.Path
         List of paths to training data.
     axes : str
         Axes of the data.
-    patch_size : Union[List[int], Tuple[int]]
+    patch_size : list or tuple of int
         Size of the patches.
     read_source_func : Callable
         Function to read the data.
 
     Returns
     -------
-    Tuple[np.ndarray, None, float, float]
-        Source and target patches, mean and standard deviation.
+    PatchedOutput
+        Dataclass holding patches and their statistics.
     """
     means, stds, num_samples = 0, 0, 0
     all_patches = []
@@ -149,21 +193,23 @@ def prepare_patches_unsupervised(
     if num_samples == 0:
         raise ValueError(f"No valid samples found in the input data: {train_files}.")
 
-    result_mean, result_std = means / num_samples, stds / num_samples
+    image_means, image_stds = compute_normalization_stats(np.concatenate(all_patches))
 
     patch_array: np.ndarray = np.concatenate(all_patches)
     logger.info(f"Extracted {patch_array.shape[0]} patches from input array.")
 
-    return patch_array, _, result_mean, result_std  # TODO return object?
+    return PatchedOutput(
+        patch_array, None, Stats(image_means, image_stds), Stats((), ())
+    )
 
 
 # called on arrays by in memory dataset
 def prepare_patches_supervised_array(
-    data: np.ndarray,
+    data: NDArray,
     axes: str,
-    data_target: np.ndarray,
-    patch_size: Union[List[int], Tuple[int]],
-) -> Tuple[np.ndarray, np.ndarray, float, float]:
+    data_target: NDArray,
+    patch_size: Union[list[int], tuple[int]],
+) -> PatchedOutput:
     """Iterate over data source and create an array of patches.
 
     This method expects an array of shape SC(Z)YX, where S and C can be singleton
@@ -173,28 +219,28 @@ def prepare_patches_supervised_array(
 
     Parameters
     ----------
-    data : np.ndarray
+    data : numpy.ndarray
         Input data array.
     axes : str
         Axes of the data.
-    data_target : np.ndarray
+    data_target : numpy.ndarray
         Target data array.
-    patch_size : Union[List[int], Tuple[int]]
+    patch_size : list or tuple of int
         Size of the patches.
 
     Returns
     -------
-    Tuple[np.ndarray, np.ndarray, float, float]
-        Source and target patches, mean and standard deviation.
+    PatchedOutput
+        Dataclass holding the source and target patches, with their statistics.
     """
-    # compute statistics
-    mean = data.mean()
-    std = data.std()
-
     # reshape array
     reshaped_sample = reshape_array(data, axes)
     reshaped_target = reshape_array(data_target, axes)
 
+    # compute statistics
+    image_means, image_stds = compute_normalization_stats(reshaped_sample)
+    target_means, target_stds = compute_normalization_stats(reshaped_target)
+
     # generate patches, return a generator
     patches, patch_targets = extract_patches_sequential(
         reshaped_sample, patch_size=patch_size, target=reshaped_target
@@ -205,20 +251,20 @@ def prepare_patches_supervised_array(
 
     logger.info(f"Extracted {patches.shape[0]} patches from input array.")
 
-    return (
+    return PatchedOutput(
         patches,
         patch_targets,
-        mean,
-        std,
+        Stats(image_means, image_stds),
+        Stats(target_means, target_stds),
     )
 
 
 # called by in memory dataset
 def prepare_patches_unsupervised_array(
-    data: np.ndarray,
+    data: NDArray,
     axes: str,
-    patch_size: Union[List[int], Tuple[int]],
-) -> Tuple[np.ndarray, None, float, float]:
+    patch_size: Union[list[int], tuple[int]],
+) -> PatchedOutput:
     """
     Iterate over data source and create an array of patches.
 
@@ -229,26 +275,25 @@ def prepare_patches_unsupervised_array(
 
     Parameters
     ----------
-    data : np.ndarray
+    data : numpy.ndarray
         Input data array.
     axes : str
         Axes of the data.
-    patch_size : Union[List[int], Tuple[int]]
+    patch_size : list or tuple of int
         Size of the patches.
 
     Returns
     -------
-    Tuple[np.ndarray, None, float, float]
-        Source patches, mean and standard deviation.
+    PatchedOutput
+        Dataclass holding the patches and their statistics.
     """
-    # calculate mean and std
-    mean = data.mean()
-    std = data.std()
-
     # reshape array
     reshaped_sample = reshape_array(data, axes)
 
+    # calculate mean and std
+    means, stds = compute_normalization_stats(reshaped_sample)
+
     # generate patches, return a generator
     patches, _ = extract_patches_sequential(reshaped_sample, patch_size=patch_size)
 
-    return patches, _, mean, std  # TODO inelegant, replace by dataclass?
+    return PatchedOutput(patches, None, Stats(means, stds), Stats((), ()))

careamics/dataset/patching/random_patching.py

@@ -13,6 +13,7 @@ def extract_patches_random(
     arr: np.ndarray,
     patch_size: Union[List[int], Tuple[int, ...]],
     target: Optional[np.ndarray] = None,
+    seed: Optional[int] = None,
 ) -> Generator[Tuple[np.ndarray, Optional[np.ndarray]], None, None]:
     """
     Generate patches from an array in a random manner.
@@ -34,12 +35,16 @@ def extract_patches_random(
         Patch sizes in each dimension.
     target : Optional[np.ndarray], optional
         Target array, by default None.
+    seed : Optional[int], optional
+        Random seed, by default None.
 
     Yields
     ------
     Generator[np.ndarray, None, None]
         Generator of patches.
     """
+    rng = np.random.default_rng(seed=seed)
+
     is_3d_patch = len(patch_size) == 3
 
     # patches sanity check
@@ -48,9 +53,6 @@ def extract_patches_random(
     # Update patch size to encompass S and C dimensions
     patch_size = [1, arr.shape[1], *patch_size]
 
-    # random generator
-    rng = np.random.default_rng()
-
     # iterate over the number of samples (S or T)
     for sample_idx in range(arr.shape[0]):
         # get sample array
@@ -113,6 +115,7 @@ def extract_patches_random_from_chunks(
     patch_size: Union[List[int], Tuple[int, ...]],
     chunk_size: Union[List[int], Tuple[int, ...]],
     chunk_limit: Optional[int] = None,
+    seed: Optional[int] = None,
 ) -> Generator[np.ndarray, None, None]:
     """
     Generate patches from an array in a random manner.
@@ -130,6 +133,8 @@ def extract_patches_random_from_chunks(
         Chunk sizes to load from the.
     chunk_limit : Optional[int], optional
         Number of chunks to load, by default None.
+    seed : Optional[int], optional
+        Random seed, by default None.
 
     Yields
     ------
@@ -141,7 +146,7 @@ def extract_patches_random_from_chunks(
     # Patches sanity check
     validate_patch_dimensions(arr, patch_size, is_3d_patch)
 
-    rng = np.random.default_rng()
+    rng = np.random.default_rng(seed=seed)
     num_chunks = chunk_limit if chunk_limit else np.prod(arr._cdata_shape)
 
     # Iterate over num chunks in the array

careamics/dataset/patching/validate_patch_dimension.py

@@ -45,18 +45,20 @@ def validate_patch_dimensions(
     if len(patch_size) != len(arr.shape[2:]):
         raise ValueError(
             f"There must be a patch size for each spatial dimensions "
-            f"(got {patch_size} patches for dims {arr.shape})."
+            f"(got {patch_size} patches for dims {arr.shape}). Check the axes order."
         )
 
     # Sanity checks on patch sizes versus array dimension
     if is_3d_patch and patch_size[0] > arr.shape[-3]:
         raise ValueError(
             f"Z patch size is inconsistent with image shape "
-            f"(got {patch_size[0]} patches for dim {arr.shape[1]})."
+            f"(got {patch_size[0]} patches for dim {arr.shape[1]}). Check the axes "
+            f"order."
         )
 
     if patch_size[-2] > arr.shape[-2] or patch_size[-1] > arr.shape[-1]:
         raise ValueError(
             f"At least one of YX patch dimensions is larger than the corresponding "
-            f"image dimension (got {patch_size} patches for dims {arr.shape[-2:]})."
+            f"image dimension (got {patch_size} patches for dims {arr.shape[-2:]}). "
+            f"Check the axes order."
         )

careamics/dataset/tiling/__init__.py

@@ -0,0 +1,10 @@
+"""Tiling functions."""
+
+__all__ = [
+    "stitch_prediction",
+    "extract_tiles",
+    "collate_tiles",
+]
+
+from .collate_tiles import collate_tiles
+from .tiled_patching import extract_tiles

careamics/dataset/tiling/collate_tiles.py

@@ -0,0 +1,33 @@
+"""Collate function for tiling."""
+
+from typing import Any, List, Tuple
+
+import numpy as np
+from torch.utils.data.dataloader import default_collate
+
+from careamics.config.tile_information import TileInformation
+
+
+def collate_tiles(batch: List[Tuple[np.ndarray, TileInformation]]) -> Any:
+    """
+    Collate tiles received from CAREamics prediction dataloader.
+
+    CAREamics prediction dataloader returns tuples of arrays and TileInformation. In
+    case of non-tiled data, this function will return the arrays. In case of tiled data,
+    it will return the arrays, the last tile flag, the overlap crop coordinates and the
+    stitch coordinates.
+
+    Parameters
+    ----------
+    batch : List[Tuple[np.ndarray, TileInformation], ...]
+        Batch of tiles.
+
+    Returns
+    -------
+    Any
+        Collated batch.
+    """
+    new_batch = [tile for tile, _ in batch]
+    tiles_batch = [tile_info for _, tile_info in batch]
+
+    return default_collate(new_batch), tiles_batch

careamics/dataset/{patching → tiling}/tiled_patching.py

@@ -84,15 +84,15 @@ def extract_tiles(
     tile_size: Union[List[int], Tuple[int, ...]],
     overlaps: Union[List[int], Tuple[int, ...]],
 ) -> Generator[Tuple[np.ndarray, TileInformation], None, None]:
-    """
-    Generate tiles from the input array with specified overlap.
+    """Generate tiles from the input array with specified overlap.
 
     The tiles cover the whole array. The method returns a generator that yields
     tuples of array and tile information, the latter includes whether
     the tile is the last one, the coordinates of the overlap crop, and the coordinates
     of the stitched tile.
 
-    The array has shape C(Z)YX, where C can be a singleton.
+    Input array should have shape SC(Z)YX, while the returned tiles have shape C(Z)YX,
+    where C can be a singleton.
 
     Parameters
     ----------
@@ -155,10 +155,10 @@ def extract_tiles(
             # create tile information
             tile_info = TileInformation(
                 array_shape=sample.squeeze().shape,
-                tiled=True,
                 last_tile=last_tile,
                 overlap_crop_coords=overlap_crop_coords,
                 stitch_coords=stitch_coords,
+                sample_id=sample_idx,
             )
 
             yield tile, tile_info

careamics/file_io/__init__.py

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

careamics/file_io/read/__init__.py

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

careamics/file_io/read/get_func.py

@@ -0,0 +1,56 @@
+"""Module to get read functions."""
+
+from pathlib import Path
+from typing import Callable, Dict, 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.")