careamics 0.1.0rc5__py3-none-any.whl → 0.1.0rc7__py3-none-any.whl

careamics/dataset/iterable_tiled_pred_dataset.py

@@ -0,0 +1,139 @@
+"""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.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, read_tiff
+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,9 +1,6 @@
-"""
-Tiling submodule.
-
-These functions are used to tile images into patches or tiles.
-"""
+"""Patching functions."""
 
+from dataclasses import dataclass
 from pathlib import Path
 from typing import Callable, List, Tuple, Union
 
@@ -11,30 +8,69 @@ import numpy as np
 
 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[np.ndarray, tuple, list, None]
+    stds: Union[np.ndarray, tuple, list, None]
+
+
+@dataclass
+class PatchedOutput:
+    """Dataclass to store patches and statistics."""
+
+    patches: Union[np.ndarray]
+    targets: Union[np.ndarray, None]
+    image_stats: Stats
+    target_stats: Stats
+
+
+@dataclass
+class StatsOutput:
+    """Dataclass to store patches and statistics."""
+
+    image_stats: Stats
+    target_stats: Stats
+
+
 # called by in memory dataset
 def prepare_patches_supervised(
     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]
+        List of paths to training data.
+    target_files : List[Path]
+        List of paths to target data.
+    axes : str
+        Axes of the data.
+    patch_size : Union[List[int], Tuple[int]]
+        Size of the patches.
+    read_source_func : Callable
+        Function to read the data.
+
     Returns
     -------
     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):
@@ -74,17 +110,18 @@ 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),
     )
 
 
@@ -94,14 +131,26 @@ def prepare_patches_unsupervised(
     axes: str,
     patch_size: Union[List[int], Tuple[int]],
     read_source_func: Callable,
-) -> Tuple[np.ndarray, None, float, float]:
-    """
-    Iterate over data source and create an array of patches.
+) -> 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]
+        List of paths to training data.
+    axes : str
+        Axes of the data.
+    patch_size : Union[List[int], Tuple[int]]
+        Size of the patches.
+    read_source_func : Callable
+        Function to read the data.
 
     Returns
     -------
-    np.ndarray
-        Array of patches.
+    Tuple[np.ndarray, None, float, float]
+        Source and target patches, mean and standard deviation.
     """
     means, stds, num_samples = 0, 0, 0
     all_patches = []
@@ -128,12 +177,14 @@ 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
@@ -142,7 +193,7 @@ def prepare_patches_supervised_array(
     axes: str,
     data_target: np.ndarray,
     patch_size: Union[List[int], Tuple[int]],
-) -> Tuple[np.ndarray, np.ndarray, float, float]:
+) -> 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
@@ -150,19 +201,30 @@ def prepare_patches_supervised_array(
 
     Patches returned are of shape SC(Z)YX, where S is now the patches dimension.
 
+    Parameters
+    ----------
+    data : np.ndarray
+        Input data array.
+    axes : str
+        Axes of the data.
+    data_target : np.ndarray
+        Target data array.
+    patch_size : Union[List[int], Tuple[int]]
+        Size of the patches.
+
     Returns
     -------
-    np.ndarray
-        Array of patches.
+    Tuple[np.ndarray, np.ndarray, float, float]
+        Source and target patches, mean and standard deviation.
     """
-    # 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
@@ -173,11 +235,11 @@ 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),
     )
 
 
@@ -186,7 +248,7 @@ def prepare_patches_unsupervised_array(
     data: np.ndarray,
     axes: str,
     patch_size: Union[List[int], Tuple[int]],
-) -> Tuple[np.ndarray, None, float, float]:
+) -> PatchedOutput:
     """
     Iterate over data source and create an array of patches.
 
@@ -195,19 +257,27 @@ def prepare_patches_unsupervised_array(
 
     Patches returned are of shape SC(Z)YX, where S is now the patches dimension.
 
+    Parameters
+    ----------
+    data : np.ndarray
+        Input data array.
+    axes : str
+        Axes of the data.
+    patch_size : Union[List[int], Tuple[int]]
+        Size of the patches.
+
     Returns
     -------
-    np.ndarray
-        Array of patches.
+    Tuple[np.ndarray, None, float, float]
+        Source patches, mean and standard deviation.
     """
-    # 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

@@ -1,3 +1,5 @@
+"""Random patching utilities."""
+
 from typing import Generator, List, Optional, Tuple, Union
 
 import numpy as np
@@ -11,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.
@@ -30,12 +33,18 @@ def extract_patches_random(
         Input image array.
     patch_size : Tuple[int]
         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
@@ -44,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
@@ -109,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.
@@ -120,10 +127,14 @@ def extract_patches_random_from_chunks(
     ----------
     arr : np.ndarray
         Input image array.
-    patch_size : Tuple[int]
+    patch_size : Union[List[int], Tuple[int, ...]]
         Patch sizes in each dimension.
-    chunk_size : Tuple[int]
+    chunk_size : Union[List[int], Tuple[int, ...]]
         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
     ------
@@ -135,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/sequential_patching.py

@@ -1,3 +1,5 @@
+"""Sequential patching functions."""
+
 from typing import List, Optional, Tuple, Union
 
 import numpy as np
@@ -14,14 +16,14 @@ def _compute_number_of_patches(
 
     Parameters
     ----------
-    arr : Tuple[int, ...]
+    arr_shape : Tuple[int, ...]
         Shape of the input array.
-    patch_sizes : Tuple[int]
+    patch_sizes : Union[List[int], Tuple[int, ...]
         Shape of the patches.
 
     Returns
     -------
-    Tuple[int]
+    Tuple[int, ...]
         Number of patches in each dimension.
     """
     if len(arr_shape) != len(patch_sizes):
@@ -55,14 +57,14 @@ def _compute_overlap(
 
     Parameters
     ----------
-    arr : Tuple[int, ...]
+    arr_shape : Tuple[int, ...]
         Input array shape.
-    patch_sizes : Tuple[int]
+    patch_sizes : Union[List[int], Tuple[int, ...]]
         Size of the patches.
 
     Returns
     -------
-    Tuple[int]
+    Tuple[int, ...]
         Overlap between patches in each dimension.
     """
     n_patches = _compute_number_of_patches(arr_shape, patch_sizes)
@@ -123,6 +125,8 @@ def _compute_patch_views(
         Steps between views.
     output_shape : Tuple[int]
         Shape of the output array.
+    target : Optional[np.ndarray], optional
+        Target array, by default None.
 
     Returns
     -------
@@ -161,11 +165,13 @@ def extract_patches_sequential(
         Input image array.
     patch_size : Tuple[int]
         Patch sizes in each dimension.
+    target : Optional[np.ndarray], optional
+        Target array, by default None.
 
     Returns
     -------
-    Generator[Tuple[np.ndarray, ...], None, None]
-        Generator of patches.
+    Tuple[np.ndarray, Optional[np.ndarray]]
+        Patches.
     """
     is_3d_patch = len(patch_size) == 3
 

careamics/dataset/patching/validate_patch_dimension.py

@@ -1,3 +1,5 @@
+"""Patch validation functions."""
+
 from typing import List, Tuple, Union
 
 import numpy as np
@@ -43,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

@@ -1,3 +1,5 @@
+"""Tiled patching utilities."""
+
 import itertools
 from typing import Generator, List, Tuple, Union
 
@@ -8,7 +10,7 @@ from careamics.config.tile_information import TileInformation
 
 def _compute_crop_and_stitch_coords_1d(
     axis_size: int, tile_size: int, overlap: int
-) -> Tuple[List[Tuple[int, ...]], ...]:
+) -> Tuple[List[Tuple[int, int]], List[Tuple[int, int]], List[Tuple[int, int]]]:
     """
     Compute the coordinates of each tile along an axis, given the overlap.
 
@@ -82,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
     ----------
@@ -153,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/dataset/zarr_dataset.py

@@ -1,3 +1,5 @@
+"""Zarr dataset."""
+
 # from itertools import islice
 # from typing import Callable, Dict, List, Optional, Tuple, Union