careamics 0.0.1__py3-none-any.whl → 0.1.0rc2__py3-none-any.whl

careamics/dataset/in_memory_dataset.py

@@ -0,0 +1,202 @@
+"""In-memory dataset module."""
+from pathlib import Path
+from typing import Callable, Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+
+from careamics.utils import normalize
+from careamics.utils.logging import get_logger
+
+from .dataset_utils import (
+    list_files,
+    read_tiff,
+)
+from .extraction_strategy import ExtractionStrategy
+from .patching import generate_patches
+
+logger = get_logger(__name__)
+
+
+class InMemoryDataset(torch.utils.data.Dataset):
+    """
+    Dataset storing data in memory and allowing generating patches from it.
+
+    Parameters
+    ----------
+    data_path : Union[str, Path]
+        Path to the data, must be a directory.
+    data_format : str
+        Extension of the data files, without period.
+    axes : str
+        Description of axes in format STCZYX.
+    patch_extraction_method : ExtractionStrategies
+        Patch extraction strategy, as defined in extraction_strategy.
+    patch_size : Union[List[int], Tuple[int]]
+        Size of the patches along each axis, must be of dimension 2 or 3.
+    patch_overlap : Optional[Union[List[int], Tuple[int]]], optional
+        Overlap of the patches, must be of dimension 2 or 3, by default None.
+    mean : Optional[float], optional
+        Expected mean of the dataset, by default None.
+    std : Optional[float], optional
+        Expected standard deviation of the dataset, by default None.
+    patch_transform : Optional[Callable], optional
+        Patch transform to apply, by default None.
+    patch_transform_params : Optional[Dict], optional
+        Patch transform parameters, by default None.
+    """
+
+    def __init__(
+        self,
+        data_path: Union[str, Path],
+        data_format: str,
+        axes: str,
+        patch_extraction_method: ExtractionStrategy,
+        patch_size: Union[List[int], Tuple[int]],
+        patch_overlap: Optional[Union[List[int], Tuple[int]]] = None,
+        mean: Optional[float] = None,
+        std: Optional[float] = None,
+        patch_transform: Optional[Callable] = None,
+        patch_transform_params: Optional[Dict] = None,
+    ) -> None:
+        """
+        Constructor.
+
+        Parameters
+        ----------
+        data_path : Union[str, Path]
+            Path to the data, must be a directory.
+        data_format : str
+            Extension of the data files, without period.
+        axes : str
+            Description of axes in format STCZYX.
+        patch_extraction_method : ExtractionStrategies
+            Patch extraction strategy, as defined in extraction_strategy.
+        patch_size : Union[List[int], Tuple[int]]
+            Size of the patches along each axis, must be of dimension 2 or 3.
+        patch_overlap : Optional[Union[List[int], Tuple[int]]], optional
+            Overlap of the patches, must be of dimension 2 or 3, by default None.
+        mean : Optional[float], optional
+            Expected mean of the dataset, by default None.
+        std : Optional[float], optional
+            Expected standard deviation of the dataset, by default None.
+        patch_transform : Optional[Callable], optional
+            Patch transform to apply, by default None.
+        patch_transform_params : Optional[Dict], optional
+            Patch transform parameters, by default None.
+
+        Raises
+        ------
+        ValueError
+            If data_path is not a directory.
+        """
+        self.data_path = Path(data_path)
+        if not self.data_path.is_dir():
+            raise ValueError("Path to data should be an existing folder.")
+
+        self.data_format = data_format
+        self.axes = axes
+
+        self.patch_transform = patch_transform
+
+        self.files = list_files(self.data_path, self.data_format)
+
+        self.patch_size = patch_size
+        self.patch_overlap = patch_overlap
+        self.patch_extraction_method = patch_extraction_method
+        self.patch_transform = patch_transform
+        self.patch_transform_params = patch_transform_params
+
+        self.mean = mean
+        self.std = std
+
+        # Generate patches
+        self.data, computed_mean, computed_std = self._prepare_patches()
+
+        if not mean or not std:
+            self.mean, self.std = computed_mean, computed_std
+            logger.info(f"Computed dataset mean: {self.mean}, std: {self.std}")
+
+        assert self.mean is not None
+        assert self.std is not None
+
+    def _prepare_patches(self) -> Tuple[np.ndarray, float, float]:
+        """
+        Iterate over data source and create an array of patches.
+
+        Returns
+        -------
+        np.ndarray
+            Array of patches.
+        """
+        means, stds, num_samples = 0, 0, 0
+        self.all_patches = []
+        for filename in self.files:
+            sample = read_tiff(filename, self.axes)
+            means += sample.mean()
+            stds += np.std(sample)
+            num_samples += 1
+
+            # generate patches, return a generator
+            patches = generate_patches(
+                sample,
+                self.patch_extraction_method,
+                self.patch_size,
+                self.patch_overlap,
+            )
+
+            # convert generator to list and add to all_patches
+            self.all_patches.extend(list(patches))
+
+            result_mean, result_std = means / num_samples, stds / num_samples
+        return np.concatenate(self.all_patches), result_mean, result_std
+
+    def __len__(self) -> int:
+        """
+        Return the length of the dataset.
+
+        Returns
+        -------
+        int
+            Length of the dataset.
+        """
+        # convert to numpy array to convince mypy that it is not a generator
+        return sum(np.array(s).shape[0] for s in self.all_patches)
+
+    def __getitem__(self, index: int) -> Tuple[np.ndarray]:
+        """
+        Return the patch corresponding to the provided index.
+
+        Parameters
+        ----------
+        index : int
+            Index of the patch to return.
+
+        Returns
+        -------
+        Tuple[np.ndarray]
+            Patch.
+
+        Raises
+        ------
+        ValueError
+            If dataset mean and std are not set.
+        """
+        patch = self.data[index].squeeze()
+
+        if self.mean is not None and self.std is not None:
+            if isinstance(patch, tuple):
+                patch = normalize(img=patch[0], mean=self.mean, std=self.std)
+                patch = (patch, *patch[1:])
+            else:
+                patch = normalize(img=patch, mean=self.mean, std=self.std)
+
+            if self.patch_transform is not None:
+                # replace None self.patch_transform_params with empty dict
+                if self.patch_transform_params is None:
+                    self.patch_transform_params = {}
+
+                patch = self.patch_transform(patch, **self.patch_transform_params)
+            return patch
+        else:
+            raise ValueError("Dataset mean and std must be set before using it.")

careamics/dataset/patching.py

@@ -0,0 +1,492 @@
+"""
+Tiling submodule.
+
+These functions are used to tile images into patches or tiles.
+"""
+import itertools
+from typing import Generator, List, Optional, Tuple, Union
+
+import numpy as np
+from skimage.util import view_as_windows
+
+from careamics.utils.logging import get_logger
+
+from .extraction_strategy import ExtractionStrategy
+
+logger = get_logger(__name__)
+
+
+def _compute_number_of_patches(
+    arr: np.ndarray, patch_sizes: Union[List[int], Tuple[int, ...]]
+) -> Tuple[int, ...]:
+    """
+    Compute the number of patches that fit in each dimension.
+
+    Array must have one dimension more than the patches (C dimension).
+
+    Parameters
+    ----------
+    arr : np.ndarray
+        Input array.
+    patch_sizes : Tuple[int]
+        Size of the patches.
+
+    Returns
+    -------
+    Tuple[int]
+        Number of patches in each dimension.
+    """
+    n_patches = [
+        np.ceil(arr.shape[i + 1] / patch_sizes[i]).astype(int)
+        for i in range(len(patch_sizes))
+    ]
+    return tuple(n_patches)
+
+
+def _compute_overlap(
+    arr: np.ndarray, patch_sizes: Union[List[int], Tuple[int, ...]]
+) -> Tuple[int, ...]:
+    """
+    Compute the overlap between patches in each dimension.
+
+    Array must be of dimensions C(Z)YX, and patches must be of dimensions YX or ZYX.
+    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 : np.ndarray
+        Input array 3 or 4 dimensions.
+    patch_sizes : Tuple[int]
+        Size of the patches.
+
+    Returns
+    -------
+    Tuple[int]
+        Overlap between patches in each dimension.
+    """
+    n_patches = _compute_number_of_patches(arr, patch_sizes)
+
+    overlap = [
+        np.ceil(
+            np.clip(n_patches[i] * patch_sizes[i] - arr.shape[i + 1], 0, None)
+            / max(1, (n_patches[i] - 1))
+        ).astype(int)
+        for i in range(len(patch_sizes))
+    ]
+    return tuple(overlap)
+
+
+def _compute_crop_and_stitch_coords_1d(
+    axis_size: int, tile_size: int, overlap: int
+) -> Tuple[List[Tuple[int, 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 a certain step
+    for i in range(0, 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((axis_size - tile_size, axis_size))
+            last_tile_end_coord = stitch_coords[-1][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 _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)
+
+
+def _compute_reshaped_view(
+    arr: np.ndarray,
+    window_shape: Tuple[int, ...],
+    step: Tuple[int, ...],
+    output_shape: Tuple[int, ...],
+) -> np.ndarray:
+    """
+    Compute reshaped views of an array, where views correspond 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()
+    patches = view_as_windows(arr, window_shape=window_shape, step=step).reshape(
+        *output_shape
+    )
+    rng.shuffle(patches, axis=0)
+    return patches
+
+
+def _patches_sanity_check(
+    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:
+    - The patch sizes must have one dimension fewer than the array (C dimension).
+    - Chack that patch sizes are smaller than array dimensions.
+
+    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[1:]):
+        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 inconsistent with image shape "
+            f"(got {patch_size} patches for dims {arr.shape[-2:]})."
+        )
+
+
+# formerly :
+# in dataloader.py#L52, 00d536c
+def _extract_patches_sequential(
+    arr: np.ndarray, patch_size: Union[List[int], Tuple[int]]
+) -> Generator[np.ndarray, None, None]:
+    """
+    Generate patches from an array in a sequential manner.
+
+    Array dimensions should be C(Z)YX, where C can be a singleton dimension. 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[np.ndarray, None, None]
+        Generator of patches.
+    """
+    # Patches sanity check
+    is_3d_patch = len(patch_size) == 3
+
+    _patches_sanity_check(arr, patch_size, is_3d_patch)
+
+    # Compute overlap
+    overlaps = _compute_overlap(arr=arr, patch_sizes=patch_size)
+
+    # Create view window and overlaps
+    window_steps = _compute_patch_steps(patch_sizes=patch_size, overlaps=overlaps)
+
+    # Correct for first dimension for computing windowed views
+    window_shape = (1, *patch_size)
+    window_steps = (1, *window_steps)
+
+    if is_3d_patch and patch_size[0] == 1:
+        output_shape = (-1,) + window_shape[1:]
+    else:
+        output_shape = (-1, *window_shape)
+
+    # 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_reshaped_view(
+        arr, window_shape=window_shape, step=window_steps, output_shape=output_shape
+    )
+    logger.info(f"Extracted {patches.shape[0]} patches from input array.")
+
+    # return a generator of patches
+    return (patches[i, ...] for i in range(patches.shape[0]))
+
+
+def _extract_patches_random(
+    arr: np.ndarray, patch_size: Union[List[int], Tuple[int]]
+) -> 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.
+
+    Yields
+    ------
+    Generator[np.ndarray, None, None]
+        Generator of patches.
+    """
+    is_3d_patch = len(patch_size) == 3
+
+    # Patches sanity check
+    _patches_sanity_check(arr, patch_size, is_3d_patch)
+
+    rng = np.random.default_rng()
+    # shuffle the array along the first axis TODO do we need shuffling?
+    rng.shuffle(arr, axis=0)
+
+    for sample_idx in range(arr.shape[0]):
+        sample = arr[sample_idx]
+        # calculate how many number of patches can image area be divided into
+        n_patches = np.ceil(np.prod(sample.shape) / np.prod(patch_size)).astype(int)
+        for _ in range(n_patches):
+            crop_coords = [
+                rng.integers(0, arr.shape[i + 1] - patch_size[i])
+                for i in range(len(patch_size))
+            ]
+            patch = (
+                sample[
+                    (
+                        ...,
+                        *[
+                            slice(c, c + patch_size[i])
+                            for i, c in enumerate(crop_coords)
+                        ],
+                    )
+                ]
+                .copy()
+                .astype(np.float32)
+            )
+            yield patch
+
+
+def _extract_tiles(
+    arr: np.ndarray,
+    tile_size: Union[List[int], Tuple[int]],
+    overlaps: Union[List[int], Tuple[int]],
+) -> Generator:
+    """
+    Generate tiles from the input array with specified overlap.
+
+    The tiles cover the whole array.
+
+    Parameters
+    ----------
+    arr : np.ndarray
+        Array of shape (S, (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
+        Tile generator that yields the tile with corresponding coordinates to stitch
+        back the tiles together.
+    """
+    # Iterate over num samples (S)
+    for sample_idx in range(arr.shape[0]):
+        sample = arr[sample_idx]
+
+        # Create an array of coordinates for cropping and stitching all axes.
+        # Shape: (axes, type_of_coord, tile_num, start/end coord)
+        crop_and_stitch_coords_list = [
+            _compute_crop_and_stitch_coords_1d(
+                sample.shape[i], 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.
+        # 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)]),
+        # where the first list is crop coordinates for 1st axis.
+        all_crop_coords, all_stitch_coords, all_overlap_crop_coords = zip(
+            *crop_and_stitch_coords_list
+        )
+
+        # 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),
+            )
+        ):
+            tile = sample[(..., *[slice(c[0], c[1]) for c in list(crop_coords)])]
+
+            # Check if we are at the end of the sample.
+            # To check that we compute the length of the array that contains all the
+            # tiles
+            if tile_idx == np.prod([len(axis) for axis in all_crop_coords]) - 1:
+                last_tile = True
+            else:
+                last_tile = False
+            yield (
+                np.expand_dims(tile.astype(np.float32), 0),
+                last_tile,
+                arr.shape[1:],
+                overlap_crop_coords,
+                stitch_coords,
+            )
+
+
+def generate_patches(
+    sample: np.ndarray,
+    patch_extraction_method: ExtractionStrategy,
+    patch_size: Optional[Union[List[int], Tuple[int]]] = None,
+    patch_overlap: Optional[Union[List[int], Tuple[int]]] = None,
+) -> Generator[np.ndarray, None, None]:
+    """
+    Generate patches from a sample.
+
+    Parameters
+    ----------
+    sample : np.ndarray
+        Input array.
+    patch_extraction_method : ExtractionStrategies
+        Patch extraction method, as defined in extraction_strategy.ExtractionStrategy.
+    patch_size : Optional[Union[List[int], Tuple[int]]]
+        Size of the patches along each dimension of the array, except the first.
+    patch_overlap : Optional[Union[List[int], Tuple[int]]]
+        Overlap between patches.
+
+    Returns
+    -------
+    Generator[np.ndarray, None, None]
+        Generator yielding patches/tiles.
+
+    Raises
+    ------
+    ValueError
+        If overlap is not specified when using tiling.
+    ValueError
+        If patches is None.
+    """
+    patches = None
+
+    if patch_size is not None:
+        patches = None
+
+        if patch_extraction_method == ExtractionStrategy.TILED:
+            if patch_overlap is None:
+                raise ValueError(
+                    "Overlaps must be specified when using tiling (got None)."
+                )
+            patches = _extract_tiles(
+                arr=sample, tile_size=patch_size, overlaps=patch_overlap
+            )
+
+        elif patch_extraction_method == ExtractionStrategy.SEQUENTIAL:
+            patches = _extract_patches_sequential(sample, patch_size=patch_size)
+
+        else:
+            # random patching
+            patches = _extract_patches_random(sample, patch_size=patch_size)
+
+        return patches
+    else:
+        # no patching, return a generator for the sample
+        return (sample for _ in range(1))