careamics 0.0.19__py3-none-any.whl

careamics/dataset_ng/microsplit_input_synth.py

@@ -0,0 +1,377 @@
+"""MicroSplit patch synthesis."""
+
+# --- PROOF OF PRINCIPLE ---
+
+
+from collections.abc import Callable, Sequence
+from typing import Any, Literal, NamedTuple
+
+import numpy as np
+from numpy.typing import NDArray
+
+from .dataset import ImageRegionData
+from .image_stack import ImageStack
+from .patch_extractor import PatchExtractor
+from .patch_filter import PatchFilterProtocol
+from .patching_strategies import PatchingStrategy, PatchSpecs
+
+
+# TODO: better name
+# mirrors format of ImageRegionData
+class UncorrelatedRegionData(NamedTuple):
+    data: NDArray
+    source: Sequence[str | Literal["array"]]
+    data_shape: Sequence[Sequence[int]]
+    dtype: Sequence[str]  # dtype should be str for collate
+    axes: Sequence[str]
+    region_spec: Sequence[PatchSpecs]
+
+
+# --- for finding empty / signal channel patches in loop
+def is_empty(filter: PatchFilterProtocol) -> Callable[[NDArray[Any]], bool]:
+    def is_empty_check(patch: NDArray[Any]) -> bool:
+        return filter.filter_out(patch)
+
+    return is_empty_check
+
+
+def is_not_empty(filter: PatchFilterProtocol) -> Callable[[NDArray[Any]], bool]:
+    def is_not_empty_check(patch: NDArray[Any]) -> bool:
+        return not filter.filter_out(patch)
+
+    return is_not_empty_check
+
+
+# ---
+
+
+def create_default_input_target(
+    idx: int,
+    patch_extractor: PatchExtractor[ImageStack],
+    patching_strategy: PatchingStrategy,
+    alphas: list[float],
+    axes: str,  # annoyingly have to supply this to image region
+) -> tuple[ImageRegionData, ImageRegionData]:
+    """
+    Create a default MicroSplit patch with synthetically summed input.
+
+    Parameters
+    ----------
+    idx: int
+        The dataset index.
+    patch_extractor: PatchExtractor
+        Used to extract patches from the data.
+    patching_strategy: PatchingStrategy
+        Patch locations will be sampled using the patching strategy.
+    alphas: list[float]
+        Weights for each channel for creating the synthetic input with summation.
+    axes: str
+        The axes of the data. This is only used to populate metadata.
+
+    Returns
+    -------
+    input_region: ImageRegionData
+        The input patch and its metadata, the data has the dimension L(Z)YX.
+    target_region: ImageRegionData
+        The target patch and its metadata, the data has the dimensions C(Z)YX.
+    """
+    patch_spec = patching_strategy.get_patch_spec(idx)
+    patches = extract_microsplit_patch(patch_extractor, patch_spec)
+
+    ndims = len(patches.shape) - 1
+    alpha_broadcast = np.array(alphas)[:, *(np.newaxis for _ in range(ndims))]
+    # weight channels by alphas then sum on the channel axis
+    # input dims will be L(Z)YX
+    input_patch = (alpha_broadcast * patches).sum(axis=0)
+    target_patch = patches[:, 0, ...]  # first L patch
+
+    data_idx = patch_spec["data_idx"]
+    input_region = ImageRegionData(
+        input_patch,
+        source=str(patch_extractor.image_stacks[data_idx].source),
+        data_shape=patch_extractor.image_stacks[data_idx].data_shape,
+        dtype=str(patch_extractor.image_stacks[data_idx].data_dtype),
+        axes=axes,
+        region_spec=patch_spec,
+        additional_metadata={},
+    )
+    target_region = ImageRegionData(
+        target_patch,
+        source=str(patch_extractor.image_stacks[data_idx].source),
+        data_shape=patch_extractor.image_stacks[data_idx].data_shape,
+        dtype=str(patch_extractor.image_stacks[data_idx].data_dtype),
+        axes=axes,
+        region_spec=patch_spec,
+        additional_metadata={},
+    )
+    return input_region, target_region
+
+
+def create_uncorrelated_input_target(
+    patches: NDArray[Any],
+    patch_specs: list[PatchSpecs],
+    alphas: list[float],
+    patch_extractor: PatchExtractor[ImageStack],  # for metadata
+    axes: str,  # mirroring imageregion
+) -> tuple[UncorrelatedRegionData, UncorrelatedRegionData]:
+    """
+    Create MicroSplit target and synthetically summed input with metadata.
+
+    Parameters
+    ----------
+    patches: NDArray
+        Patches with dimensions LC(Z)YX, where L contains the lateral context at
+        multiple scales.
+    patch_specs: list[PatchSpecs]
+        The patch specs for each channel.
+    alphas: list[float]
+        Weights for each channel for creating the synthetic input with summation.
+    patch_extractor: PatchExtractor
+        The patch extractor the patches were extracted from. Used for additional
+        metadata.
+
+    Returns
+    -------
+    input_region: UncorrelatedRegionData
+        The input patch and its metadata, the data has the dimension L(Z)YX.
+    target_region: UncorrelatedRegionData
+        The target patch and its metadata, the data has the dimensions C(Z)YX.
+    """
+    ndims = len(patches.shape) - 1
+    alpha_broadcast = np.array(alphas)[:, *(np.newaxis for _ in range(ndims))]
+    # weight channels by alphas then sum on the channel axis
+    # input dims will be L(Z)YX
+    input_patch = (alpha_broadcast * patches).sum(axis=0)
+    target_patch = patches[:, 0, ...]  # first L patch
+
+    input_stacks = [
+        patch_extractor.image_stacks[patch_spec["data_idx"]]
+        for patch_spec in patch_specs
+    ]
+    source = [str(stack.source) for stack in input_stacks]
+    data_shape = [stack.data_shape for stack in input_stacks]
+    dtype = [str(stack.data_dtype) for stack in input_stacks]
+
+    input_region = UncorrelatedRegionData(
+        data=input_patch,
+        source=source,
+        data_shape=data_shape,
+        dtype=dtype,
+        region_spec=patch_specs,
+        axes=axes,
+    )
+    target_region = UncorrelatedRegionData(
+        data=target_patch,
+        source=source,
+        data_shape=data_shape,
+        dtype=dtype,
+        region_spec=patch_specs,
+        axes=axes,
+    )
+    return input_region, target_region
+
+
+def get_random_channel_patches(
+    idx: int,  # TODO: is this needed it makes it work the same as original dataset
+    patch_extractor: PatchExtractor[ImageStack],
+    patching_strategy: PatchingStrategy,
+    rng: np.random.Generator | None,
+) -> tuple[NDArray[Any], list[PatchSpecs]]:
+    """
+    Select patches form random patch locations for each channel.
+
+    Parameters
+    ----------
+    idx: int
+        The dataset index.
+    patch_extractor: PatchExtractor
+        Used to extract patches from the data.
+    patching_strategy: PatchingStrategy
+        Patch locations will be sampled using the patching strategy.
+    rng: numpy.random.Generator | None
+        Useful for seeding the process. If `None` the default random number generator
+        will be used.
+
+    Returns
+    -------
+    NDArray[Any]
+        The resulting patches with dimensions LC(Z)YX, where L contains the lateral
+        context at multiple scales.
+    list[PatchSpecs]
+        A list of patch specification, one for each channel.
+    """
+    if rng is None:
+        rng = np.random.default_rng()
+
+    n_channels = patch_extractor.n_channels
+
+    # in the original dataset, new random indices are chosen for each channel
+    # the other channels can come from anywhere in the entire dataset
+    indices = (idx, *rng.integers(patching_strategy.n_patches, size=(n_channels - 1)))
+
+    # get n different patch specs for n different channels
+    patch_specs = [patching_strategy.get_patch_spec(i) for i in indices]
+    patches = extract_microsplit_patch(patch_extractor, patch_specs)
+
+    return patches, patch_specs
+
+
+# TODO: better name
+def get_empty_channel_patches(
+    idx: int,
+    patch_extractor: PatchExtractor,
+    patching_strategy: PatchingStrategy,
+    signal_channels: dict[int, PatchFilterProtocol],
+    empty_channels: dict[int, PatchFilterProtocol],
+    patience: int,
+    rng: np.random.Generator | None,
+) -> tuple[NDArray[Any], list[PatchSpecs]]:
+    """
+    Select patches, specifying which channels should have signal and which should not.
+
+    Parameters
+    ----------
+    idx: int
+        The dataset index.
+    patch_extractor: PatchExtractor
+        Used to extract patches from the data.
+    patching_strategy: PatchingStrategy
+        Patch locations will be sampled using the patching strategy.
+    signal_channels: dict[int, PatchFilterProtocol]
+        A dictionary to specify the channels that should have signal and how they should
+        be filtered. The keys are the channel index and the values are the patch filters
+        used to determine if the channel patch is empty or not.
+    empty_channels: dict[int, PatchFilterProtocol]
+        A dictionary to specify the channels that should not have signal. Similar to
+        the `signal_channels`.
+    patience: int
+        New patches are selected at random until a patch with signal or without is
+        found, the `patience` determines how many times to look before giving up.
+    rng: numpy.random.Generator | None
+        Useful for seeding the process. If `None` the default random number generator
+        will be used.
+
+    Returns
+    -------
+    NDArray[Any]
+        The resulting patches with dimensions LC(Z)YX, where L contains the lateral
+        context at multiple scales.
+    list[PatchSpecs]
+        A list of patch specification, one for each channel.
+    """
+    if rng is None:
+        rng = np.random.default_rng()
+
+    # if a channel is not selected to be empty or filled it will from idx
+    filled = set(signal_channels.keys())
+    empty = set(empty_channels.keys())
+    if len(intersect := filled.intersection(empty)) != 0:
+        raise ValueError(
+            "Channels cannot be selected as both empty and filled, the following "
+            f"channels were selected as both {intersect}."
+        )
+
+    n_channels = patch_extractor.n_channels
+
+    # start with random initial patches
+    patches, patch_specs = get_random_channel_patches(
+        idx, patch_extractor, patching_strategy, rng
+    )
+
+    # for each channel sample patches until they are empty or not empty
+    for c in range(n_channels):
+
+        # criterion for the while loop
+        criterion: Callable[[NDArray[Any]], bool]
+        filter_: PatchFilterProtocol
+        if c in empty_channels:
+            filter_ = empty_channels[c]
+            criterion = is_not_empty(filter_)
+        elif c in signal_channels:
+            filter_ = signal_channels[c]
+            criterion = is_empty(filter_)
+        else:
+            break
+
+        patch = patches[c]
+        patch_spec = patch_specs[c]
+        patience_ = patience
+        # only check if primary input is empty
+        while criterion(patch[0]) and patience_ > 0:
+            # sample random indices from anywhere in the dataset
+            new_idx = rng.integers(patching_strategy.n_patches)
+            patch_spec = patching_strategy.get_patch_spec(new_idx.item())
+            patch = patch_extractor.extract_channel_patch(
+                data_idx=patch_spec["data_idx"],
+                sample_idx=patch_spec["sample_idx"],
+                channels=[c],
+                coords=patch_spec["coords"],
+                patch_size=patch_spec["patch_size"],
+            )[0]
+            # ^ removing channel dim
+            patience_ -= 1
+        if patience <= 0:
+            # TODO: log properly
+            print(f"Out of patience finding patch for channel {c}")
+
+        patches[c] = patch
+        patch_specs[c] = patch_spec
+
+    return patches, patch_specs
+
+
+def extract_microsplit_patch(
+    patch_extractor: PatchExtractor[ImageStack],
+    patch_specs: PatchSpecs | list[PatchSpecs],
+) -> NDArray[Any]:
+    """
+    Extract a MicroSplit patch with the dimensions LC(Z)YX.
+
+    This patch can be used to synthesis an input patch by summing the C dimension, and
+    it can be used to create a target patch by selecting the primary input from the
+    L dimension, where L is to store lateral context patches.
+
+    Parameters
+    ----------
+    patch_extractor: PatchExtractor
+        Used to extract patches from the data.
+    patch_specs: PatchSpec | list[PatchSpecs]
+        A patch specification or a list of patch specifications — one for each channel.
+        Different patch specs can be used or each channel to create uncorrelated channel
+        patches.
+
+    Returns
+    -------
+    NDArray[Any]
+        The resulting patches with dimensions LC(Z)YX, where L contains the lateral
+        context at multiple scales.
+    """
+    if isinstance(patch_specs, list):
+        patches = np.concat(
+            [
+                patch_extractor.extract_channel_patch(
+                    data_idx=patch_spec["data_idx"],
+                    sample_idx=patch_spec["sample_idx"],
+                    channels=[c],
+                    coords=patch_spec["coords"],
+                    patch_size=patch_spec["patch_size"],
+                )
+                for c, patch_spec in enumerate(patch_specs)
+            ],
+            axis=0,
+        )
+    else:
+        patches = patch_extractor.extract_patch(
+            data_idx=patch_specs["data_idx"],
+            sample_idx=patch_specs["sample_idx"],
+            coords=patch_specs["coords"],
+            patch_size=patch_specs["patch_size"],
+        )
+    # Add L dimension if not present
+    n_spatial_dims = patch_extractor.n_spatial_dims
+    lateral_context_present = len(patches.shape) - n_spatial_dims == 2
+    if not lateral_context_present:
+        # insert a L dim
+        patches = patches[:, np.newaxis]
+
+    return patches

careamics/dataset_ng/patch_extractor/__init__.py

@@ -0,0 +1,7 @@
+__all__ = [
+    "LimitFilesPatchExtractor",
+    "PatchExtractor",
+]
+
+from .limit_file_extractor import LimitFilesPatchExtractor
+from .patch_extractor import PatchExtractor

careamics/dataset_ng/patch_extractor/limit_file_extractor.py

@@ -0,0 +1,50 @@
+from collections.abc import Sequence
+
+from numpy.typing import NDArray
+
+from ..image_stack import FileImageStack
+from .patch_construction import PatchConstructor, default_patch_constr
+from .patch_extractor import PatchExtractor
+
+
+class LimitFilesPatchExtractor(PatchExtractor[FileImageStack]):
+    """
+    A patch extractor that limits the number of files that have their data loaded.
+
+    This is useful for when not all of the data will fit into memory.
+    """
+
+    def __init__(
+        self,
+        image_stacks: Sequence[FileImageStack],
+        patch_constructor: PatchConstructor = default_patch_constr,
+    ):
+        """
+        Parameters
+        ----------
+        image_stacks: Sequence of `FileImageStack`
+        """
+        super().__init__(image_stacks, patch_constructor)
+        self.loaded_stacks: list[int] = []
+
+    def extract_channel_patch(
+        self,
+        data_idx: int,
+        sample_idx: int,
+        channels: Sequence[int] | None,
+        coords: Sequence[int],
+        patch_size: Sequence[int],
+    ) -> NDArray:
+        if data_idx not in self.loaded_stacks:
+            # TODO: make maximum images loaded configurable?
+            if len(self.loaded_stacks) >= 1:
+                # get the idx that was added longest ago
+                idx_to_close = self.loaded_stacks.pop(0)
+                self.image_stacks[idx_to_close].close()
+
+            self.image_stacks[data_idx].load()
+            self.loaded_stacks.append(data_idx)
+
+        return super().extract_channel_patch(
+            data_idx, sample_idx, channels, coords, patch_size
+        )

careamics/dataset_ng/patch_extractor/patch_construction.py

@@ -0,0 +1,151 @@
+from collections.abc import Sequence
+from typing import Any, Literal, Protocol
+
+import numpy as np
+from numpy.typing import NDArray
+from skimage.transform import resize
+
+from ..image_stack import ImageStack
+
+
+class PatchConstructor(Protocol):
+    """
+    A callable that modifies how patches are constructed in the PatchExtractor.
+
+    This protocol defines the signature of a callable that is passed as an argument to
+    the `PatchExtractor`. It can be used to modify how patches are constructed, for
+    example creating patches with multiple lateral context levels for MicroSplit.
+    """
+
+    def __call__(
+        self,
+        image_stack: ImageStack,
+        sample_idx: int,
+        channels: Sequence[int] | None,  # `channels = None` to select all channels
+        coords: Sequence[int],
+        patch_size: Sequence[int],
+    ) -> NDArray[Any]:
+        """
+        Parameters
+        ----------
+        image_stack: ImageStack
+            The image stack to construct a patch from.
+        sample_idx: int
+            Sample index. The first dimension of the image data will be indexed at this
+            value.
+        coords: Sequence of int
+            The coordinates that define the start of a patch.
+        patch_size: Sequence of int
+            The size of the patch in each spatial dimension.
+
+        Returns
+        -------
+        numpy.ndarray
+            The patch.
+        """
+        ...
+
+
+def default_patch_constr(
+    image_stack: ImageStack,
+    sample_idx: int,
+    channels: Sequence[int] | None,  # `channels = None` to select all channels
+    coords: Sequence[int],
+    patch_size: Sequence[int],
+) -> NDArray[Any]:
+    return image_stack.extract_channel_patch(
+        sample_idx=sample_idx,
+        channels=channels,
+        coords=coords,
+        patch_size=patch_size,
+    )
+
+
+# closure to create constructor funcs with particular multiscale_count and padding mode
+def lateral_context_patch_constr(
+    # TODO: will we stick with this as the parameter name
+    multiscale_count: int,
+    # TODO: add other modes?
+    padding_mode: Literal["reflect", "wrap"],
+) -> PatchConstructor:
+    """
+    Create a lateral context `PatchConstructor` for MicroSplit.
+
+    Parameters
+    ----------
+    multiscale_count : int
+        The number of multiscale inputs that will be created including the original
+        image size.
+    padding_mode : {"reflect", "wrap"}
+        How lateral context inputs will be padded at the edge of the image. See
+        [`numpy.pad`](https://numpy.org/devdocs/reference/generated/numpy.pad.html) for
+        more information.
+
+    Returns
+    -------
+    PatchConstructor
+        The patch constructor function. It will return patches with the dimensions
+        (C, L, (Z), Y, X) where L will be equal to `multiscale_count`, C is the number
+        of channels in the image, and (Z), Y, X are the patch size.
+    """
+
+    def constructor_func(
+        image_stack: ImageStack,
+        sample_idx: int,
+        channels: Sequence[int] | None,  # `channels = None` to select all channels
+        coords: Sequence[int],
+        patch_size: Sequence[int],
+    ) -> NDArray[Any]:
+        if channels is not None and len(channels) > 1:
+            raise NotImplementedError(
+                "Selecting multiple channels is currently not implemented for lateral "
+                "context patches. Select a single channel or pass `channels=None` to "
+                "select all channels."
+            )
+
+        shape = image_stack.data_shape
+        spatial_shape = shape[2:]
+        n_channels = shape[1] if channels is None else 1
+
+        # There will now be an additional lc dimension,
+        # this has to be handled correctly by the dataset
+        # TODO: maybe we want to limit this constructor to only images with 1 channel
+        #   then we can put LCs in the channel dimension
+        #   but not sure if this artificially limits potential use-cases
+        patch = np.zeros((n_channels, multiscale_count, *patch_size))
+        for scale in range(multiscale_count):
+            lc_patch_size = np.array(patch_size) * (2**scale)
+            lc_start = np.array(coords) + np.array(patch_size) // 2 - lc_patch_size // 2
+            lc_end = lc_start + np.array(lc_patch_size)
+
+            start_clipped = np.clip(
+                lc_start, np.zeros_like(spatial_shape), np.array(spatial_shape)
+            )
+            end_clipped = np.clip(
+                lc_end, np.zeros_like(spatial_shape), np.array(spatial_shape)
+            )
+            size_clipped = end_clipped - start_clipped
+
+            lc_patch = image_stack.extract_channel_patch(
+                sample_idx, channels, start_clipped, size_clipped
+            )
+            pad_before = start_clipped - lc_start
+            pad_after = lc_end - end_clipped
+            pad_width = np.concat(
+                [
+                    # zeros to not pad the channel axis
+                    np.zeros((1, 2), dtype=int),
+                    np.stack([pad_before, pad_after], axis=-1),
+                ]
+            )
+            lc_patch = np.pad(
+                lc_patch,
+                pad_width,
+                mode=padding_mode,
+            )
+            # TODO: test different downscaling? skimage suggests downscale_local_mean
+            lc_patch = resize(lc_patch, (n_channels, *patch_size))
+            patch[:, scale, ...] = lc_patch
+        return patch
+
+    return constructor_func

careamics/dataset_ng/patch_extractor/patch_extractor.py

@@ -0,0 +1,117 @@
+from collections.abc import Sequence
+from typing import Generic
+
+from numpy.typing import NDArray
+
+from ..image_stack import GenericImageStack
+from .patch_construction import PatchConstructor, default_patch_constr
+
+
+class PatchExtractor(Generic[GenericImageStack]):
+    """
+    A class for extracting patches from multiple image stacks.
+    """
+
+    def __init__(
+        self,
+        image_stacks: Sequence[GenericImageStack],
+        patch_constructor: PatchConstructor = default_patch_constr,
+    ):
+        self.patch_constructor = patch_constructor
+        self.image_stacks: list[GenericImageStack] = list(image_stacks)
+
+        # check all image stacks have the same number of dimensions
+        # check all image stacks have the same number of channels
+        self.n_spatial_dims = len(self.image_stacks[0].data_shape) - 2  # SC(Z)YX
+        self.n_channels = self.image_stacks[0].data_shape[1]
+        for i, image_stack in enumerate(image_stacks):
+            if (ndims := len(image_stack.data_shape) - 2) != self.n_spatial_dims:
+                raise ValueError(
+                    "All `ImageStack` objects in a `PatchExtractor` must have the same "
+                    "number of spatial dimensions. The first image stack is "
+                    f"{self.n_spatial_dims}D but found a {ndims}D image stack at index "
+                    f"{i}."
+                )
+            if (n_channels := image_stack.data_shape[1]) != self.n_channels:
+                raise ValueError(
+                    "All `ImageStack` objects in a `PatchExtractor` must have the same "
+                    f"number of channels. The first image stack has {self.n_channels} "
+                    f"but found an image stack with {n_channels} at index {i}."
+                )
+
+    def extract_patch(
+        self,
+        data_idx: int,
+        sample_idx: int,
+        coords: Sequence[int],
+        patch_size: Sequence[int],
+    ) -> NDArray:
+        """Extract a patch from the specified image stack across all channels.
+
+        Eqauivalent to calling `extract_channel_patch` with `channels=None`.
+
+        Parameters
+        ----------
+        data_idx : int
+            Index of the image stack to extract the patch from.
+        sample_idx : int
+            Sample index. The first dimension of the image data will be indexed at this
+            value.
+        coords : Sequence of int
+            The coordinates that define the start of a patch.
+        patch_size : Sequence of int
+            The size of the patch in each spatial dimension.
+
+        Returns
+        -------
+        numpy.ndarray
+            The extracted patch.
+        """
+        return self.extract_channel_patch(
+            data_idx=data_idx,
+            sample_idx=sample_idx,
+            channels=None,
+            coords=coords,
+            patch_size=patch_size,
+        )
+
+    def extract_channel_patch(
+        self,
+        data_idx: int,
+        sample_idx: int,
+        channels: Sequence[int] | None,
+        coords: Sequence[int],
+        patch_size: Sequence[int],
+    ) -> NDArray:
+        """Extract a patch from the specified image stack.
+
+        Parameters
+        ----------
+        data_idx : int
+            Index of the image stack to extract the patch from.
+        sample_idx : int
+            Sample index. The first dimension of the image data will be indexed at this
+            value.
+        channels : Sequence of int | None
+            Channels to extract. If `None`, all channels are extracted.
+        coords : Sequence of int
+            The coordinates that define the start of a patch.
+        patch_size : Sequence of int
+            The size of the patch in each spatial dimension.
+
+        Returns
+        -------
+        numpy.ndarray
+            The extracted patch.
+        """
+        return self.patch_constructor(
+            self.image_stacks[data_idx],
+            sample_idx=sample_idx,
+            channels=channels,
+            coords=coords,
+            patch_size=patch_size,
+        )
+
+    @property
+    def shapes(self) -> list[Sequence[int]]:
+        return [stack.data_shape for stack in self.image_stacks]