careamics 0.0.19__py3-none-any.whl

careamics/dataset_ng/factory.py

@@ -0,0 +1,180 @@
+from functools import partial
+from typing import Any
+
+from typing_extensions import ParamSpec
+
+from careamics.config.data.ng_data_config import NGDataConfig
+from careamics.config.support import SupportedData
+from careamics.file_io.read import ReadFunc
+
+from .dataset import CareamicsDataset
+from .image_stack import (
+    GenericImageStack,
+    ImageStack,
+)
+from .image_stack_loader import (
+    ImageStackLoader,
+    load_arrays,
+    load_custom_file,
+    load_czis,
+    load_iter_tiff,
+    load_tiffs,
+    load_zarrs,
+)
+from .patch_extractor import LimitFilesPatchExtractor, PatchExtractor
+
+P = ParamSpec("P")
+
+
+# convenience function but should use `create_dataloader` function instead
+# For lazy loading custom batch sampler also needs to be set.
+def create_dataset(
+    config: NGDataConfig,
+    inputs: Any,
+    targets: Any,
+    masks: Any = None,
+    read_func: ReadFunc | None = None,
+    read_kwargs: dict[str, Any] | None = None,
+    image_stack_loader: ImageStackLoader | None = None,
+    image_stack_loader_kwargs: dict[str, Any] | None = None,
+) -> CareamicsDataset[ImageStack]:
+    """
+    Convenience function to create the CAREamicsDataset.
+
+    Parameters
+    ----------
+    config : DataConfig or InferenceConfig
+        The data configuration.
+    inputs : Any
+        The input sources to the dataset.
+    targets : Any, optional
+        The target sources to the dataset.
+    masks : Any, optional
+        The mask sources used to filter patches.
+    read_func : ReadFunc, optional
+        A function that can that can be used to load custom data. This argument is
+        ignored unless the `data_type` in the `config` is "custom".
+    read_kwargs : dict of {str, Any}, optional
+        Additional key-word arguments to pass to the `read_func`.
+    image_stack_loader : ImageStackLoader, optional
+        A function for custom image stack loading. This argument is ignored unless the
+        `data_type` in the `config` is "custom".
+    image_stack_loader_kwargs : {str, Any}, optional
+        Additional key-word arguments to pass to the `image_stack_loader`.
+    """
+    image_stack_loader = select_image_stack_loader(
+        data_type=SupportedData(config.data_type),
+        in_memory=config.in_memory,
+        read_func=read_func,
+        read_kwargs=read_kwargs,
+        image_stack_loader=image_stack_loader,
+        image_stack_loader_kwargs=image_stack_loader_kwargs,
+    )
+    patch_extractor_type = select_patch_extractor_type(
+        data_type=SupportedData(config.data_type), in_memory=config.in_memory
+    )
+    input_extractor = init_patch_extractor(
+        patch_extractor_type, image_stack_loader, inputs, config.axes
+    )
+    if targets is not None:
+        target_extractor = init_patch_extractor(
+            patch_extractor_type, image_stack_loader, targets, config.axes
+        )
+    else:
+        target_extractor = None
+    if masks is not None:
+        mask_extractor = init_patch_extractor(
+            patch_extractor_type, image_stack_loader, masks, config.axes
+        )
+    else:
+        mask_extractor = None
+    return CareamicsDataset(
+        data_config=config,
+        input_extractor=input_extractor,
+        target_extractor=target_extractor,
+        mask_extractor=mask_extractor,
+    )
+
+
+def init_patch_extractor(
+    patch_extractor: type[PatchExtractor],
+    image_stack_loader: ImageStackLoader[..., GenericImageStack],
+    source: Any,
+    axes: str,
+) -> PatchExtractor[GenericImageStack]:
+    image_stacks = image_stack_loader(source, axes)
+    return patch_extractor(image_stacks)
+
+
+def select_patch_extractor_type(
+    data_type: SupportedData,
+    in_memory: bool,
+) -> type[PatchExtractor]:
+    """Select the appropriate PatchExtractor type based on data type and memory mode.
+
+    If `in_memory` is True, or `data_type` is ZARR or CZI, the standard
+    `PatchExtractor` is selected, otherwise the `LimitFilesPatchExtractor` will be used.
+
+    Parameters
+    ----------
+    data_type : SupportedData
+        The type of data being handled.
+    in_memory : bool
+        Indicates whether data is to be loaded into memory.
+
+    Returns
+    -------
+    type[PatchExtractor]
+        The selected PatchExtractor type.
+    """
+    if not in_memory and data_type in (SupportedData.TIFF, SupportedData.CUSTOM):
+        return LimitFilesPatchExtractor
+    else:
+        return PatchExtractor
+
+
+def select_image_stack_loader(
+    data_type: SupportedData,
+    in_memory: bool,
+    read_func: ReadFunc | None = None,
+    read_kwargs: dict[str, Any] | None = None,
+    image_stack_loader: ImageStackLoader | None = None,
+    image_stack_loader_kwargs: dict[str, Any] | None = None,
+) -> ImageStackLoader:
+    match data_type:
+        case SupportedData.ARRAY:
+            return load_arrays
+        case SupportedData.TIFF:
+            if in_memory:
+                return load_tiffs
+            else:
+                return load_iter_tiff
+        case SupportedData.CUSTOM:
+            if (read_func is not None) and (image_stack_loader is None):
+                read_kwargs = {} if read_kwargs is None else read_kwargs
+                return partial(
+                    load_custom_file, read_func=read_func, read_kwargs=read_kwargs
+                )
+            elif (read_func is None) and (image_stack_loader is not None):
+                image_stack_loader_kwargs = (
+                    {}
+                    if image_stack_loader_kwargs is None
+                    else image_stack_loader_kwargs
+                )
+                return partial(image_stack_loader, **image_stack_loader_kwargs)
+            else:
+                raise ValueError(
+                    "Found `data_type='custom'` **one** of `read_func` or "
+                    "`image_stack_loader` must be provided."
+                )
+        case SupportedData.ZARR:
+            # TODO: in_memory or not
+            return load_zarrs
+        case SupportedData.CZI:
+            # TODO: in_memory or not
+            return load_czis
+        case _:
+            raise NotImplementedError(
+                f"Selecting an image stack for data type '{data_type}' has not been "
+                "implemented yet."
+            )

careamics/dataset_ng/grouped_index_sampler.py

@@ -0,0 +1,73 @@
+"""Module for the `GroupedIndexSampler`."""
+
+from collections.abc import Iterator, Sequence
+from typing import Self
+
+import numpy as np
+from numpy.random import Generator, default_rng
+from torch.utils.data import Sampler
+
+from careamics.dataset_ng.dataset import CareamicsDataset
+
+
+class GroupedIndexSampler(Sampler):
+    """
+    A PyTorch Sampler iterates through groups of indices.
+
+    The order of the groups will be shuffled and the order of the indices within the
+    groups will be shuffled.
+
+    This sampler is useful for iterative file loading — one file should be loaded at a
+    time so indices belonging to the same file should be grouped, but the order of the
+    files and the order of the indices should be shuffled.
+    """
+
+    def __init__(self, grouped_indices: Sequence[Sequence[int]], rng: Generator | None):
+        """
+        Parameters
+        ----------
+        grouped_indices : Sequence of (Sequence of int)
+            The indices that should be iterated through in groups.
+        """
+        super().__init__()
+        if rng is None:
+            self.rng = default_rng()
+        else:
+            self.rng = rng
+        # TODO: validate indices are unique across groups
+        self.grouped_indices = grouped_indices
+
+    @classmethod
+    def from_dataset(
+        cls, dataset: CareamicsDataset, rng: Generator | None = None
+    ) -> Self:
+        """
+        Create the sampler from a CareamicsDataset.
+
+        The grouped indices will be retrieved from the dataset's patching strategy.
+
+        Parameters
+        ----------
+        dataset: CareamicsDataset
+            An instance of the CareamicsDataset to create the sampler for.
+        rng: numpy.random.Generator, optional
+            Numpy random number generator that can be used to seed the sampler.
+        """
+        n_data_samples = len(dataset.input_extractor.shapes)
+        grouped_indices: list[Sequence[int]] = [
+            dataset.patching_strategy.get_patch_indices(i)
+            for i in range(n_data_samples)
+        ]
+        return cls(grouped_indices=grouped_indices, rng=rng)
+
+    def __iter__(self) -> Iterator[int]:
+
+        # shuffle the groups and the sub groups but keep indices in a group adjacent
+        group_order = np.arange(len(self.grouped_indices))
+        self.rng.shuffle(group_order)
+        for group_idx in group_order:
+            group = self.grouped_indices[group_idx.item()]
+            index_order = np.arange(len(group))
+            self.rng.shuffle(index_order)
+            for idx in index_order:
+                yield group[idx.item()]

careamics/dataset_ng/image_stack/__init__.py

@@ -0,0 +1,14 @@
+__all__ = [
+    "CziImageStack",
+    "FileImageStack",
+    "GenericImageStack",
+    "ImageStack",
+    "InMemoryImageStack",
+    "ZarrImageStack",
+]
+
+from .czi_image_stack import CziImageStack
+from .file_image_stack import FileImageStack
+from .image_stack_protocol import GenericImageStack, ImageStack
+from .in_memory_image_stack import InMemoryImageStack
+from .zarr_image_stack import ZarrImageStack

careamics/dataset_ng/image_stack/czi_image_stack.py

@@ -0,0 +1,396 @@
+from __future__ import annotations
+
+import re
+from collections.abc import Iterable, Sequence
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Literal
+
+import numpy as np
+from numpy.typing import NDArray
+
+try:
+    from pylibCZIrw.czi import CziReader, Rectangle, open_czi
+
+    pyczi_available = True
+except ImportError:
+    pyczi_available = False
+
+if TYPE_CHECKING:
+    try:
+        from pylibCZIrw.czi import CziReader, Rectangle, open_czi
+    except ImportError:
+        CziReader = Rectangle = open_czi = None  # type: ignore
+
+
+class CziImageStack:
+    """
+    A class for extracting patches from an image stack that is stored as a CZI file.
+
+    Parameters
+    ----------
+    data_path : str or Path
+        Path to the CZI file.
+
+    scene : int, optional
+        Index of the scene to extract.
+
+        A single CZI file can contain multiple "scenes", which are stored alongside each
+        other at different coordinates in the image plane, often separated by empty
+        space. Specifying this argument will read only the single scene with that index
+        from the file. Think of it as cropping the CZI file to the region where that
+        scene is located.
+
+        If no scene index is specified, the entire image will be read. In case it
+        contains multiple scenes, they will all be present in the resulting image.
+        This is usually not desirable due to the empty space between them.
+        In general, only omit this argument or set it to `None` if you know that
+        your CZI file does not contain any scenes.
+
+        The static function :py:meth:`get_bounding_rectangles` can be used to find out
+        how many scenes a given file contains and what their bounding rectangles are.
+
+        The scene can also be provided as part of `data_path` by appending an `"@"`
+        followed by the scene index to the filename.
+
+    depth_axis : {"none", "Z", "T"}, default: "none"
+        Which axis to use as depth-axis for providing 3-D patches.
+
+        - `"none"`: Only provide 2-D patches. If a Z or T dimension is present in the
+          data, they will be combined into the sample dimension `S`.
+        - `"Z"`: Use the Z-axis as depth-axis. If a T axis is present as well, it will
+          be merged into the sample dimensions `S`.
+        - `"T"`: Use the T-axis as depth-axis. If a Z axis is present as well, it will
+          be merged into the sample dimensions `S`.
+
+    Attributes
+    ----------
+    source : Path
+        Path to the CZI file, including the scene index if specified.
+    data_path : Path
+        Path to the CZI file without scene index.
+    scene : int or None
+        Index of the scene to extract, or None if not specified.
+    data_shape : Sequence[int]
+        The shape of the data in the order `(SC(Z)YX)`.
+    axes : str
+        The axes in the CZI file corresponding to the dimensions in `data_shape`.
+        The following values can occur:
+
+        - "SCZYX" for 3-D volumes if `depth_axis` is `"Z"`.
+        - "SCTYX" for time-series if `depth_axis` is `"T"`.
+        - "SCYX" if `depth_axis` is `"none"`.
+
+        The axis `S` (sample) is the only one not mapping one-to-one to an axis in the
+        CZI file but combines all remaining axes present in the file into one.
+
+    Examples
+    --------
+    Create an image stack for the first scene in a CZI file:
+    >>> stack = CziImageStack("path/to/file.czi", scene=0)  # doctest: +SKIP
+
+    Alternatively, the scene index can also be provided as part of the filename.
+    This is mainly intended for re-creating an image stack from the `source` property:
+    >>> stack = CziImageStack("path/to/file.czi@0")  # doctest: +SKIP
+    >>> stack2 = CziImageStack(stack.source)  # doctest: +SKIP
+
+    If the CZI file contains a third dimension (Z or T) and you want to perform 3-D
+    denoising, you need to explicitly set `depth_axis` to `"Z"` or `"T"`:
+    >>> stack_2d = CziImageStack("path/to/file.czi", scene=0)  # doctest: +SKIP
+    >>> stack_2d.axes, stack_2d.data_shape  # doctest: +SKIP
+    ('SCYX', [40, 1, 512, 512])
+    >>> stack_3d = CziImageStack(  # doctest: +SKIP
+    ...     "path/to/file.czi", scene=0, depth_axis="Z"
+    ... )
+    >>> stack_3d.axes, stack_3d.data_shape  # doctest: +SKIP
+    ('SCZYX', [4, 1, 10, 512, 512])
+    """
+
+    def __init__(
+        self,
+        data_path: str | Path,
+        scene: int | None = None,
+        depth_axis: Literal["none", "Z", "T"] = "none",
+    ) -> None:
+        if not pyczi_available:
+            raise ImportError(
+                "The CZI image stack requires the `pylibCZIrw` package to be installed."
+                " Please install it with `pip install careamics[czi]`."
+            )
+
+        _data_path = Path(data_path)
+
+        # Check for scene encoded in filename.
+        # Normally, file path and scene should be provided as separate arguments but
+        # we would also like to support using the `source` property to re-create the
+        # CZI image stack. In this case, the scene index is encoded in the file path.
+        scene_matches = re.match(r"^(.*)@(\d+)$", _data_path.name)
+        if scene_matches:
+            if scene is not None:
+                raise ValueError(
+                    f"Scene index is specified in the filename ({_data_path.name}) and "
+                    f"as an argument ({scene}). Please specify only one."
+                )
+            _data_path = _data_path.parent / scene_matches.group(1)
+            scene = int(scene_matches.group(2))
+
+        # Set variables
+        self.data_path = _data_path
+        self.scene = scene
+        self._depth_axis = depth_axis
+
+        # Open CZI file
+        self._czi = CziReader(str(self.data_path))
+
+        # Determine metadata
+        self.axes, self.data_shape, self._bounding_rectangle, self._sample_axes = (
+            self._get_shape()
+        )
+        self.data_dtype = np.float32
+
+    def __del__(self):
+        if hasattr(self, "_czi"):
+            # Close CZI file
+            self._czi.close()
+
+    def __getstate__(self) -> dict[str, Any]:
+        # Remove CziReader object from state to avoid pickling issues
+        state = self.__dict__.copy()
+        del state["_czi"]
+        return state
+
+    def __setstate__(self, state: dict[str, Any]) -> None:
+        # Reopen CZI file after unpickling
+        self.__dict__.update(state)
+        self._czi = CziReader(str(self.data_path))
+
+    # TODO: we append the scene index to the file name
+    #       - not sure if this is a good approach
+    @property
+    def source(self) -> Path:
+        filename = self.data_path.name
+        if self.scene is not None:
+            filename = f"{filename}@{self.scene}"
+        return self.data_path.parent / filename
+
+    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:
+        # 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."
+                    )
+
+        # Determine 3rd dimension (T, Z or none)
+        if len(coords) == 3:
+            if len(self.axes) != 5:
+                raise ValueError(
+                    f"Requested a 3D patch from a 2D image stack with axes {self.axes}."
+                )
+            third_dim = self.axes[2]
+            third_dim_offset, third_dim_size = coords[0], patch_size[0]
+        else:
+            if len(self.axes) != 4:
+                raise ValueError(
+                    f"Requested a 2D patch from a 3D image stack with axes {self.axes}."
+                )
+            third_dim = None
+            third_dim_offset, third_dim_size = 0, 1
+
+        # Set up ROI to extract from each plane as (x, y, w, h)
+        roi = (
+            self._bounding_rectangle.x + coords[-1],
+            self._bounding_rectangle.y + coords[-2],
+            patch_size[-1],
+            patch_size[-2],
+        )
+
+        # Create output array of shape (C, Z, Y, X)
+        n_channels = self.data_shape[1] if channels is None else len(channels)
+        patch = np.empty(
+            (n_channels, third_dim_size, *patch_size[-2:]), dtype=np.float32
+        )
+
+        # Set up plane to index `sample_idx`
+        sample_shape = list(self._sample_axes.values())
+        sample_indices = np.unravel_index(sample_idx, sample_shape)
+        plane = {
+            dimension: int(index)
+            for dimension, index in zip(
+                self._sample_axes.keys(), sample_indices, strict=False
+            )
+        }
+
+        # Read XY planes sequentially
+        channel_iter: Iterable
+        if channels is None:
+            channel_iter = range(self.data_shape[1])  # iter over number of requested C
+        else:
+            channel_iter = list(channels)
+
+        # for each channel
+        for patch_channel, data_channel in enumerate(channel_iter):
+            # pull plane with the given channel and 3rd dim index
+            for third_dim_index in range(third_dim_size):
+                plane["C"] = data_channel
+                if third_dim is not None:
+                    plane[third_dim] = third_dim_offset + third_dim_index
+
+                # read plane
+                extracted_roi = self._czi.read(roi=roi, plane=plane, scene=self.scene)
+                if extracted_roi.ndim == 3:
+                    if extracted_roi.shape[-1] > 1:
+                        raise ValueError(
+                            "CZI files with RGB channels are currently not supported."
+                        )
+
+                    # remove channel dimension
+                    extracted_roi = extracted_roi.squeeze(-1)
+
+                # add requested channel into the patch
+                patch[patch_channel, third_dim_index] = extracted_roi
+
+        # Remove dummy 3rd dimension for 2-D data
+        if third_dim is None:
+            patch = patch.squeeze(1)
+
+        return patch
+
+    def _get_shape(self) -> tuple[str, list[int], Rectangle, dict[str, int]]:
+        """Determines the shape of the selected scene.
+
+        Returns
+        -------
+        axes : str
+            String specifying the axis order. Examples:
+
+            - "SCZYX" for 3-D volumes if `depth_axis` is `"Z"`.
+            - "SCTYX" for time-series if `depth_axis` is `"T"`.
+            - "SCYX" if `depth_axis` is `"none"`.
+
+            The axis `S` is the sample dimension and combines all remaining axes
+            present in the data.
+
+        shape : list[int]
+            The size of each axis, in the order listed in `axes`.
+
+        bounding_rectangle : Rectangle
+            The bounding rectangle of the scene in pixels. The rectangle is
+            defined by its top-left corner (x, y) and its width and height (w, h).
+
+        sample_axes : dict[str, int]
+            A dictionary with information about the remaining axes used for the
+            sample dimension.
+            The keys are the axis names (e.g., "T", "Z") and the values are their
+            respective sizes.
+        """
+        # Get CZI dimensions
+        total_bbox = self._czi.total_bounding_box_no_pyramid
+        if self.scene is None:
+            bounding_rectangle = self._czi.total_bounding_rectangle_no_pyramid
+        else:
+            bounding_rectangle = self._czi.scenes_bounding_rectangle_no_pyramid[
+                self.scene
+            ]
+
+        # Determine if T and Z axis are present
+        # Note: An axis of size 1 is as good as no axis since we cannot use it for 3-D
+        # denoising.
+        has_time = "T" in total_bbox and (total_bbox["T"][1] - total_bbox["T"][0]) > 1
+        has_depth = "Z" in total_bbox and (total_bbox["Z"][1] - total_bbox["Z"][0]) > 1
+
+        # Determine axis order depending on `depth_axis`
+        if self._depth_axis == "Z":
+            axes = "SCZYX"
+            if not has_depth:
+                raise RuntimeError(
+                    f"The CZI file {self.data_path} does not contain a Z axis to use "
+                    'for 3-D denoising. Consider setting `axes="YX"` or '
+                    '`depth_axis="none"` to perform 2-D denoising instead.'
+                )
+        elif self._depth_axis == "T":
+            axes = "SCTYX"
+            if not has_time:
+                raise RuntimeError(
+                    f"The CZI file {self.data_path} does not contain a T axis to use "
+                    'for 3-D denoising. Consider setting `axes="YX"` or '
+                    '`depth_axis="none"` to perform 2-D denoising instead.'
+                )
+        else:
+            axes = "SCYX"
+
+        # Calculcate size of sample dimension S, combining all axes not used elsewhere.
+        # This could, for example, be a time axis. If we only perform 2-D denoising, a
+        # potentially present Z axis would also be used as sample dimension. If both,
+        # T and Z, are present, both need to be combined into the sample dimension.
+        # The same needs to be done to any other potentially present axis in the CZI
+        # file which is not a spatial or channel axis.
+        # The following code calculates the size of the combined sample axis.
+        sample_axes = {}
+        sample_size = 1
+        for dimension, (start, end) in total_bbox.items():
+            if dimension not in axes:
+                sample_axes[dimension] = end - start
+                sample_size *= end - start
+
+        # Determine data shape
+        shape = []
+        for dimension in axes:
+            if dimension == "S":
+                shape.append(sample_size)
+            elif dimension == "Y":
+                shape.append(bounding_rectangle.h)
+            elif dimension == "X":
+                shape.append(bounding_rectangle.w)
+            elif dimension in total_bbox:
+                shape.append(total_bbox[dimension][1] - total_bbox[dimension][0])
+            else:
+                shape.append(1)
+
+        return axes, shape, bounding_rectangle, sample_axes
+
+    @classmethod
+    def get_bounding_rectangles(
+        cls, czi: Path | str | CziReader
+    ) -> dict[int | None, Rectangle]:
+        """Gets the bounding rectangles of all scenes in a CZI file.
+
+        Parameters
+        ----------
+        czi : Path or str or pyczi.CziReader
+            Path to the CZI file or an already opened file as CziReader object.
+
+        Returns
+        -------
+        dict[int | None, Rectangle]
+            A dictionary mapping scene indices to their bounding rectangles in the
+            format `(x, y, w, h)`.
+            If no scenes are present in the CZI file, the returned dictionary will
+            have only one entry with key `None`, whose bounding rectangle covers the
+            entire image.
+        """
+        if not isinstance(czi, CziReader):
+            with open_czi(str(czi)) as czi_reader:
+                return cls.get_bounding_rectangles(czi_reader)
+
+        scenes_bounding_rectangle = czi.scenes_bounding_rectangle_no_pyramid
+        if len(scenes_bounding_rectangle) >= 1:
+            # Ensure keys are int | None for type compatibility
+            return {int(k): v for k, v in scenes_bounding_rectangle.items()}
+        else:
+            return {None: czi.total_bounding_rectangle_no_pyramid}