careamics 0.1.0rc2__py3-none-any.whl → 0.1.0rc3__py3-none-any.whl

careamics/dataset/patching/random_patching.py

@@ -0,0 +1,190 @@
+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,
+) -> 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.
+
+    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)
+
+    # 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
+        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,
+) -> 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 : Tuple[int]
+        Patch sizes in each dimension.
+    chunk_size : Tuple[int]
+        Chunk sizes to load from the.
+
+    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()
+    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,206 @@
+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 : Tuple[int, ...]
+        Shape of the input array.
+    patch_sizes : 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 : Tuple[int, ...]
+        Input array shape.
+    patch_sizes : 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.
+
+    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 = [arr.shape[0], -1, arr.shape[2], *output_shape[2:]]
+
+    patches = view_as_windows(arr, window_shape=window_shape, step=step).reshape(
+        *output_shape
+    )
+    if target is not None:
+        rng.shuffle(patches, axis=1)
+    else:
+        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.
+
+    Returns
+    -------
+    Generator[Tuple[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)
+
+    # 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

careamics/dataset/patching/tiled_patching.py

@@ -0,0 +1,158 @@
+import itertools
+from typing import Generator, List, Tuple, Union
+
+import numpy as np
+
+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, ...]], ...]:
+    """
+    Compute the coordinates of each tile along an axis, given the overlap.
+
+    Parameters
+    ----------
+    axis_size : int
+        Length of the axis.
+    tile_size : int
+        Size of the tile for the given axis.
+    overlap : int
+        Size of the overlap for the given axis.
+
+    Returns
+    -------
+    Tuple[Tuple[int, ...], ...]
+        Tuple of all coordinates for given axis.
+    """
+    # Compute the step between tiles
+    step = tile_size - overlap
+    crop_coords = []
+    stitch_coords = []
+    overlap_crop_coords = []
+
+    # Iterate over the axis with step
+    for i in range(0, max(1, axis_size - overlap), step):
+        # Check if the tile fits within the axis
+        if i + tile_size <= axis_size:
+            # Add the coordinates to crop one tile
+            crop_coords.append((i, i + tile_size))
+
+            # Add the pixel coordinates of the cropped tile in the original image space
+            stitch_coords.append(
+                (
+                    i + overlap // 2 if i > 0 else 0,
+                    i + tile_size - overlap // 2
+                    if crop_coords[-1][1] < axis_size
+                    else axis_size,
+                )
+            )
+
+            # Add the coordinates to crop the overlap from the prediction.
+            overlap_crop_coords.append(
+                (
+                    overlap // 2 if i > 0 else 0,
+                    tile_size - overlap // 2
+                    if crop_coords[-1][1] < axis_size
+                    else tile_size,
+                )
+            )
+
+        # If the tile does not fit within the axis, perform the abovementioned
+        # operations starting from the end of the axis
+        else:
+            # if (axis_size - tile_size, axis_size) not in crop_coords:
+            crop_coords.append((max(0, axis_size - tile_size), axis_size))
+            last_tile_end_coord = stitch_coords[-1][1] if stitch_coords else 1
+            stitch_coords.append((last_tile_end_coord, axis_size))
+            overlap_crop_coords.append(
+                (tile_size - (axis_size - last_tile_end_coord), tile_size)
+            )
+            break
+    return crop_coords, stitch_coords, overlap_crop_coords
+
+
+def extract_tiles(
+    arr: np.ndarray,
+    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.
+
+    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.
+
+    Parameters
+    ----------
+    arr : np.ndarray
+        Array of shape (S, C, (Z), Y, X).
+    tile_size : Union[List[int], Tuple[int]]
+        Tile sizes in each dimension, of length 2 or 3.
+    overlaps : Union[List[int], Tuple[int]]
+        Overlap values in each dimension, of length 2 or 3.
+
+    Yields
+    ------
+    Generator[Tuple[np.ndarray, TileInformation], None, None]
+        Tile generator, yields the tile and additional information.
+    """
+    # Iterate over num samples (S)
+    for sample_idx in range(arr.shape[0]):
+        sample: np.ndarray = arr[sample_idx, ...]
+
+        # Create a list of coordinates for cropping and stitching all axes.
+        # [crop coordinates, stitching coordinates, overlap crop coordinates]
+        # For axis of size 35 and patch size of 32 compute_crop_and_stitch_coords_1d
+        # will output ([(0, 32), (3, 35)], [(0, 20), (20, 35)], [(0, 20), (17, 32)])
+        crop_and_stitch_coords_list = [
+            _compute_crop_and_stitch_coords_1d(
+                sample.shape[i + 1], tile_size[i], overlaps[i]
+            )
+            for i in range(len(tile_size))
+        ]
+
+        # Rearrange crop coordinates from a list of coordinate pairs per axis to a list
+        # grouped by type.
+        all_crop_coords, all_stitch_coords, all_overlap_crop_coords = zip(
+            *crop_and_stitch_coords_list
+        )
+
+        # Maximum tile index
+        max_tile_idx = np.prod([len(axis) for axis in all_crop_coords]) - 1
+
+        # Iterate over generated coordinate pairs:
+        for tile_idx, (crop_coords, stitch_coords, overlap_crop_coords) in enumerate(
+            zip(
+                itertools.product(*all_crop_coords),
+                itertools.product(*all_stitch_coords),
+                itertools.product(*all_overlap_crop_coords),
+            )
+        ):
+            # Extract tile from the sample
+            tile: np.ndarray = sample[
+                (..., *[slice(c[0], c[1]) for c in list(crop_coords)])  # type: ignore
+            ]
+
+            # Check if we are at the end of the sample by computing the length of the
+            # array that contains all the tiles
+            if tile_idx == max_tile_idx:
+                last_tile = True
+            else:
+                last_tile = False
+
+            # 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,
+            )
+
+            yield tile, tile_info

careamics/dataset/patching/validate_patch_dimension.py

@@ -0,0 +1,60 @@
+from typing import List, Tuple, Union
+
+import numpy as np
+
+
+def validate_patch_dimensions(
+    arr: np.ndarray,
+    patch_size: Union[List[int], Tuple[int, ...]],
+    is_3d_patch: bool,
+) -> None:
+    """
+    Check patch size and array compatibility.
+
+    This method validates the patch sizes with respect to the array dimensions:
+
+    - Patch must have two dimensions fewer than the array (S and C).
+    - Patch sizes are smaller than the corresponding array dimensions.
+
+    If one of these conditions is not met, a ValueError is raised.
+
+    This method should be called after inputs have been resized.
+
+    Parameters
+    ----------
+    arr : np.ndarray
+        Input array.
+    patch_size : Union[List[int], Tuple[int, ...]]
+        Size of the patches along each dimension of the array, except the first.
+    is_3d_patch : bool
+        Whether the patch is 3D or not.
+
+    Raises
+    ------
+    ValueError
+        If the patch size is not consistent with the array shape (one more array
+        dimension).
+    ValueError
+        If the patch size in Z is larger than the array dimension.
+    ValueError
+        If either of the patch sizes in X or Y is larger than the corresponding array
+        dimension.
+    """
+    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})."
+        )
+
+    # 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]})."
+        )
+
+    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:]})."
+        )