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

careamics/dataset/patching/patching.py

@@ -0,0 +1,299 @@
+"""Patching functions."""
+
+from dataclasses import dataclass
+from pathlib import Path
+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],
+    axes: str,
+    patch_size: Union[list[int], tuple[int, ...]],
+    read_source_func: Callable,
+) -> 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 of pathlib.Path
+        List of paths to training data.
+    target_files : list of pathlib.Path
+        List of paths to target data.
+    axes : str
+        Axes of the data.
+    patch_size : list or tuple of int
+        Size of the patches.
+    read_source_func : Callable
+        Function to read the data.
+
+    Returns
+    -------
+    np.ndarray
+        Array of patches.
+    """
+    means, stds, num_samples = 0, 0, 0
+    all_patches, all_targets = [], []
+    for train_filename, target_filename in zip(train_files, target_files):
+        try:
+            sample: np.ndarray = read_source_func(train_filename, axes)
+            target: np.ndarray = read_source_func(target_filename, axes)
+            means += sample.mean()
+            stds += sample.std()
+            num_samples += 1
+
+            # reshape array
+            sample = reshape_array(sample, axes)
+            target = reshape_array(target, axes)
+
+            # generate patches, return a generator
+            patches, targets = extract_patches_sequential(
+                sample, patch_size=patch_size, target=target
+            )
+
+            # convert generator to list and add to all_patches
+            all_patches.append(patches)
+
+            # ensure targets are not None (type checking)
+            if targets is not None:
+                all_targets.append(targets)
+            else:
+                raise ValueError(f"No target found for {target_filename}.")
+
+        except Exception as e:
+            # emit warning and continue
+            logger.error(f"Failed to read {train_filename} or {target_filename}: {e}")
+
+    # raise error if no valid samples found
+    if num_samples == 0 or len(all_patches) == 0:
+        raise ValueError(
+            f"No valid samples found in the input data: {train_files} and "
+            f"{target_files}."
+        )
+
+    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 PatchedOutput(
+        patch_array,
+        target_array,
+        Stats(image_means, image_stds),
+        Stats(target_means, target_stds),
+    )
+
+
+# called by in_memory_dataset
+def prepare_patches_unsupervised(
+    train_files: list[Path],
+    axes: str,
+    patch_size: Union[list[int], tuple[int]],
+    read_source_func: Callable,
+) -> 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 of pathlib.Path
+        List of paths to training data.
+    axes : str
+        Axes of the data.
+    patch_size : list or tuple of int
+        Size of the patches.
+    read_source_func : Callable
+        Function to read the data.
+
+    Returns
+    -------
+    PatchedOutput
+        Dataclass holding patches and their statistics.
+    """
+    means, stds, num_samples = 0, 0, 0
+    all_patches = []
+    for filename in train_files:
+        try:
+            sample: np.ndarray = read_source_func(filename, axes)
+            means += sample.mean()
+            stds += sample.std()
+            num_samples += 1
+
+            # reshape array
+            sample = reshape_array(sample, axes)
+
+            # generate patches, return a generator
+            patches, _ = extract_patches_sequential(sample, patch_size=patch_size)
+
+            # convert generator to list and add to all_patches
+            all_patches.append(patches)
+        except Exception as e:
+            # emit warning and continue
+            logger.error(f"Failed to read {filename}: {e}")
+
+    # raise error if no valid samples found
+    if num_samples == 0:
+        raise ValueError(f"No valid samples found in the input data: {train_files}.")
+
+    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 PatchedOutput(
+        patch_array, None, Stats(image_means, image_stds), Stats((), ())
+    )
+
+
+# called on arrays by in memory dataset
+def prepare_patches_supervised_array(
+    data: NDArray,
+    axes: str,
+    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
+    dimensions.
+
+    Patches returned are of shape SC(Z)YX, where S is now the patches dimension.
+
+    Parameters
+    ----------
+    data : numpy.ndarray
+        Input data array.
+    axes : str
+        Axes of the data.
+    data_target : numpy.ndarray
+        Target data array.
+    patch_size : list or tuple of int
+        Size of the patches.
+
+    Returns
+    -------
+    PatchedOutput
+        Dataclass holding the source and target patches, with their statistics.
+    """
+    # 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
+    )
+
+    if patch_targets is None:
+        raise ValueError("No target extracted.")
+
+    logger.info(f"Extracted {patches.shape[0]} patches from input array.")
+
+    return PatchedOutput(
+        patches,
+        patch_targets,
+        Stats(image_means, image_stds),
+        Stats(target_means, target_stds),
+    )
+
+
+# called by in memory dataset
+def prepare_patches_unsupervised_array(
+    data: NDArray,
+    axes: str,
+    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
+    dimensions.
+
+    Patches returned are of shape SC(Z)YX, where S is now the patches dimension.
+
+    Parameters
+    ----------
+    data : numpy.ndarray
+        Input data array.
+    axes : str
+        Axes of the data.
+    patch_size : list or tuple of int
+        Size of the patches.
+
+    Returns
+    -------
+    PatchedOutput
+        Dataclass holding the patches and their statistics.
+    """
+    # 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 PatchedOutput(patches, None, Stats(means, stds), Stats((), ()))

careamics/dataset/patching/random_patching.py

@@ -0,0 +1,201 @@
+"""Random patching utilities."""
+
+from typing import Generator, List, Optional, Tuple, Union
+
+import numpy as np
+import zarr
+
+from .validate_patch_dimension import validate_patch_dimensions
+
+
+# TOOD split in testable functions
+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.
+
+    The method calculates how many patches the image can be divided into and then
+    extracts an equal number of random patches.
+
+    It returns a generator that yields the following:
+
+    - patch: np.ndarray, dimension C(Z)YX.
+    - target_patch: np.ndarray, dimension C(Z)YX, if the target is present, None
+        otherwise.
+
+    Parameters
+    ----------
+    arr : np.ndarray
+        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
+    validate_patch_dimensions(arr, patch_size, is_3d_patch)
+
+    # Update patch size to encompass S and C dimensions
+    patch_size = [1, arr.shape[1], *patch_size]
+
+    # iterate over the number of samples (S or T)
+    for sample_idx in range(arr.shape[0]):
+        # get sample array
+        sample: np.ndarray = arr[sample_idx, ...]
+
+        # same for target
+        if target is not None:
+            target_sample: np.ndarray = target[sample_idx, ...]
+
+        # calculate the number of patches
+        n_patches = np.ceil(np.prod(sample.shape) / np.prod(patch_size)).astype(int)
+
+        # iterate over the number of patches
+        for _ in range(n_patches):
+            # get crop coordinates
+            crop_coords = [
+                rng.integers(0, sample.shape[i] - patch_size[1:][i], endpoint=True)
+                for i in range(len(patch_size[1:]))
+            ]
+
+            # extract patch
+            patch = (
+                sample[
+                    (
+                        ...,  # type: ignore
+                        *[  # type: ignore
+                            slice(c, c + patch_size[1:][i])
+                            for i, c in enumerate(crop_coords)
+                        ],
+                    )
+                ]
+                .copy()
+                .astype(np.float32)
+            )
+
+            # same for target
+            if target is not None:
+                target_patch = (
+                    target_sample[
+                        (
+                            ...,  # type: ignore
+                            *[  # type: ignore
+                                slice(c, c + patch_size[1:][i])
+                                for i, c in enumerate(crop_coords)
+                            ],
+                        )
+                    ]
+                    .copy()
+                    .astype(np.float32)
+                )
+                # return patch and target patch
+                yield patch, target_patch
+            else:
+                # return patch
+                yield patch, None
+
+
+def extract_patches_random_from_chunks(
+    arr: zarr.Array,
+    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.
+
+    The method calculates how many patches the image can be divided into and then
+    extracts an equal number of random patches.
+
+    Parameters
+    ----------
+    arr : np.ndarray
+        Input image array.
+    patch_size : Union[List[int], Tuple[int, ...]]
+        Patch sizes in each dimension.
+    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
+    ------
+    Generator[np.ndarray, None, None]
+        Generator of patches.
+    """
+    is_3d_patch = len(patch_size) == 3
+
+    # Patches sanity check
+    validate_patch_dimensions(arr, patch_size, is_3d_patch)
+
+    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
+    for _ in range(num_chunks):
+        chunk_crop_coords = [
+            rng.integers(0, max(0, arr.shape[i] - chunk_size[i]), endpoint=True)
+            for i in range(len(chunk_size))
+        ]
+        chunk = arr[
+            (
+                ...,
+                *[slice(c, c + chunk_size[i]) for i, c in enumerate(chunk_crop_coords)],
+            )
+        ].squeeze()
+
+        # Add a singleton dimension if the chunk does not have a sample dimension
+        if len(chunk.shape) == len(patch_size):
+            chunk = np.expand_dims(chunk, axis=0)
+
+        # Iterate over num samples (S)
+        for sample_idx in range(chunk.shape[0]):
+            spatial_chunk = chunk[sample_idx]
+            assert len(spatial_chunk.shape) == len(
+                patch_size
+            ), "Requested chunk shape is not equal to patch size"
+
+            n_patches = np.ceil(
+                np.prod(spatial_chunk.shape) / np.prod(patch_size)
+            ).astype(int)
+
+            # Iterate over the number of patches
+            for _ in range(n_patches):
+                patch_crop_coords = [
+                    rng.integers(
+                        0, spatial_chunk.shape[i] - patch_size[i], endpoint=True
+                    )
+                    for i in range(len(patch_size))
+                ]
+                patch = (
+                    spatial_chunk[
+                        (
+                            ...,
+                            *[
+                                slice(c, c + patch_size[i])
+                                for i, c in enumerate(patch_crop_coords)
+                            ],
+                        )
+                    ]
+                    .copy()
+                    .astype(np.float32)
+                )
+                yield patch

careamics/dataset/patching/sequential_patching.py

@@ -0,0 +1,212 @@
+"""Sequential patching functions."""
+
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+from skimage.util import view_as_windows
+
+from .validate_patch_dimension import validate_patch_dimensions
+
+
+def _compute_number_of_patches(
+    arr_shape: Tuple[int, ...], patch_sizes: Union[List[int], Tuple[int, ...]]
+) -> Tuple[int, ...]:
+    """
+    Compute the number of patches that fit in each dimension.
+
+    Parameters
+    ----------
+    arr_shape : Tuple[int, ...]
+        Shape of the input array.
+    patch_sizes : Union[List[int], Tuple[int, ...]
+        Shape of the patches.
+
+    Returns
+    -------
+    Tuple[int, ...]
+        Number of patches in each dimension.
+    """
+    if len(arr_shape) != len(patch_sizes):
+        raise ValueError(
+            f"Array shape {arr_shape} and patch size {patch_sizes} should have the "
+            f"same dimension, including singleton dimension for S and equal dimension "
+            f"for C."
+        )
+
+    try:
+        n_patches = [
+            np.ceil(arr_shape[i] / patch_sizes[i]).astype(int)
+            for i in range(len(patch_sizes))
+        ]
+    except IndexError as e:
+        raise ValueError(
+            f"Patch size {patch_sizes} is not compatible with array shape {arr_shape}"
+        ) from e
+
+    return tuple(n_patches)
+
+
+def _compute_overlap(
+    arr_shape: Tuple[int, ...], patch_sizes: Union[List[int], Tuple[int, ...]]
+) -> Tuple[int, ...]:
+    """
+    Compute the overlap between patches in each dimension.
+
+    If the array dimensions are divisible by the patch sizes, then the overlap is
+    0. Otherwise, it is the result of the division rounded to the upper value.
+
+    Parameters
+    ----------
+    arr_shape : Tuple[int, ...]
+        Input array shape.
+    patch_sizes : Union[List[int], Tuple[int, ...]]
+        Size of the patches.
+
+    Returns
+    -------
+    Tuple[int, ...]
+        Overlap between patches in each dimension.
+    """
+    n_patches = _compute_number_of_patches(arr_shape, patch_sizes)
+
+    overlap = [
+        np.ceil(
+            np.clip(n_patches[i] * patch_sizes[i] - arr_shape[i], 0, None)
+            / max(1, (n_patches[i] - 1))
+        ).astype(int)
+        for i in range(len(patch_sizes))
+    ]
+    return tuple(overlap)
+
+
+def _compute_patch_steps(
+    patch_sizes: Union[List[int], Tuple[int, ...]], overlaps: Tuple[int, ...]
+) -> Tuple[int, ...]:
+    """
+    Compute steps between patches.
+
+    Parameters
+    ----------
+    patch_sizes : Tuple[int]
+        Size of the patches.
+    overlaps : Tuple[int]
+        Overlap between patches.
+
+    Returns
+    -------
+    Tuple[int]
+        Steps between patches.
+    """
+    steps = [
+        min(patch_sizes[i] - overlaps[i], patch_sizes[i])
+        for i in range(len(patch_sizes))
+    ]
+    return tuple(steps)
+
+
+# TODO why stack the target here and not on a different dimension before this function?
+def _compute_patch_views(
+    arr: np.ndarray,
+    window_shape: List[int],
+    step: Tuple[int, ...],
+    output_shape: List[int],
+    target: Optional[np.ndarray] = None,
+) -> np.ndarray:
+    """
+    Compute views of an array corresponding to patches.
+
+    Parameters
+    ----------
+    arr : np.ndarray
+        Array from which the views are extracted.
+    window_shape : Tuple[int]
+        Shape of the views.
+    step : Tuple[int]
+        Steps between views.
+    output_shape : Tuple[int]
+        Shape of the output array.
+    target : Optional[np.ndarray], optional
+        Target array, by default None.
+
+    Returns
+    -------
+    np.ndarray
+        Array with views dimension.
+    """
+    rng = np.random.default_rng()
+
+    if target is not None:
+        arr = np.stack([arr, target], axis=0)
+        window_shape = [arr.shape[0], *window_shape]
+        step = (arr.shape[0], *step)
+        output_shape = [-1, arr.shape[0], arr.shape[2], *output_shape[2:]]
+
+    patches = view_as_windows(arr, window_shape=window_shape, step=step).reshape(
+        *output_shape
+    )
+    rng.shuffle(patches, axis=0)
+    return patches
+
+
+def extract_patches_sequential(
+    arr: np.ndarray,
+    patch_size: Union[List[int], Tuple[int, ...]],
+    target: Optional[np.ndarray] = None,
+) -> Tuple[np.ndarray, Optional[np.ndarray]]:
+    """
+    Generate patches from an array in a sequential manner.
+
+    Array dimensions should be SC(Z)YX, where S and C can be singleton dimensions. The
+    patches are generated sequentially and cover the whole array.
+
+    Parameters
+    ----------
+    arr : np.ndarray
+        Input image array.
+    patch_size : Tuple[int]
+        Patch sizes in each dimension.
+    target : Optional[np.ndarray], optional
+        Target array, by default None.
+
+    Returns
+    -------
+    Tuple[np.ndarray, Optional[np.ndarray]]
+        Patches.
+    """
+    is_3d_patch = len(patch_size) == 3
+
+    # Patches sanity check
+    validate_patch_dimensions(arr, patch_size, is_3d_patch)
+
+    # Update patch size to encompass S and C dimensions
+    patch_size = [1, arr.shape[1], *patch_size]
+
+    # Compute overlap
+    overlaps = _compute_overlap(arr_shape=arr.shape, patch_sizes=patch_size)
+
+    # Create view window and overlaps
+    window_steps = _compute_patch_steps(patch_sizes=patch_size, overlaps=overlaps)
+
+    output_shape = [
+        -1,
+    ] + patch_size[1:]
+
+    # Generate a view of the input array containing pre-calculated number of patches
+    # in each dimension with overlap.
+    # Resulting array is resized to (n_patches, C, Z, Y, X) or (n_patches, C, Y, X)
+    patches = _compute_patch_views(
+        arr,
+        window_shape=patch_size,
+        step=window_steps,
+        output_shape=output_shape,
+        target=target,
+    )
+
+    if target is not None:
+        # target was concatenated to patches in _compute_reshaped_view
+        return (
+            patches[:, 0, ...],
+            patches[:, 1, ...],
+        )  # TODO  in _compute_reshaped_view?
+    else:
+        return patches, None