careamics 0.0.11__py3-none-any.whl → 0.0.13__py3-none-any.whl

careamics/dataset_ng/patch_extractor/patch_extractor_factory.py

@@ -1,29 +1,30 @@
 from collections.abc import Sequence
 from pathlib import Path
-from typing import Any, Literal, Optional, Union, overload
+from typing import Any, Literal
 
 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 import (
+    CziImageStack,
+    GenericImageStack,
+    InMemoryImageStack,
+    ZarrImageStack,
+)
 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:
+def create_array_extractor(
+    source: Sequence[NDArray[Any]], axes: str
+) -> PatchExtractor[InMemoryImageStack]:
     """
     Create a patch extractor from a sequence of numpy arrays.
 
@@ -31,57 +32,119 @@ def create_patch_extractor(
     ----------
     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`.
+    axes: str
+        The original axes of the data, must be a subset of "STCZYX".
 
     Returns
     -------
     PatchExtractor
     """
+    image_stacks = [
+        InMemoryImageStack.from_array(data=array, axes=axes) for array in source
+    ]
+    return PatchExtractor(image_stacks)
+
+
+# TIFF case
+def create_tiff_extractor(
+    source: Sequence[Path], axes: str
+) -> PatchExtractor[InMemoryImageStack]:
+    """
+    Create a patch extractor from a sequence of TIFF files.
 
+    Parameters
+    ----------
+    source: sequence of Path
+        The source files for the data.
+    axes: str
+        The original axes of the data, must be a subset of "STCZYX".
+
+    Returns
+    -------
+    PatchExtractor
+    """
+    image_stacks = [
+        InMemoryImageStack.from_tiff(path=path, axes=axes) for path in source
+    ]
+    return PatchExtractor(image_stacks)
 
-# TIFF and ZARR case
-@overload
-def create_patch_extractor(
+
+# ZARR case
+def create_ome_zarr_extractor(
     source: Sequence[Path],
     axes: str,
-    data_type: Literal[SupportedData.TIFF, SupportedDataDev.ZARR],
-) -> PatchExtractor:
+) -> PatchExtractor[ZarrImageStack]:
+    """
+    Create a patch extractor from a sequence of OME ZARR files.
+
+    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.
+    axes: str
+        The original axes of the data, must be a subset of "STCZYX".
+
+    Returns
+    -------
+    PatchExtractor
     """
-    Create a patch extractor from a sequence of files that match our supported types.
+    # NOTE: axes is unused here, in from_ome_zarr the axes are automatically retrieved
+    image_stacks = [ZarrImageStack.from_ome_zarr(path) for path in source]
+    return PatchExtractor(image_stacks)
+
 
-    Supported file types include TIFF and ZARR.
+# CZI case
+def create_czi_extractor(
+    source: Sequence[Path],
+    axes: str,
+) -> PatchExtractor[CziImageStack]:
+    """
+    Create a patch extractor from a sequence of CZI files.
 
-    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).
+    If the CZI files contain multiple scenes, one patch extractor will be created for
+    each scene.
 
     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`.
+    axes: str
+        Specifies which axes of the data to use and how.
+        If this string ends with `"ZYX"` or `"TYX"`, the data will consist of 3-D
+        patches, using `Z` or `T` as third dimension, respectively.
+        If the string does not end with "ZYX", the data will consist of 2-D patches.
 
     Returns
     -------
     PatchExtractor
     """
+    depth_axis: Literal["none", "Z", "T"] = "none"
+    if axes.endswith("TYX"):
+        depth_axis = "T"
+    elif axes.endswith("ZYX"):
+        depth_axis = "Z"
+
+    image_stacks: list[CziImageStack] = []
+    for path in source:
+        scene_rectangles = CziImageStack.get_bounding_rectangles(path)
+        image_stacks.extend(
+            CziImageStack(path, scene=scene, depth_axis=depth_axis)
+            for scene in scene_rectangles.keys()
+        )
+    return PatchExtractor(image_stacks)
 
 
 # Custom file type case (loaded into memory)
-@overload
-def create_patch_extractor(
-    source: Any,
+def create_custom_file_extractor(
+    source: Sequence[Path],
     axes: str,
-    data_type: Literal[SupportedData.CUSTOM],
     *,
     read_func: ReadFunc,
     read_kwargs: dict[str, Any],
-) -> PatchExtractor:
+) -> PatchExtractor[InMemoryImageStack]:
     """
     Create a patch extractor from a sequence of files of a custom type.
 
@@ -89,8 +152,8 @@ def create_patch_extractor(
     ----------
     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".
+    axes: str
+        The original axes of the data, must be a subset of "STCZYX".
     read_func : ReadFunc
         A function to read the custom file type, see the `ReadFunc` protocol.
     read_kwargs : dict of {str: Any}
@@ -100,18 +163,28 @@ def create_patch_extractor(
     -------
     PatchExtractor
     """
+    # TODO: lazy loading custom files
+    image_stacks = [
+        InMemoryImageStack.from_custom_file_type(
+            path=path,
+            axes=axes,
+            read_func=read_func,
+            **read_kwargs,
+        )
+        for path in source
+    ]
+
+    return PatchExtractor(image_stacks)
 
 
 # Custom ImageStackLoader case
-@overload
-def create_patch_extractor(
+def create_custom_image_stack_extractor(
     source: Any,
     axes: str,
-    data_type: Literal[SupportedData.CUSTOM],
-    image_stack_loader: ImageStackLoader[P],
+    image_stack_loader: ImageStackLoader[P, GenericImageStack],
     *args: P.args,
     **kwargs: P.kwargs,
-) -> PatchExtractor:
+) -> PatchExtractor[GenericImageStack]:
     """
     Create a patch extractor using a custom `ImageStackLoader`.
 
@@ -127,8 +200,8 @@ def create_patch_extractor(
     ----------
     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".
+    axes: str
+        The original axes of the data, must be a subset of "STCZYX".
     image_stack_loader: ImageStackLoader
         A custom image stack loader callable.
     *args: Any
@@ -140,69 +213,5 @@ def create_patch_extractor(
     -------
     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)
+    image_stacks = image_stack_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

@@ -4,8 +4,13 @@ __all__ = [
     "PatchingStrategy",
     "RandomPatchingStrategy",
     "SequentialPatchingStrategy",
+    "TileSpecs",
+    "TilingStrategy",
+    "WholeSamplePatchingStrategy",
 ]
 
-from .patching_strategy_protocol import PatchingStrategy, PatchSpecs
+from .patching_strategy_protocol import PatchingStrategy, PatchSpecs, TileSpecs
 from .random_patching import FixedRandomPatchingStrategy, RandomPatchingStrategy
 from .sequential_patching import SequentialPatchingStrategy
+from .tiling_strategy import TilingStrategy
+from .whole_sample import WholeSamplePatchingStrategy

careamics/dataset_ng/patching_strategies/patching_strategy_protocol.py

@@ -28,6 +28,37 @@ class PatchSpecs(TypedDict):
     patch_size: Sequence[int]
 
 
+class TileSpecs(PatchSpecs):
+    """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.
+    crop_coords: sequence of int
+        The top-left side of where the tile will be cropped, in coordinates relative
+        to the tile.
+    crop_size: sequence of int
+        The size of the cropped tile.
+    stitch_coords: sequence of int
+        Where the tile will be stitched back into an image, taking into account
+        that the tile will be cropped, in coords relative to the image.
+    """
+
+    crop_coords: Sequence[int]
+    crop_size: Sequence[int]
+    stitch_coords: Sequence[int]
+
+
 class PatchingStrategy(Protocol):
     """
     An interface for patching strategies.

careamics/dataset_ng/patching_strategies/random_patching.py

@@ -335,4 +335,8 @@ def _calc_n_patches(spatial_shape: Sequence[int], patch_size: Sequence[int]) ->
             f"spatial dimensions {len(spatial_shape)}, for `patch_size={patch_size}` "
             f"and `spatial_shape={spatial_shape}`."
         )
-    return int(np.ceil(np.prod(spatial_shape) / np.prod(patch_size)))
+    patches_per_dim = [
+        np.ceil(s / p) for s, p in zip(spatial_shape, patch_size, strict=False)
+    ]
+    total_patches = int(np.prod(patches_per_dim))
+    return total_patches

careamics/dataset_ng/patching_strategies/sequential_patching.py

@@ -18,13 +18,13 @@ class SequentialPatchingStrategy:
         self,
         data_shapes: Sequence[Sequence[int]],
         patch_size: Sequence[int],
-        overlap: Optional[Sequence[int]] = None,
+        overlaps: Optional[Sequence[int]] = None,
     ):
         self.data_shapes = data_shapes
         self.patch_size = patch_size
-        if overlap is None:
-            overlap = [0] * len(patch_size)
-        self.overlap = np.asarray(overlap)
+        if overlaps is None:
+            overlaps = [0] * len(patch_size)
+        self.overlaps = np.asarray(overlaps)
 
         self.patch_specs: list[PatchSpecs] = self._initialize_patch_specs()
 
@@ -58,7 +58,7 @@ class SequentialPatchingStrategy:
             data_spatial_shape = data_shape[-len(self.patch_size) :]
             coords_list = [
                 self._compute_coords_1d(
-                    self.patch_size[i], data_spatial_shape[i], self.overlap[i]
+                    self.patch_size[i], data_spatial_shape[i], self.overlaps[i]
                 )
                 for i in range(len(self.patch_size))
             ]

careamics/dataset_ng/patching_strategies/tiling_strategy.py

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

careamics/dataset_ng/patching_strategies/whole_sample.py

@@ -0,0 +1,36 @@
+from collections.abc import Sequence
+
+from .patching_strategy_protocol import PatchSpecs
+
+
+class WholeSamplePatchingStrategy:
+    # TODO: warn this strategy should only be used with batch size = 1
+    #   for the case of multiple image stacks with different dimensions
+
+    # TODO: docs
+    def __init__(self, data_shapes: Sequence[Sequence[int]]):
+        self.data_shapes = data_shapes
+
+        self.patch_specs: list[PatchSpecs] = self._initialize_patch_specs()
+
+    @property
+    def n_patches(self) -> int:
+        return len(self.patch_specs)
+
+    def get_patch_spec(self, index: int) -> PatchSpecs:
+        return self.patch_specs[index]
+
+    def _initialize_patch_specs(self) -> list[PatchSpecs]:
+        patch_specs: list[PatchSpecs] = []
+        for data_idx, data_shape in enumerate(self.data_shapes):
+            spatial_shape = data_shape[2:]
+            for sample_idx in range(data_shape[0]):
+                patch_specs.append(
+                    {
+                        "data_idx": data_idx,
+                        "sample_idx": sample_idx,
+                        "coords": tuple(0 for _ in spatial_shape),
+                        "patch_size": spatial_shape,
+                    }
+                )
+        return patch_specs

careamics/file_io/read/get_func.py

@@ -1,7 +1,8 @@
 """Module to get read functions."""
 
+from collections.abc import Callable
 from pathlib import Path
-from typing import Callable, Protocol, Union
+from typing import Protocol, Union
 
 from numpy.typing import NDArray
 

careamics/lightning/dataset_ng/__init__.py

@@ -0,0 +1 @@
+"""Next-Generation DataModules for Careamics."""