careamics 0.0.19__py3-none-any.whl

careamics/dataset_ng/patching_strategies/random_patching.py

@@ -0,0 +1,393 @@
+"""A module for random patching strategies."""
+
+from collections.abc import Sequence
+
+import numpy as np
+
+from .patching_strategy_protocol import PatchSpecs
+
+
+class RandomPatchingStrategy:
+    """
+    A patching strategy for sampling random patches, it implements the
+    `PatchingStrategy` `Protocol`.
+
+    The output of `get_patch_spec` will be random, i.e. if the same index is given
+    twice the two outputs can be different.
+
+    However the strategy still ensures that there will be a known number of patches for
+    each sample in each image stack. This is achieved through defining a set of bins
+    that map to each sample in each image stack. Whichever bin an `index` passed to
+    `get_patch_spec` falls into, determines the `"data_idx"` and `"sample_idx"` in
+    the returned `PatchSpecs`, but the `"coords"` will be random.
+
+    The number of patches in each sample is based on the number of patches that would
+    fit if they were sampled sequentially, non-overlapping, and covering the entire
+    array.
+    """
+
+    def __init__(
+        self,
+        data_shapes: Sequence[Sequence[int]],
+        patch_size: Sequence[int],
+        seed: int | None = None,
+    ):
+        """
+        A patching strategy for sampling random patches.
+
+        Parameters
+        ----------
+        data_shapes : sequence of (sequence of int)
+            The shapes of the underlying data. Each element is the dimension of the
+            axes SC(Z)YX.
+        patch_size : sequence of int
+            The size of the patch. The sequence will have length 2 or 3, for 2D and 3D
+            data respectively.
+        seed : int, optional
+            An optional seed to ensure the reproducibility of the random patches.
+        """
+        self.rng = np.random.default_rng(seed=seed)
+        self.patch_size = patch_size
+        self.data_shapes = data_shapes
+
+        # these bins will determine which image stack and sample a patch comes from
+        # the image_stack_cumulative_patches map a patch index to each image stack
+        # the sample_cumulative_patches map a patch index to each sample
+        # the image_stack_cumulative_samples map a sample index to each image stack
+        (
+            self.image_stack_cumulative_patches,
+            self.sample_cumulative_patches,
+            self.image_stack_cumulative_samples,
+        ) = self._calc_bins(self.data_shapes, self.patch_size)
+
+    @property
+    def n_patches(self) -> int:
+        """
+        The number of patches that this patching strategy will return.
+
+        It also determines the maximum index that can be given to `get_patch_spec`.
+        """
+        # last bin boundary will be total patches
+        return self.image_stack_cumulative_patches[-1]
+
+    def get_patch_spec(self, index: int) -> PatchSpecs:
+        """Return the patch specs for a given index.
+
+        Parameters
+        ----------
+        index : int
+            A patch index.
+
+        Returns
+        -------
+        PatchSpecs
+            A dictionary that specifies a single patch in a series of `ImageStacks`.
+        """
+        # TODO: break into smaller testable functions?
+        if index >= self.n_patches:
+            raise IndexError(
+                f"Index {index} out of bounds for RandomPatchingStrategy with number "
+                f"of patches {self.n_patches}"
+            )
+        # digitize returns the bin that `index` belongs to
+        data_index = np.digitize(index, bins=self.image_stack_cumulative_patches).item()
+        # maps to a particular sample within the whole series of image stacks
+        #   (not just a single image stack)
+        total_samples_index = np.digitize(
+            index, bins=self.sample_cumulative_patches
+        ).item()
+
+        data_shape = self.data_shapes[data_index]
+        spatial_shape = data_shape[2:]
+
+        # calculate sample index relative to image stack:
+        #   subtract the total number of samples in the previous image stacks
+        if data_index == 0:
+            n_previous_samples = 0
+        else:
+            n_previous_samples = self.image_stack_cumulative_samples[data_index - 1]
+        sample_index = total_samples_index - n_previous_samples
+        coords = _generate_random_coords(spatial_shape, self.patch_size, self.rng)
+        return {
+            "data_idx": data_index,
+            "sample_idx": sample_index,
+            "coords": coords,
+            "patch_size": self.patch_size,
+        }
+
+    # Note: this is used by the FileIterSampler
+    def get_patch_indices(self, data_idx: int) -> Sequence[int]:
+        """
+        Get the patch indices will return patches for a specific `image_stack`.
+
+        The `image_stack` corresponds to the given `data_idx`.
+
+        Parameters
+        ----------
+        data_idx : int
+            An index that corresponds to a given `image_stack`.
+
+        Returns
+        -------
+        sequence of int
+            A sequence of patch indices, that when used to index the `CAREamicsDataset
+            will return a patch that comes from the `image_stack` corresponding to the
+            given `data_idx`.
+        """
+        # return all the values in the corresponding bin
+        if data_idx == 0:
+            start = 0
+        else:
+            start = self.image_stack_cumulative_patches[data_idx - 1]
+
+        return np.arange(start, self.image_stack_cumulative_patches[data_idx]).tolist()
+
+    @staticmethod
+    def _calc_bins(
+        data_shapes: Sequence[Sequence[int]], patch_size: Sequence[int]
+    ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
+        """Calculate bins used to map an index to an image_stack and a sample.
+
+        The number of patches in each sample is based on the number of patches that
+        would fit if they were sampled sequentially.
+
+        Parameters
+        ----------
+        data_shapes : sequence of (sequence of int)
+            The shapes of the underlying data. Each element is the dimension of the
+            axes SC(Z)YX.
+        patch_size : sequence of int
+            The size of the patch. The sequence will have length 2 or 3, for 2D and 3D
+            data respectively.
+
+        Returns
+        -------
+        image_stack_cumulative_patches: tuple of int
+            The bins that map a patch index to an image stack. E.g. if a patch index
+            falls below the first bin boundary it belongs to the first image stack, if
+            a patch index falls between the first bin boundary and the second bin
+            boundary it belongs to the second image stack, and so on.
+        sample_cumulative_patches: tuple of int
+            The bins that map a patch index to a sample. E.g. if a patch index
+            falls below the first bin boundary it belongs to the first sample, if
+            a patch index falls between the first bin boundary and the second bin
+            boundary it belongs to the second sample, and so on.
+        image_stack_cumulative_samples: tuple of int
+            The bins that map a sample index to an image stack. E.g. if a sample index
+            falls below the first bin boundary it belongs to the first image stack, if
+            a patch index falls between the first bin boundary and the second bin
+            boundary it belongs to the second image stack, and so on.
+        """
+        patches_per_image_stack: list[int] = []
+        patches_per_sample: list[int] = []
+        samples_per_image_stack: list[int] = []
+        for data_shape in data_shapes:
+            spatial_shape = data_shape[2:]
+            n_single_sample_patches = _calc_n_patches(spatial_shape, patch_size)
+            # multiply by number of samples in image_stack
+            patches_per_image_stack.append(n_single_sample_patches * data_shape[0])
+            # list of length `sample` filled with `n_single_sample_patches`
+            patches_per_sample.extend([n_single_sample_patches] * data_shape[0])
+            # number of samples in each image stack
+            samples_per_image_stack.append(data_shape[0])
+
+        # cumulative sum creates the bins
+        image_stack_cumulative_patches = np.cumsum(patches_per_image_stack)
+        sample_cumulative_patches = np.cumsum(patches_per_sample)
+        image_stack_cumulative_samples = np.cumsum(samples_per_image_stack)
+        return (
+            tuple(image_stack_cumulative_patches),
+            tuple(sample_cumulative_patches),
+            tuple(image_stack_cumulative_samples),
+        )
+
+
+class FixedRandomPatchingStrategy:
+    """
+    A patching strategy for sampling random patches it implements the `PatchingStrategy`
+    `Protocol`.
+
+    The output of `get_patch_spec` will be deterministic, i.e. if the same index is
+    given twice the two outputs will be the same.
+
+    The number of patches in each sample is based on the number of patches that would
+    fit if they were sampled sequentially, non-overlapping, and covering the entire
+    array.
+    """
+
+    def __init__(
+        self,
+        data_shapes: Sequence[Sequence[int]],
+        patch_size: Sequence[int],
+        seed: int | None = None,
+    ):
+        """A patching strategy for sampling random patches.
+
+        Parameters
+        ----------
+        data_shapes : sequence of (sequence of int)
+            The shapes of the underlying data. Each element is the dimension of the
+            axes SC(Z)YX.
+        patch_size : sequence of int
+            The size of the patch. The sequence will have length 2 or 3, for 2D and 3D
+            data respectively.
+        seed : int, optional
+            An optional seed to ensure the reproducibility of the random patches.
+        """
+        self.rng = np.random.default_rng(seed=seed)
+        self.patch_size = patch_size
+        self.data_shapes = data_shapes
+
+        # simply generate all the patches at initialisation, so they will be fixed
+        self.fixed_patch_specs: list[PatchSpecs] = []
+        for data_idx, data_shape in enumerate(self.data_shapes):
+            spatial_shape = data_shape[2:]
+            n_patches = _calc_n_patches(spatial_shape, self.patch_size)
+            for sample_idx in range(data_shape[0]):
+                for _ in range(n_patches):
+                    random_coords = _generate_random_coords(
+                        spatial_shape, self.patch_size, self.rng
+                    )
+                    patch_specs: PatchSpecs = {
+                        "data_idx": data_idx,
+                        "sample_idx": sample_idx,
+                        "coords": random_coords,
+                        "patch_size": self.patch_size,
+                    }
+                    self.fixed_patch_specs.append(patch_specs)
+
+    @property
+    def n_patches(self):
+        """
+        The number of patches that this patching strategy will return.
+
+        It also determines the maximum index that can be given to `get_patch_spec`.
+        """
+        return len(self.fixed_patch_specs)
+
+    def get_patch_spec(self, index: int) -> PatchSpecs:
+        """Return the patch specs for a given index.
+
+        Parameters
+        ----------
+        index : int
+            A patch index.
+
+        Returns
+        -------
+        PatchSpecs
+            A dictionary that specifies a single patch in a series of `ImageStacks`.
+        """
+        if index >= self.n_patches:
+            raise IndexError(
+                f"Index {index} out of bounds for FixedRandomPatchingStrategy with "
+                f"number of patches, {self.n_patches}"
+            )
+        # simply index the pre-generated patches to get the correct patch
+        return self.fixed_patch_specs[index]
+
+    # Note: this is used by the FileIterSampler
+    def get_patch_indices(self, data_idx: int) -> Sequence[int]:
+        """
+        Get the patch indices will return patches for a specific `image_stack`.
+
+        The `image_stack` corresponds to the given `data_idx`.
+
+        Parameters
+        ----------
+        data_idx : int
+            An index that corresponds to a given `image_stack`.
+
+        Returns
+        -------
+        sequence of int
+            A sequence of patch indices, that when used to index the `CAREamicsDataset
+            will return a patch that comes from the `image_stack` corresponding to the
+            given `data_idx`.
+        """
+        return [
+            i
+            for i, patch_spec in enumerate(self.fixed_patch_specs)
+            if patch_spec["data_idx"] == data_idx
+        ]
+
+
+def _generate_random_coords(
+    spatial_shape: Sequence[int], patch_size: Sequence[int], rng: np.random.Generator
+) -> tuple[int, ...]:
+    """Generate random patch coordinates for a given `spatial_shape` and `patch_size`.
+
+    The coords are the top-left (and first z-slice for 3D data) of a patch. The
+    sequence will have length 2 or 3, for 2D and 3D data respectively.
+
+    Parameters
+    ----------
+    spatial_shape : sequence of int
+        The dimension of the axes (Z)YX, a sequence of length 2 or 3, for 2D and 3D
+        data respectively.
+    patch_size : sequence of int
+        The size of the patch. The sequence will have length 2 or 3, for 2D and 3D
+        data respectively.
+    rng : numpy.random.Generator
+        A numpy generator to ensure the reproducibility of the random patches.
+
+    Returns
+    -------
+    coords: tuple of int
+        The top-left (and first z-slice for 3D data) coords of a patch. The tuple will
+        have length 2 or 3, for 2D and 3D data respectively.
+
+    Raises
+    ------
+    ValueError
+        Raises if the number of spatial dimensions do not match the number of patch
+        dimensions.
+    """
+    if len(patch_size) != len(spatial_shape):
+        raise ValueError(
+            f"Number of patch dimension {len(patch_size)}, do not match the number of "
+            f"spatial dimensions {len(spatial_shape)}, for `patch_size={patch_size}` "
+            f"and `spatial_shape={spatial_shape}`."
+        )
+    return tuple(
+        rng.integers(
+            np.zeros(len(patch_size), dtype=int),
+            np.clip(np.array(spatial_shape) - np.array(patch_size), 0, None),
+            endpoint=True,
+            dtype=int,
+        ).tolist()
+    )
+
+
+def _calc_n_patches(spatial_shape: Sequence[int], patch_size: Sequence[int]) -> int:
+    """
+    Calculates the number of patches for a given `spatial_shape` and `patch_size`.
+
+    This is based on the number of patches that would fit if they were sampled
+    sequentially.
+
+    Parameters
+    ----------
+    spatial_shape : sequence of int
+        The dimension of the axes (Z)YX, a sequence of length 2 or 3, for 2D and 3D
+        data respectively.
+    patch_size : sequence of int
+        The size of the patch. The sequence will have length 2 or 3, for 2D and 3D
+        data respectively.
+
+    Returns
+    -------
+    int
+        The number of patches.
+    """
+    if len(patch_size) != len(spatial_shape):
+        raise ValueError(
+            f"Number of patch dimension {len(patch_size)}, do not match the number of "
+            f"spatial dimensions {len(spatial_shape)}, for `patch_size={patch_size}` "
+            f"and `spatial_shape={spatial_shape}`."
+        )
+    patches_per_dim = [
+        np.ceil(s / p) for s, p in zip(spatial_shape, patch_size, strict=False)
+    ]
+    total_patches = int(np.prod(patches_per_dim))
+    return total_patches

careamics/dataset_ng/patching_strategies/sequential_patching.py

@@ -0,0 +1,99 @@
+import itertools
+from collections.abc import Sequence
+
+import numpy as np
+from typing_extensions import ParamSpec
+
+from .patching_strategy_protocol import PatchSpecs
+
+P = ParamSpec("P")
+
+
+# TODO: this is an unfinished prototype based on current tiling implementation
+#  not guaranteed to work!
+class SequentialPatchingStrategy:
+    # TODO: docs
+    def __init__(
+        self,
+        data_shapes: Sequence[Sequence[int]],
+        patch_size: Sequence[int],
+        overlaps: Sequence[int] | None = None,
+    ):
+        self.data_shapes = data_shapes
+        self.patch_size = patch_size
+        if overlaps is None:
+            overlaps = [0] * len(patch_size)
+        self.overlaps = np.asarray(overlaps)
+
+        self.patch_specs: list[PatchSpecs] = self._initialize_patch_specs()
+
+    @property
+    def n_patches(self) -> int:
+        return len(self.patch_specs)
+
+    def get_patch_spec(self, index: int) -> PatchSpecs:
+        return self.patch_specs[index]
+
+    # Note: this is used by the FileIterSampler
+    def get_patch_indices(self, data_idx: int) -> Sequence[int]:
+        """
+        Get the patch indices will return patches for a specific `image_stack`.
+
+        The `image_stack` corresponds to the given `data_idx`.
+
+        Parameters
+        ----------
+        data_idx : int
+            An index that corresponds to a given `image_stack`.
+
+        Returns
+        -------
+        sequence of int
+            A sequence of patch indices, that when used to index the `CAREamicsDataset
+            will return a patch that comes from the `image_stack` corresponding to the
+            given `data_idx`.
+        """
+        return [
+            i
+            for i, patch_spec in enumerate(self.patch_specs)
+            if patch_spec["data_idx"] == data_idx
+        ]
+
+    def _compute_coords_1d(
+        self, patch_size: int, spatial_shape: int, overlap: int
+    ) -> list[tuple[int, int]]:
+        step = patch_size - overlap
+        crop_coords = []
+
+        current_pos = 0
+        while current_pos <= spatial_shape - patch_size:
+            crop_coords.append((current_pos, current_pos + patch_size))
+            current_pos += step
+
+        if crop_coords[-1][1] < spatial_shape:
+            crop_coords.append((spatial_shape - patch_size, spatial_shape))
+
+        return crop_coords
+
+    def _initialize_patch_specs(self) -> list[PatchSpecs]:
+        patch_specs: list[PatchSpecs] = []
+        for data_idx, data_shape in enumerate(self.data_shapes):
+
+            data_spatial_shape = data_shape[-len(self.patch_size) :]
+            coords_list = [
+                self._compute_coords_1d(
+                    self.patch_size[i], data_spatial_shape[i], self.overlaps[i]
+                )
+                for i in range(len(self.patch_size))
+            ]
+            for sample_idx in range(data_shape[0]):
+                for crop_coord in itertools.product(*coords_list):
+                    patch_specs.append(
+                        PatchSpecs(
+                            data_idx=data_idx,
+                            sample_idx=sample_idx,
+                            coords=tuple(coord[0] for coord in crop_coord),
+                            patch_size=self.patch_size,
+                        )
+                    )
+        return patch_specs

careamics/dataset_ng/patching_strategies/tiling_strategy.py

@@ -0,0 +1,207 @@
+"""Module for the `TilingStrategy` class."""
+
+import itertools
+from collections.abc import Sequence
+from math import prod
+
+from .patching_strategy_protocol import TileSpecs
+
+
+class TilingStrategy:
+    """
+    The tiling strategy should be used for prediction. The `get_patch_specs`
+    method returns `TileSpec` dictionaries that contains information on how to
+    stitch the tiles back together to create the full image.
+    """
+
+    def __init__(
+        self,
+        data_shapes: Sequence[Sequence[int]],
+        patch_size: Sequence[int],
+        overlaps: Sequence[int],
+    ):
+        """
+        The tiling strategy should be used for prediction. The `get_patch_specs`
+        method returns `TileSpec` dictionaries that contains information on how to
+        stitch the tiles back together to create the full image.
+
+        Parameters
+        ----------
+        data_shapes : sequence of (sequence of int)
+            The shapes of the underlying data. Each element is the dimension of the
+            axes SC(Z)YX.
+        patch_size : sequence of int
+            The size of the tile. The sequence will have length 2 or 3, for 2D and 3D
+            data respectively.
+        overlaps : sequence of int
+            How much a tile will overlap with adjacent tiles in each spatial dimension.
+        """
+        self.data_shapes = data_shapes
+        self.patch_size = patch_size
+        self.overlaps = overlaps
+        # patch_size and overlap should have same length validated in pydantic configs
+        self.tile_specs: list[TileSpecs] = self._generate_specs()
+
+    @property
+    def n_patches(self) -> int:
+        """
+        The number of patches that this patching strategy will return.
+
+        It also determines the maximum index that can be given to `get_patch_spec`.
+        """
+        return len(self.tile_specs)
+
+    def get_patch_spec(self, index: int) -> TileSpecs:
+        """Return the tile specs for a given index.
+
+        Parameters
+        ----------
+        index : int
+            A patch index.
+
+        Returns
+        -------
+        TileSpecs
+            A dictionary that specifies a single patch in a series of `ImageStacks`.
+        """
+        return self.tile_specs[index]
+
+    # Note: this is used by the FileIterSampler
+    def get_patch_indices(self, data_idx: int) -> Sequence[int]:
+        """
+        Get the patch indices will return patches for a specific `image_stack`.
+
+        The `image_stack` corresponds to the given `data_idx`.
+
+        Parameters
+        ----------
+        data_idx : int
+            An index that corresponds to a given `image_stack`.
+
+        Returns
+        -------
+        sequence of int
+            A sequence of patch indices, that when used to index the `CAREamicsDataset
+            will return a patch that comes from the `image_stack` corresponding to the
+            given `data_idx`.
+        """
+        return [
+            i
+            for i, patch_spec in enumerate(self.tile_specs)
+            if patch_spec["data_idx"] == data_idx
+        ]
+
+    def _generate_specs(self) -> list[TileSpecs]:
+        tile_specs: list[TileSpecs] = []
+        for data_idx, data_shape in enumerate(self.data_shapes):
+            spatial_shape = data_shape[2:]
+
+            # spec info for each axis
+            axis_specs: list[tuple[list[int], list[int], list[int], list[int]]] = [
+                self._compute_1d_coords(
+                    axis_size, self.patch_size[axis_idx], self.overlaps[axis_idx]
+                )
+                for axis_idx, axis_size in enumerate(spatial_shape)
+            ]
+
+            # combine by using zip
+            all_coords, all_stitch_coords, all_crop_coords, all_crop_size = zip(
+                *axis_specs, strict=False
+            )
+
+            # number of tiles for this data_idx
+            n_tiles = prod(len(dim) for dim in all_coords) * data_shape[0]
+
+            # patches will be the same for each sample in a stack
+            for sample_idx in range(data_shape[0]):
+                # iterate through all combinations using itertools.product
+                for coords, stitch_coords, crop_coords, crop_size in zip(
+                    itertools.product(*all_coords),
+                    itertools.product(*all_stitch_coords),
+                    itertools.product(*all_crop_coords),
+                    itertools.product(*all_crop_size),
+                    strict=False,
+                ):
+                    tile_specs.append(
+                        {
+                            # PatchSpecs
+                            "data_idx": data_idx,
+                            "sample_idx": sample_idx,
+                            "coords": coords,
+                            "patch_size": self.patch_size,
+                            # TileSpecs additional fields
+                            "crop_coords": crop_coords,
+                            "crop_size": crop_size,
+                            "stitch_coords": stitch_coords,
+                            "total_tiles": n_tiles,
+                        }
+                    )
+
+        return tile_specs
+
+    @staticmethod
+    def _compute_1d_coords(
+        axis_size: int, patch_size: int, overlap: int
+    ) -> tuple[list[int], list[int], list[int], list[int]]:
+        """
+        Computes the TileSpec information for a single axis.
+
+        Parameters
+        ----------
+        axis_size : int
+            The size of the axis.
+        patch_size : int
+            The tile size.
+        overlap : int
+            The tile overlap.
+
+        Returns
+        -------
+        coords: list of int
+            The top-left (and first z-slice for 3D data) of a tile, in coords relative
+            to the image.
+        stitch_coords: list of int
+            Where the tile will be stitched back into an image, taking into account
+            that the tile will be cropped, in coords relative to the image.
+        crop_coords: list of int
+            The top-left side of where the tile will be cropped, in coordinates relative
+            to the tile.
+        crop_size: list of int
+            The size of the cropped tile.
+        """
+        coords: list[int] = []
+        stitch_coords: list[int] = []
+        crop_coords: list[int] = []
+        crop_size: list[int] = []
+
+        step = patch_size - overlap
+        for i in range(0, max(1, axis_size - overlap), step):
+            if i == 0:
+                coords.append(i)
+                crop_coords.append(0)
+                stitch_coords.append(0)
+                if axis_size <= patch_size:
+                    crop_size.append(axis_size)
+                else:
+                    crop_size.append(patch_size - overlap // 2)
+            elif (0 < i) and (i + patch_size < axis_size):
+                coords.append(i)
+                crop_coords.append(overlap // 2)
+                stitch_coords.append(coords[-1] + crop_coords[-1])
+                crop_size.append(patch_size - overlap)
+            else:
+                previous_crop_size = crop_size[-1] if crop_size else 1
+                previous_stitch_coord = stitch_coords[-1] if stitch_coords else 0
+                previous_tile_end = previous_stitch_coord + previous_crop_size
+
+                coords.append(max(0, axis_size - patch_size))
+                stitch_coords.append(previous_tile_end)
+                crop_coords.append(stitch_coords[-1] - coords[-1])
+                crop_size.append(axis_size - stitch_coords[-1])
+
+        return (
+            coords,
+            stitch_coords,
+            crop_coords,
+            crop_size,
+        )