careamics 0.0.9__py3-none-any.whl → 0.0.10__py3-none-any.whl

careamics/dataset_ng/patch_extractor/image_stack/zarr_image_stack.py

@@ -0,0 +1,163 @@
+from collections.abc import Sequence
+from pathlib import Path
+from typing import Union
+
+import zarr
+import zarr.storage
+from numpy.typing import NDArray
+from typing_extensions import Self
+
+from careamics.dataset.dataset_utils import reshape_array
+
+
+class ZarrImageStack:
+    """
+    A class for extracting patches from an image stack that is stored as a zarr array.
+    """
+
+    # TODO: keeping store type narrow so that it has the path attribute
+    #   base zarr store is zarr.storage.Store, includes MemoryStore
+    def __init__(self, store: zarr.storage.FSStore, data_path: str, axes: str):
+        self._store = store
+        self._array = zarr.open_array(store=self._store, path=data_path, mode="r")
+        # TODO: validate axes
+        #   - must contain XY
+        #   - must be subset of STCZYX
+        self._original_axes = axes
+        self._original_data_shape: tuple[int, ...] = self._array.shape
+        self.data_shape = _reshaped_array_shape(axes, self._original_data_shape)
+        self.data_dtype = self._array.dtype
+
+    # TODO: not sure if this is useful
+    # TODO: potential solution using different metadata class for each ImageStack type
+    #   - see #399
+    @property
+    def source(self) -> Path:
+        return Path(self._store.path) / self._array.path
+
+    # automatically finds axes from metadata
+    # based on implementation in ome-zarr python package
+    # https://github.com/ome/ome-zarr-py/blob/f7096b0f2c1fc8edf4d7304e33caf8d279d99dbb/ome_zarr/reader.py#L294-L316
+    @classmethod
+    def from_ome_zarr(cls, path: Union[Path, str]) -> Self:
+        """
+        Will only use the first resolution in the hierarchy.
+
+        Assumes the path only contains 1 image.
+
+        Path can be to a local file, or it can be a URL to a zarr stored in the cloud.
+        """
+        store = zarr.storage.FSStore(url=path)
+        group = zarr.open_group(store=store, mode="r")
+        if "multiscales" not in group.attrs:
+            raise ValueError(
+                f"Zarr at path '{path}' cannot be loaded as an OME-Zarr because it "
+                "does not contain the attribute 'multiscales'."
+            )
+        # TODO: why is this a list of length 1? 0 index also in ome-zarr-python
+        # https://github.com/ome/ome-zarr-py/blob/f7096b0f2c1fc8edf4d7304e33caf8d279d99dbb/ome_zarr/reader.py#L286
+        multiscales_metadata = group.attrs["multiscales"][0]
+
+        # get axes
+        axes_list = [axes_data["name"] for axes_data in multiscales_metadata["axes"]]
+        axes = "".join(axes_list).upper()
+
+        first_multiscale_path = multiscales_metadata["datasets"][0]["path"]
+
+        return cls(store=store, data_path=first_multiscale_path, axes=axes)
+
+    def extract_patch(
+        self, sample_idx: int, coords: Sequence[int], patch_size: Sequence[int]
+    ) -> NDArray:
+        # original axes assumed to be any subset of STCZYX (containing YX), in any order
+        # arguments must be transformed to index data in original axes order
+        # to do this: loop through original axes and append correct index/slice
+        #   for each case: STCZYX
+        #   Note: if any axis is not present in original_axes it is skipped.
+
+        # guard for no S and T in original axes
+        if ("S" not in self._original_axes) and ("T" not in self._original_axes):
+            if sample_idx not in [0, -1]:
+                raise IndexError(
+                    f"Sample index {sample_idx} out of bounds for S axes with size "
+                    f"{self.data_shape[0]}"
+                )
+
+        patch_slice: list[Union[int, slice]] = []
+        for d in self._original_axes:
+            if d == "S":
+                patch_slice.append(self._get_S_index(sample_idx))
+            elif d == "T":
+                patch_slice.append(self._get_T_index(sample_idx))
+            elif d == "C":
+                patch_slice.append(slice(None, None))
+            elif d == "Z":
+                patch_slice.append(slice(coords[0], coords[0] + patch_size[0]))
+            elif d == "Y":
+                y_idx = 0 if "Z" not in self._original_axes else 1
+                patch_slice.append(
+                    slice(coords[y_idx], coords[y_idx] + patch_size[y_idx])
+                )
+            elif d == "X":
+                x_idx = 1 if "Z" not in self._original_axes else 2
+                patch_slice.append(
+                    slice(coords[x_idx], coords[x_idx] + patch_size[x_idx])
+                )
+            else:
+                raise ValueError(f"Unrecognised axis '{d}', axes should be in STCZYX.")
+
+        patch = self._array[tuple(patch_slice)]
+        patch_axes = self._original_axes.replace("S", "").replace("T", "")
+        return reshape_array(patch, patch_axes)[0]  # remove first sample dim
+
+    def _get_T_index(self, sample_idx: int) -> int:
+        """Get T index given `sample_idx`."""
+        if "T" not in self._original_axes:
+            raise ValueError("No 'T' axis specified in original data axes.")
+        axis_idx = self._original_axes.index("T")
+        dim = self._original_data_shape[axis_idx]
+
+        # new S' = S*T
+        # T_idx = S_idx' // T_size
+        # S_idx = S_idx' % T_size
+        # - floor divide finds the row
+        # - modulus finds how far along the row i.e. the column
+        return sample_idx % dim
+
+    def _get_S_index(self, sample_idx: int) -> int:
+        """Get S index given `sample_idx`."""
+        if "S" not in self._original_axes:
+            raise ValueError("No 'S' axis specified in original data axes.")
+        if "T" in self._original_axes:
+            T_axis_idx = self._original_axes.index("T")
+            T_dim = self._original_data_shape[T_axis_idx]
+
+            # new S' = S*T
+            # T_idx = S_idx' // T_size
+            # S_idx = S_idx' % T_size
+            # - floor divide finds the row
+            # - modulus finds how far along the row i.e. the column
+            return sample_idx // T_dim
+        else:
+            return sample_idx
+
+
+# TODO: move to dataset_utils, better name?
+def _reshaped_array_shape(axes: str, shape: Sequence[int]) -> tuple[int, ...]:
+    """Find resulting shape if reshaping array with given `axes` and `shape`."""
+    target_axes = "SCZYX"
+    target_shape = []
+    for d in target_axes:
+        if d in axes:
+            idx = axes.index(d)
+            target_shape.append(shape[idx])
+        elif (d != axes) and (d != "Z"):
+            target_shape.append(1)
+        else:
+            pass
+
+    if "T" in axes:
+        idx = axes.index("T")
+        target_shape[0] = target_shape[0] * shape[idx]
+
+    return tuple(target_shape)

careamics/dataset_ng/patch_extractor/image_stack_loader.py

@@ -0,0 +1,140 @@
+from collections.abc import Sequence
+from pathlib import Path
+from typing import (
+    Any,
+    Optional,
+    Protocol,
+    Union,
+)
+
+from numpy.typing import NDArray
+from typing_extensions import ParamSpec
+
+from careamics.config.support import SupportedData
+from careamics.file_io.read import ReadFunc
+from careamics.utils import BaseEnum
+
+from .image_stack import ImageStack, InMemoryImageStack, ZarrImageStack
+
+P = ParamSpec("P")
+
+
+class SupportedDataDev(str, BaseEnum):
+    ZARR = "zarr"
+
+
+class ImageStackLoader(Protocol[P]):
+    """
+    Protocol to define how `ImageStacks` should be loaded.
+
+    An `ImageStackLoader` is a callable that must take the `source` of the data as the
+    first argument, and the data `axes` as the second argument.
+
+    Additional `*args` and `**kwargs` are allowed, but they should only be used to
+    determine _how_ the data is loaded, not _what_ data is loaded. The `source`
+    argument has to wholly determine _what_ data is loaded, this is because,
+    downstream, both an input-source and a target-source have to be specified but they
+    will share `*args` and `**kwargs`.
+
+    An `ImageStackLoader` must return a sequence of the `ImageStack` class. This could
+    be a sequence of one of the existing concrete implementations, such as
+    `ZarrImageStack`, or a custom user defined `ImageStack`.
+
+    Example
+    -------
+    The following example demonstrates how an `ImageStackLoader` could be defined
+    for loading non-OME Zarr images. Returning a list of `ZarrImageStack` instances.
+
+    >>> from typing import TypedDict
+
+    >>> from zarr.storage import FSStore
+
+    >>> from careamics.config import DataConfig
+    >>> from careamics.dataset_ng.patch_extractor.image_stack import ZarrImageStack
+
+    >>> # Define a zarr source
+    >>> # It encompasses multiple arguments that determine what data will be loaded
+    >>> class ZarrSource(TypedDict):
+    ...     store: FSStore
+    ...     data_paths: Sequence[str]
+
+    >>> def custom_image_stack_loader(
+    ...     source: ZarrSource, axes: str, *args, **kwargs
+    ... ) -> list[ZarrImageStack]:
+    ...     image_stacks = [
+    ...         ZarrImageStack(store=source["store"], data_path=data_path, axes=axes)
+    ...         for data_path in source["data_paths"]
+    ...     ]
+    ...     return image_stacks
+
+    TODO: show example use in the `CAREamicsDataset`
+
+    The example above defines a `ZarrSource` dict because to determine _which_ ZARR
+    images will be loaded both a ZARR store and the internal data paths need to be
+    specified.
+    """
+
+    def __call__(
+        self, source: Any, axes: str, *args: P.args, **kwargs: P.kwargs
+    ) -> Sequence[ImageStack]: ...
+
+
+def from_arrays(
+    source: Sequence[NDArray], axes: str, *args, **kwargs
+) -> list[InMemoryImageStack]:
+    return [InMemoryImageStack.from_array(data=array, axes=axes) for array in source]
+
+
+# TODO: change source to directory path? Like in current implementation
+#   Advantage of having a list is the user can match input and target order themselves
+def from_tiff_files(
+    source: Sequence[Path], axes: str, *args, **kwargs
+) -> list[InMemoryImageStack]:
+    return [InMemoryImageStack.from_tiff(path=path, axes=axes) for path in source]
+
+
+# TODO: change source to directory path? Like in current implementation
+#   Advantage of having a list is the user can match input and target order themselves
+def from_custom_file_type(
+    source: Sequence[Path],
+    axes: str,
+    read_func: ReadFunc,
+    read_kwargs: dict[str, Any],
+    *args,
+    **kwargs,
+) -> list[InMemoryImageStack]:
+    return [
+        InMemoryImageStack.from_custom_file_type(
+            path=path,
+            axes=axes,
+            read_func=read_func,
+            **read_kwargs,
+        )
+        for path in source
+    ]
+
+
+def from_ome_zarr_files(
+    source: Sequence[Path], axes: str, *args, **kwargs
+) -> list[ZarrImageStack]:
+    # NOTE: axes is unused here, in from_ome_zarr the axes are automatically retrieved
+    return [ZarrImageStack.from_ome_zarr(path) for path in source]
+
+
+def get_image_stack_loader(
+    data_type: Union[SupportedData, SupportedDataDev],
+    image_stack_loader: Optional[ImageStackLoader] = None,
+) -> ImageStackLoader:
+    if data_type == SupportedData.ARRAY:
+        return from_arrays
+    elif data_type == SupportedData.TIFF:
+        return from_tiff_files
+    elif data_type == "zarr":  # temp for testing until zarr is added to SupportedData
+        return from_ome_zarr_files
+    elif data_type == SupportedData.CUSTOM:
+        if image_stack_loader is None:
+            return from_custom_file_type
+        else:
+            return image_stack_loader
+    else:
+        raise ValueError

careamics/dataset_ng/patch_extractor/patch_extractor.py

@@ -0,0 +1,29 @@
+from collections.abc import Sequence
+
+from numpy.typing import NDArray
+
+from .image_stack import ImageStack
+
+
+class PatchExtractor:
+    """
+    A class for extracting patches from multiple image stacks.
+    """
+
+    def __init__(self, image_stacks: Sequence[ImageStack]):
+        self.image_stacks: list[ImageStack] = list(image_stacks)
+
+    def extract_patch(
+        self,
+        data_idx: int,
+        sample_idx: int,
+        coords: Sequence[int],
+        patch_size: Sequence[int],
+    ) -> NDArray:
+        return self.image_stacks[data_idx].extract_patch(
+            sample_idx=sample_idx, coords=coords, patch_size=patch_size
+        )
+
+    @property
+    def shape(self):
+        return [stack.data_shape for stack in self.image_stacks]

careamics/dataset_ng/patch_extractor/patch_extractor_factory.py

@@ -0,0 +1,208 @@
+from collections.abc import Sequence
+from pathlib import Path
+from typing import Any, Literal, Optional, Union, overload
+
+from numpy.typing import NDArray
+from typing_extensions import ParamSpec
+
+from careamics.config.support import SupportedData
+from careamics.dataset_ng.patch_extractor import PatchExtractor
+from careamics.file_io.read import ReadFunc
+
+from .image_stack_loader import (
+    ImageStackLoader,
+    SupportedDataDev,
+    get_image_stack_loader,
+)
+
+P = ParamSpec("P")
+
+
+# Define overloads for each implemented ImageStackLoader case
+# Array case
+@overload
+def create_patch_extractor(
+    source: Sequence[NDArray], axes: str, data_type: Literal[SupportedData.ARRAY]
+) -> PatchExtractor:
+    """
+    Create a patch extractor from a sequence of numpy arrays.
+
+    Parameters
+    ----------
+    source: sequence of numpy.ndarray
+        The source arrays of the data.
+    data_config: DataConfig
+        The data configuration, `data_config.data_type` should have the value "array",
+        and `data_config.axes` should describe the axes of every array in the `source`.
+
+    Returns
+    -------
+    PatchExtractor
+    """
+
+
+# TIFF and ZARR case
+@overload
+def create_patch_extractor(
+    source: Sequence[Path],
+    axes: str,
+    data_type: Literal[SupportedData.TIFF, SupportedDataDev.ZARR],
+) -> PatchExtractor:
+    """
+    Create a patch extractor from a sequence of files that match our supported types.
+
+    Supported file types include TIFF and ZARR.
+
+    If the files are ZARR files they must follow the OME standard. If you have ZARR
+    files that do not follow the OME standard, see documentation on how to create
+    a custom `image_stack_loader`. (TODO: Add link).
+
+    Parameters
+    ----------
+    source: sequence of Path
+        The source files for the data.
+    data_config: DataConfig
+        The data configuration, `data_config.data_type` should have the value "tiff" or
+        "zarr", and `data_config.axes` should describe the axes of every image in the
+        `source`.
+
+    Returns
+    -------
+    PatchExtractor
+    """
+
+
+# Custom file type case (loaded into memory)
+@overload
+def create_patch_extractor(
+    source: Any,
+    axes: str,
+    data_type: Literal[SupportedData.CUSTOM],
+    *,
+    read_func: ReadFunc,
+    read_kwargs: dict[str, Any],
+) -> PatchExtractor:
+    """
+    Create a patch extractor from a sequence of files of a custom type.
+
+    Parameters
+    ----------
+    source: sequence of Path
+        The source files for the data.
+    data_config: DataConfig
+        The data configuration, `data_config.data_type` should have the value "custom".
+    read_func : ReadFunc
+        A function to read the custom file type, see the `ReadFunc` protocol.
+    read_kwargs : dict of {str: Any}
+        Kwargs that will be passed to the custom `read_func`.
+
+    Returns
+    -------
+    PatchExtractor
+    """
+
+
+# Custom ImageStackLoader case
+@overload
+def create_patch_extractor(
+    source: Any,
+    axes: str,
+    data_type: Literal[SupportedData.CUSTOM],
+    image_stack_loader: ImageStackLoader[P],
+    *args: P.args,
+    **kwargs: P.kwargs,
+) -> PatchExtractor:
+    """
+    Create a patch extractor using a custom `ImageStackLoader`.
+
+    The custom image stack loader must follow the `ImageStackLoader` protocol, i.e.
+    it must have the following function signature:
+    ```
+    def image_loader_example(
+        source: Any, data_config: DataConfig, *args, **kwargs
+    ) -> Sequence[ImageStack]:
+    ```
+
+    Parameters
+    ----------
+    source: sequence of Path
+        The source files for the data.
+    data_config: DataConfig
+        The data configuration, `data_config.data_type` should have the value "custom".
+    image_stack_loader: ImageStackLoader
+        A custom image stack loader callable.
+    *args: Any
+        Positional arguments that will be passed to the custom image stack loader.
+    **kwargs: Any
+        Keyword arguments that will be passed to the custom image stack loader.
+
+    Returns
+    -------
+    PatchExtractor
+    """
+
+
+# final overload to match the implentation function signature
+# Need this so it works later in the code
+#   (bec there aren't created overloads for create_patch_extractors below)
+@overload
+def create_patch_extractor(
+    source: Any,
+    axes: str,
+    data_type: Union[SupportedData, SupportedDataDev],
+    image_stack_loader: Optional[ImageStackLoader[P]] = None,
+    *args: P.args,
+    **kwargs: P.kwargs,
+) -> PatchExtractor: ...
+
+
+def create_patch_extractor(
+    source: Any,
+    axes: str,
+    data_type: Union[SupportedData, SupportedDataDev],
+    image_stack_loader: Optional[ImageStackLoader[P]] = None,
+    *args: P.args,
+    **kwargs: P.kwargs,
+) -> PatchExtractor:
+    # TODO: Do we need to catch data_config.data_type and source mismatches?
+    #   e.g. data_config.data_type is "array" but source is not Sequence[NDArray]
+    loader: ImageStackLoader[P] = get_image_stack_loader(data_type, image_stack_loader)
+    image_stacks = loader(source, axes, *args, **kwargs)
+    return PatchExtractor(image_stacks)
+
+
+# TODO: Remove this and just call `create_patch_extractor` within the Dataset class
+# Keeping for consistency for now
+def create_patch_extractors(
+    source: Any,
+    target_source: Optional[Any],
+    axes: str,
+    data_type: Union[SupportedData, SupportedDataDev],
+    image_stack_loader: Optional[ImageStackLoader] = None,
+    *args,
+    **kwargs,
+) -> tuple[PatchExtractor, Optional[PatchExtractor]]:
+
+    # --- data extractor
+    patch_extractor: PatchExtractor = create_patch_extractor(
+        source,
+        axes,
+        data_type,
+        image_stack_loader,
+        *args,
+        **kwargs,
+    )
+    # --- optional target extractor
+    if target_source is not None:
+        target_patch_extractor = create_patch_extractor(
+            target_source,
+            axes,
+            data_type,
+            image_stack_loader,
+            *args,
+            **kwargs,
+        )
+
+        return patch_extractor, target_patch_extractor
+
+    return patch_extractor, None

careamics/dataset_ng/patching_strategies/__init__.py

@@ -0,0 +1,11 @@
+__all__ = [
+    "FixedRandomPatchingStrategy",
+    "PatchSpecs",
+    "PatchingStrategy",
+    "RandomPatchingStrategy",
+    "SequentialPatchingStrategy",
+]
+
+from .patching_strategy_protocol import PatchingStrategy, PatchSpecs
+from .random_patching import FixedRandomPatchingStrategy, RandomPatchingStrategy
+from .sequential_patching import SequentialPatchingStrategy

careamics/dataset_ng/patching_strategies/patching_strategy_protocol.py

@@ -0,0 +1,82 @@
+"""A module to contain type definitions relating to patching strategies."""
+
+from collections.abc import Sequence
+from typing import Protocol, TypedDict
+
+
+class PatchSpecs(TypedDict):
+    """A dictionary that specifies a single patch in a series of `ImageStacks`.
+
+    Attributes
+    ----------
+    data_idx: int
+        Determines which `ImageStack` a patch belongs to, within a series of
+        `ImageStack`s.
+    sample_idx: int
+        Determines which sample a patch belongs to, within an `ImageStack`.
+    coords: sequence of int
+        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.
+    patch_size: sequence of int
+        The size of the patch. The sequence will have length 2 or 3, for 2D and 3D data
+        respectively.
+    """
+
+    data_idx: int
+    sample_idx: int
+    coords: Sequence[int]
+    patch_size: Sequence[int]
+
+
+class PatchingStrategy(Protocol):
+    """
+    An interface for patching strategies.
+
+    Patching strategies are a component of the `CAREamicsDataset`; they determine
+    how patches are extracted from the underlying data.
+
+    Attributes
+    ----------
+    n_patches: int
+        The number of patches that the patching strategy will return.
+
+    Methods
+    -------
+    get_patch_spec(index: int) -> PatchSpecs
+        Get a patch specification for a given patch index.
+    """
+
+    @property
+    def n_patches(self) -> int:
+        """
+        The number of patches that the patching strategy will return.
+
+        It also determines the maximum index that can be given to `get_patch_spec`,
+        and the length of the `CAREamicsDataset`.
+
+        Returns
+        -------
+        int
+            Number of patches.
+        """
+        ...
+
+    def get_patch_spec(self, index: int) -> PatchSpecs:
+        """
+        Get a patch specification for a given patch index.
+
+        This method is intended to be called from within the
+        `CAREamicsDataset.__getitem__`. The index will be passed through from this
+        method.
+
+        Parameters
+        ----------
+        index : int
+            A patch index.
+
+        Returns
+        -------
+        PatchSpecs
+            A dictionary that specifies a single patch in a series of `ImageStacks`.
+        """
+        ...