careamics 0.0.19__py3-none-any.whl

careamics/dataset_ng/image_stack/file_image_stack.py

@@ -0,0 +1,140 @@
+from collections.abc import Sequence
+from pathlib import Path
+from typing import Any, Self
+
+import numpy as np
+import tifffile
+from numpy.typing import DTypeLike, NDArray
+
+from careamics.dataset.dataset_utils import reshape_array
+from careamics.file_io.read import ReadFunc, read_tiff
+
+from .image_utils.image_stack_utils import channel_slice, pad_patch, reshape_array_shape
+
+
+class FileImageStack:
+    """
+    An ImageStack implementation for data that is coming from a file.
+
+    The data will not be loaded until the `load` method is called. The `close` method
+    can be used to remove the internal reference to the data.
+    """
+
+    def __init__(
+        self,
+        source: Path,
+        axes: str,
+        data_shape: tuple[int, ...],
+        data_dtype: DTypeLike,
+        read_func: ReadFunc,
+        read_kwargs: dict[str, Any] | Any = None,
+    ):
+        self.source = source
+        self.axes = axes
+        self.data_shape = data_shape
+        self.data_dtype = data_dtype
+        self.read_func = read_func
+        self.read_kwargs = read_kwargs
+        self._data: NDArray | None = None
+
+    def extract_patch(
+        self, sample_idx: int, coords: Sequence[int], patch_size: Sequence[int]
+    ) -> NDArray:
+        return self.extract_channel_patch(sample_idx, None, coords, patch_size)
+
+    def extract_channel_patch(
+        self,
+        sample_idx: int,
+        channels: Sequence[int] | None,  # `channels = None` to select all channels
+        coords: Sequence[int],
+        patch_size: Sequence[int],
+    ) -> NDArray:
+        if self._data is None:
+            raise ValueError(
+                "Cannot extract patch because data has not been loaded from "
+                f"'{self.source}', the `load` method must be called first."
+            )
+
+        if (coord_dims := len(coords)) != (patch_dims := len(patch_size)):
+            raise ValueError(
+                "Patch coordinates and patch size must have the same dimensions but "
+                f"found {coord_dims} and {patch_dims}."
+            )
+
+        # check that channels are within bounds
+        if channels is not None:
+            max_channel = self.data_shape[1] - 1  # channel is second dimension
+            for ch in channels:
+                if ch > max_channel:
+                    raise ValueError(
+                        f"Channel index {ch} is out of bounds for data with "
+                        f"{self.data_shape[1]} channels. Check the provided `channels` "
+                        f"parameter in the configuration for erroneous channel "
+                        f"indices."
+                    )
+
+        patch_data = self._data[
+            (
+                sample_idx,  # type: ignore
+                # use channel slice so that channel dimension is kept
+                channel_slice(channels),  # type: ignore
+                *[
+                    slice(
+                        np.clip(c, 0, self.data_shape[2 + i]),
+                        np.clip(c + ps, 0, self.data_shape[2 + i]),
+                    )
+                    for i, (c, ps) in enumerate(zip(coords, patch_size, strict=False))
+                ],  # type: ignore
+            )  # type: ignore
+        ]
+        patch = pad_patch(coords, patch_size, self.data_shape, patch_data)
+
+        return patch
+
+    def load(self):
+        """Load the data stored in a file."""
+        data = self.read_func(self.source)
+        self._data = reshape_array(data, self.axes)
+
+    # TODO: maybe this should be called something else
+    def close(self):
+        """Remove the internal reference to the data to clear up memory."""
+        # will get cleaned up by the garbage collector since there is no longer a ref
+        self._data = None
+
+    @property
+    def is_loaded(self):
+        return self._data is not None
+
+    @classmethod
+    def from_tiff(
+        cls,
+        path: Path,
+        axes: str,
+    ) -> Self:
+        """
+        Construct the `ImageStack` from a TIFF file.
+
+        Parameters
+        ----------
+        path : Path
+            Path to the TIFF file.
+        axes : str
+            The original axes of the data, must be a subset of STCZYX.
+
+        Returns
+        -------
+        Self
+            The `ImageStack` with the underlying data being from a TIFF file.
+        """
+        # TODO: think this is correct but need more examples to test
+        file = tifffile.TiffFile(path)
+        data_shape = reshape_array_shape(axes, file.series[0].shape)
+        dtype = file.series[0].dtype
+        return cls(
+            source=path,
+            axes=axes,
+            data_shape=data_shape,
+            data_dtype=dtype,
+            read_func=read_tiff,
+        )

careamics/dataset_ng/image_stack/image_stack_protocol.py

@@ -0,0 +1,93 @@
+from collections.abc import Sequence
+from pathlib import Path
+from typing import Literal, Protocol, TypeVar, Union
+
+from numpy.typing import DTypeLike, NDArray
+
+
+class ImageStack(Protocol):
+    """
+    An interface for extracting patches from an image stack.
+
+    Attributes
+    ----------
+    source: Path or "array"
+        Origin of the image data.
+    data_shape: Sequence[int]
+        The shape of the data, it is expected to be in the order (SC(Z)YX).
+    data_dtype: DTypeLike
+        The data type of the image data.
+    """
+
+    @property
+    def source(self) -> Union[str, Path, Literal["array"]]: ...
+
+    """Source of the image data."""
+
+    @property
+    def data_shape(self) -> Sequence[int]: ...
+
+    """Shape of the image data."""
+
+    @property
+    def data_dtype(self) -> DTypeLike: ...
+
+    """Data type of the image data."""
+
+    def extract_patch(
+        self, sample_idx: int, coords: Sequence[int], patch_size: Sequence[int]
+    ) -> NDArray:
+        """
+        Extract a patch for a given sample within the image stack.
+
+        Parameters
+        ----------
+        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
+            A patch of the image data from a particlular sample. It will have the
+            dimensions C(Z)YX.
+        """
+        ...
+
+    def extract_channel_patch(
+        self,
+        sample_idx: int,
+        channels: Sequence[int] | None,
+        coords: Sequence[int],
+        patch_size: Sequence[int],
+    ) -> NDArray:
+        """
+        Extract a patch of a single channel for a given sample within the image stack.
+
+        Parameters
+        ----------
+        sample_idx: int
+            Sample index. The first dimension of the image data will be indexed at this
+            value.
+        channels: Sequence[int] | None
+            Channel indices to extract. If `None` is given all channels will be
+            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
+            A patch of the image data from a particlular sample. It will have the
+            dimensions C(Z)YX.
+        """
+        ...
+
+
+GenericImageStack = TypeVar("GenericImageStack", bound=ImageStack, covariant=True)

careamics/dataset_ng/image_stack/image_utils/__init__.py

@@ -0,0 +1,6 @@
+"""Image stack utility functions."""
+
+__all__ = ["channel_slice", "pad_patch", "reshape_array_shape"]
+
+
+from .image_stack_utils import channel_slice, pad_patch, reshape_array_shape

careamics/dataset_ng/image_stack/image_utils/image_stack_utils.py

@@ -0,0 +1,125 @@
+from collections.abc import Sequence
+from types import EllipsisType
+from typing import TypeVar
+
+import numpy as np
+from numpy.typing import NDArray
+
+T = TypeVar("T", bound=np.generic)
+
+
+def channel_slice(
+    channels: Sequence[int] | None,
+) -> EllipsisType | Sequence[int]:
+    """Create a slice or sequence for indexing channels while preserving dimensions.
+
+    Parameters
+    ----------
+    channels : Sequence[int] | None
+        The channel indices to select, or None to select all channels.
+
+    Returns
+    -------
+    EllipsisType | Sequence[int]
+        An indexing object that can be used to index the channel dimension while
+        preserving it.
+    """
+    if channels is None:
+        return ...
+
+    if len(channels) == 0:
+        raise ValueError("Channel index sequence cannot be empty.")
+
+    return channels
+
+
+# TODO: add tests
+# TODO: move to dataset_utils, better name?
+def reshape_array_shape(
+    original_axes: str, shape: Sequence[int], add_singleton: bool = True
+) -> tuple[int, ...]:
+    """Find resulting shape if reshaping array to SC(Z)YX.
+
+    If `T` is present in the original axes, its size is multiplied into `S`, as both
+    axes are multiplexed.
+
+    Setting `add_singleton` to `False` will only include axes that are present in
+    `original_axes` in the output shape.
+
+    Parameters
+    ----------
+    original_axes : str
+        The axes of the original array, e.g. "TCZYX", "SCYX", etc.
+    shape : Sequence[int]
+        The shape of the original array.
+    add_singleton : bool, default=True
+        Whether to add singleton dimensions for missing axes. When `False`, only axes
+        present in `original_axes` will be included in the output shape. When `True`,
+        missing mandatory axes (`S` and `C`) will be added as singleton dimensions.
+    """
+    target_axes = "SCZYX"
+    target_shape = []
+    for d in target_axes:
+        if d in original_axes:
+            idx = original_axes.index(d)
+            target_shape.append(shape[idx])
+        elif d != "Z":
+            if add_singleton:
+                target_shape.append(1)
+
+    if "T" in original_axes:
+        idx = original_axes.index("T")
+        if "S" in original_axes or add_singleton:
+            target_shape[0] = target_shape[0] * shape[idx]
+        else:
+            target_shape.insert(0, shape[idx])
+
+    return tuple(target_shape)
+
+
+def pad_patch(
+    coords: Sequence[int],
+    patch_size: Sequence[int],
+    data_shape: Sequence[int],
+    patch_data: NDArray[T],
+) -> NDArray[T]:
+    """
+    Pad patch data with zeros where it is outside the bounds of it's source image.
+
+    This ensures the patch data is contained in an array with the expected patch size.
+
+    If `coords` are negative, the start of the patch will be padded with zeros up until
+    where the start of the image would be, and this is where the patch data starts.
+
+    If the `coords + patch_size` are greater than the bounds of the image then the
+    end of the patch will be filled with zeros.
+
+    Parameters
+    ----------
+    coords : Sequence[int]
+        The coordinates that describe where the patch starts in the spatial dimension of
+        the image
+    patch_size : Sequence[int]
+        The size of the patch in the spatial dimensions.
+    data_shape : Sequence[int]
+        The shape of the image the patch originates from, must be in the format SC(Z)YX.
+    patch_data : NDArray[T]
+        The patch data to be padded.
+
+    Returns
+    -------
+    NDArray[T]
+        The resulting padded patch.
+    """
+    coords_ = np.array(coords)
+    patch = np.zeros((patch_data.shape[0], *patch_size), dtype=patch_data.dtype)
+    # data start will be zero unless coords are negative
+    data_start = np.clip(coords_, 0, None) - coords_
+    data_end = data_start + np.array(patch_data.shape[1:])
+    patch[
+        (
+            slice(None, None, None),  # channel slice
+            *tuple(slice(s, t) for s, t in zip(data_start, data_end, strict=False)),
+        )
+    ] = patch_data
+    return patch

careamics/dataset_ng/image_stack/in_memory_image_stack.py

@@ -0,0 +1,93 @@
+from collections.abc import Sequence
+from pathlib import Path
+from typing import Any, Literal, Self, Union
+
+import numpy as np
+from numpy.typing import DTypeLike, NDArray
+
+from careamics.dataset.dataset_utils import reshape_array
+from careamics.file_io.read import ReadFunc, read_tiff
+
+from .image_utils.image_stack_utils import channel_slice, pad_patch
+
+
+class InMemoryImageStack:
+    """
+    A class for extracting patches from an image stack that has been loaded into memory.
+    """
+
+    def __init__(self, source: Union[Path, Literal["array"]], data: NDArray):
+        self.source: Union[str, Path, Literal["array"]] = source
+        # data expected to be in SC(Z)YX shape, reason to use from_array constructor
+        self._data: NDArray = data
+        self.data_shape: Sequence[int] = self._data.shape
+        self.data_dtype: DTypeLike = self._data.dtype
+
+    def extract_patch(
+        self, sample_idx: int, coords: Sequence[int], patch_size: Sequence[int]
+    ) -> NDArray:
+        return self.extract_channel_patch(sample_idx, None, coords, patch_size)
+
+    def extract_channel_patch(
+        self,
+        sample_idx: int,
+        channels: Sequence[int] | None,  # `channels = None` to select all channels
+        coords: Sequence[int],
+        patch_size: Sequence[int],
+    ) -> NDArray:
+        if (coord_dims := len(coords)) != (patch_dims := len(patch_size)):
+            raise ValueError(
+                "Patch coordinates and patch size must have the same dimensions but "
+                f"found {coord_dims} ({coords}) and {patch_dims} ({patch_size})."
+            )
+
+        # check that channels are within bounds
+        if channels is not None:
+            max_channel = self.data_shape[1] - 1  # channel is second dimension
+            for ch in channels:
+                if ch > max_channel:
+                    raise ValueError(
+                        f"Channel index {ch} is out of bounds for data with "
+                        f"{self.data_shape[1]} channels. Check the provided `channels` "
+                        f"parameter in the configuration for erroneous channel "
+                        f"indices."
+                    )
+
+        # TODO: test for 2D or 3D?
+
+        patch_data = self._data[
+            (
+                sample_idx,  # type: ignore
+                # use channel slice so that channel dimension is kept
+                channel_slice(channels),  # type: ignore
+                *[
+                    slice(
+                        np.clip(c, 0, self.data_shape[2 + i]),
+                        np.clip(c + ps, 0, self.data_shape[2 + i]),
+                    )
+                    for i, (c, ps) in enumerate(zip(coords, patch_size, strict=False))
+                ],  # type: ignore
+            )  # type: ignore
+        ]
+        patch = pad_patch(coords, patch_size, self.data_shape, patch_data)
+
+        return patch
+
+    @classmethod
+    def from_array(cls, data: NDArray, axes: str) -> Self:
+        data = reshape_array(data, axes)
+        return cls(source="array", data=data)
+
+    @classmethod
+    def from_tiff(cls, path: Path, axes: str) -> Self:
+        data = read_tiff(path)
+        data = reshape_array(data, axes)
+        return cls(source=path, data=data)
+
+    @classmethod
+    def from_custom_file_type(
+        cls, path: Path, axes: str, read_func: ReadFunc, **read_kwargs: Any
+    ) -> Self:
+        data = read_func(path, **read_kwargs)
+        data = reshape_array(data, axes)
+        return cls(source=path, data=data)

careamics/dataset_ng/image_stack/zarr_image_stack.py

@@ -0,0 +1,170 @@
+from collections.abc import Sequence
+
+import zarr
+from numpy.typing import DTypeLike, NDArray
+
+from careamics.dataset.dataset_utils import reshape_array
+
+from .image_utils.image_stack_utils import channel_slice, pad_patch, reshape_array_shape
+
+
+class ZarrImageStack:
+    """
+    A class for extracting patches from an image stack that is stored as a zarr array.
+    """
+
+    def __init__(self, group: zarr.Group, data_path: str, axes: str):
+        if not isinstance(group, zarr.Group):
+            raise TypeError(f"group must be a zarr.Group instance, got {type(group)}.")
+
+        self._group = group
+        self._store = str(group.store_path)
+        try:
+            self._array = group[data_path]
+        except KeyError as e:
+            raise ValueError(
+                f"Did not find array at '{data_path}' in store '{self._store}'."
+            ) from e
+
+        if not isinstance(self._array, zarr.Array):
+            raise TypeError(
+                f"data at path '{data_path}' must be a zarr.Array instance, "
+                f"got {type(self._array)}."
+            )
+
+        self._source = self._array.store_path
+
+        # 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 = reshape_array_shape(axes, self._original_data_shape)
+        self._data_dtype = self._array.dtype
+        self._chunk_size = reshape_array_shape(
+            axes, self._array.chunks, add_singleton=False
+        )
+        self._shard_size = (
+            reshape_array_shape(axes, self._array.shards, add_singleton=False)
+            if self._array.shards is not None
+            else None
+        )
+
+    # Used to identify the source of the data and write to similar path during pred
+    @property
+    def source(self) -> str:
+        # e.g. file://data/bsd68.zarr/train/
+        return str(self._source)
+
+    @property
+    def chunks(self) -> Sequence[int]:
+        """Chunks size in the order of data_shape (SC(Z)YX)."""
+        return self._chunk_size
+
+    @property
+    def shards(self) -> Sequence[int] | None:
+        """Shard size in the order of data_shape (SC(Z)YX)."""
+        return self._shard_size
+
+    @property
+    def data_dtype(self) -> DTypeLike:
+        return self._data_dtype
+
+    def extract_patch(
+        self, sample_idx: int, coords: Sequence[int], patch_size: Sequence[int]
+    ) -> NDArray:
+        return self.extract_channel_patch(sample_idx, None, coords, patch_size)
+
+    def extract_channel_patch(
+        self,
+        sample_idx: int,
+        channels: Sequence[int] | None,  # `channels = None` to select all channels,
+        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]}"
+                )
+
+        # check that channels are within bounds
+        if channels is not None:
+            max_channel = self.data_shape[1] - 1  # channel is second dimension
+            for ch in channels:
+                if ch > max_channel:
+                    raise ValueError(
+                        f"Channel index {ch} is out of bounds for data with "
+                        f"{self.data_shape[1]} channels. Check the provided `channels` "
+                        f"parameter in the configuration for erroneous channel "
+                        f"indices."
+                    )
+
+        patch_slice: list[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(channel_slice(channels))  # type: ignore
+            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_data: NDArray = self._array[tuple(patch_slice)]  # type: ignore
+        patch_axes = self._original_axes.replace("S", "").replace("T", "")
+        patch_data = reshape_array(patch_data, patch_axes)[0]  # remove first sample dim
+        patch = pad_patch(coords, patch_size, self.data_shape, patch_data)
+
+        return patch
+
+    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

careamics/dataset_ng/image_stack_loader/__init__.py

@@ -0,0 +1,19 @@
+__all__ = [
+    "ImageStackLoader",
+    "load_arrays",
+    "load_custom_file",
+    "load_czis",
+    "load_iter_tiff",
+    "load_tiffs",
+    "load_zarrs",
+]
+
+from .image_stack_loader_protocol import ImageStackLoader
+from .image_stack_loaders import (
+    load_arrays,
+    load_custom_file,
+    load_czis,
+    load_iter_tiff,
+    load_tiffs,
+    load_zarrs,
+)