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

careamics/dataset_ng/legacy_interoperability.py

@@ -0,0 +1,168 @@
+"""
+A module for utility functions that adapts the new dataset outputs to work with previous
+code until it is updated.
+"""
+
+from collections.abc import Sequence
+from typing import cast
+
+import numpy as np
+from numpy.typing import NDArray
+
+from careamics.config.tile_information import TileInformation
+
+from .dataset import ImageRegionData
+from .patching_strategies import TileSpecs
+
+
+def imageregions_to_tileinfos(
+    image_regions: Sequence[ImageRegionData],
+) -> list[tuple[NDArray, list[TileInformation]]]:
+    """
+    Converts a series of `TileSpecs` dictionaries to `TileInformation` pydantic class.
+
+    Parameters
+    ----------
+    image_regions : sequence of ImageRegionData
+        A list of ImageRegionData, it must have an instance of `TileSpecs` as it's
+        `region_data` field.
+
+    Returns
+    -------
+    list of TileInformation
+        The converted tile information.
+    """
+
+    tile_infos: list[TileInformation] = []
+
+    data = [image_region.data for image_region in image_regions]
+    tile_specs = [image_region.region_spec for image_region in image_regions]
+
+    data_indices: NDArray[np.int_] = np.array(
+        [tile_spec["data_idx"] for tile_spec in tile_specs], dtype=int
+    )
+    unique_data_indices = np.unique(data_indices)
+    # data_idx denotes which image stack a patch belongs to
+    # separate TileSpecs by image_stack
+    for data_idx in unique_data_indices:
+        # collect all ImageRegions
+        data_image_regions: list[ImageRegionData] = [
+            image_region
+            for image_region in image_regions
+            if image_region.region_spec["data_idx"] == data_idx
+        ]
+
+        # --- find last indices
+        # make sure tiles belonging to the same sample are together
+        data_image_regions.sort(
+            key=lambda image_region: image_region.region_spec["sample_idx"]
+        )
+        sample_indices = np.array(
+            [
+                image_region.region_spec["sample_idx"]
+                for image_region in data_image_regions
+            ]
+        )
+        # reverse array so indices returned are at far edge
+        _, unique_indices = np.unique(sample_indices[::-1], return_index=True)
+        # un reverse indices
+        last_indices = len(sample_indices) - 1 - unique_indices
+
+        # convert each ImageRegionData to tile_info
+        for i, image_region in enumerate(data_image_regions):
+            last_tile = i in last_indices
+            tile_info = _imageregion_to_tileinfo(image_region, last_tile)
+            tile_infos.append(tile_info)
+
+    return [(data, [tile_info]) for data, tile_info in zip(data, tile_infos)]
+
+
+def _imageregion_to_tileinfo(
+    image_region: ImageRegionData, last_tile: bool
+) -> TileInformation:
+    """
+    Convert a single `ImageRegionData` instance to a `TileInformation` instance. Whether
+    it is the last tile in a sequence needs to be supplied.
+
+    Parameters
+    ----------
+    image_region : ImageRegionData
+        An instance of `ImageRegionData`, it must have an instance of `TileSpecs` as
+        it's `region_data` field.
+    last_tile : bool
+        Whether a tile is the last tile in a sequence, for stitching.
+
+    Returns
+    -------
+    TileInformation
+        A tile information object.
+
+    Raises
+    ------
+    KeyError
+        If `image_region.region_spec` does not contain the keys: {'crop_coords',
+        'crop_size', 'stitch_coords'}.
+    """
+    patch_spec = image_region.region_spec
+    data_shape = image_region.data_shape
+
+    # TODO: In python 3.11 and greater, NamedTuples can inherit from Generic
+    #   so we could do image_region: ImageRegionData[TileSpecs]
+    #   and not have to do this check here + cast
+    # make sure image_region.region_spec is TileSpec
+    if (
+        ("crop_coords" not in patch_spec)
+        or ("crop_size" not in patch_spec)
+        or ("stitch_coords" not in patch_spec)
+    ):
+        raise KeyError(
+            "Could not find all keys: {'crop_coords', 'crop_size', 'stitch_coords'} in "
+            "`image_region.region_spec`."
+        )
+    tile_spec = cast(TileSpecs, patch_spec)  # ugly cast for mypy
+    return _tilespec_to_tileinfo(tile_spec, data_shape, last_tile)
+
+
+def _tilespec_to_tileinfo(
+    tile_spec: TileSpecs, data_shape: Sequence[int], last_tile: bool
+) -> TileInformation:
+    """
+    Convert a single `TileSpec` to a `TileInformation`. Whether it is the last tile
+    needs to be supplied.
+
+    Parameters
+    ----------
+    tile_spec : TileSpecs
+        A tile spec dictionary.
+    data_shape : sequence of int
+        The original shape of the data the tile came from, labeling the dimensions of
+        axes SC(Z)YX.
+    last_tile : bool
+        Whether a tile is the last tile in a sequence, for stitching.
+
+    Returns
+    -------
+    TileInformation
+        A tile information object.
+    """
+    overlap_crop_coords = tuple(
+        (
+            tile_spec["crop_coords"][i],
+            tile_spec["crop_coords"][i] + tile_spec["crop_size"][i],
+        )
+        for i in range(len(tile_spec["crop_coords"]))
+    )
+    stitch_coords = tuple(
+        (
+            tile_spec["stitch_coords"][i],
+            tile_spec["stitch_coords"][i] + tile_spec["crop_size"][i],
+        )
+        for i in range(len(tile_spec["crop_coords"]))
+    )
+    return TileInformation(
+        array_shape=tuple(data_shape[1:]),  # remove sample dimension
+        last_tile=last_tile,
+        overlap_crop_coords=overlap_crop_coords,
+        stitch_coords=stitch_coords,
+        sample_id=tile_spec["sample_idx"],
+    )

careamics/dataset_ng/patch_extractor/__init__.py

@@ -1,10 +1,5 @@
-__all__ = [
-    "ImageStackLoader",
-    "PatchExtractor",
-    "create_patch_extractor",
-    "get_image_stack_loader",
-]
+__all__ = ["GenericImageStack", "ImageStackLoader", "PatchExtractor"]
 
-from .image_stack_loader import ImageStackLoader, get_image_stack_loader
+from .image_stack import GenericImageStack
+from .image_stack_loader import ImageStackLoader
 from .patch_extractor import PatchExtractor
-from .patch_extractor_factory import create_patch_extractor

careamics/dataset_ng/patch_extractor/demo_custom_image_stack_loader.py

@@ -11,9 +11,11 @@ from zarr.storage import FSStore
 
 from careamics.config import DataConfig
 from careamics.config.support import SupportedData
-from careamics.dataset_ng.patch_extractor import create_patch_extractor
 from careamics.dataset_ng.patch_extractor.image_stack import ZarrImageStack
 from careamics.dataset_ng.patch_extractor.image_stack_loader import ImageStackLoader
+from careamics.dataset_ng.patch_extractor.patch_extractor_factory import (
+    create_custom_image_stack_extractor,
+)
 
 
 # %%
@@ -94,12 +96,12 @@ image_stack_loader: ImageStackLoader = custom_image_stack_loader
 
 # %%
 # So pylance knows that datatype is custom to match function overloads
-assert data_config.data_type is SupportedData.CUSTOM
+assert SupportedData(data_config.data_type) is SupportedData.CUSTOM
 
-patch_extractor = create_patch_extractor(
+patch_extractor = create_custom_image_stack_extractor(
     source={"store": store, "data_paths": data_paths},
     axes=data_config.axes,
-    data_type=data_config.data_type,
+    data_type=SupportedData(data_config.data_type),
     image_stack_loader=custom_image_stack_loader,
 )
 

careamics/dataset_ng/patch_extractor/image_stack/__init__.py

@@ -1,9 +1,10 @@
 __all__ = [
+    "GenericImageStack",
     "ImageStack",
     "InMemoryImageStack",
     "ZarrImageStack",
 ]
 
-from .image_stack_protocol import ImageStack
+from .image_stack_protocol import GenericImageStack, ImageStack
 from .in_memory_image_stack import InMemoryImageStack
 from .zarr_image_stack import ZarrImageStack

careamics/dataset_ng/patch_extractor/image_stack/image_stack_protocol.py

@@ -1,6 +1,6 @@
 from collections.abc import Sequence
 from pathlib import Path
-from typing import Literal, Protocol, Union
+from typing import Literal, Protocol, TypeVar, Union
 
 from numpy.typing import DTypeLike, NDArray
 
@@ -51,3 +51,7 @@ class ImageStack(Protocol):
             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/patch_extractor/image_stack_loader.py

@@ -1,20 +1,11 @@
 from collections.abc import Sequence
-from pathlib import Path
-from typing import (
-    Any,
-    Optional,
-    Protocol,
-    Union,
-)
-
-from numpy.typing import NDArray
+from typing import Any, Protocol
+
 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
+from .image_stack import GenericImageStack
 
 P = ParamSpec("P")
 
@@ -23,7 +14,7 @@ class SupportedDataDev(str, BaseEnum):
     ZARR = "zarr"
 
 
-class ImageStackLoader(Protocol[P]):
+class ImageStackLoader(Protocol[P, GenericImageStack]):
     """
     Protocol to define how `ImageStacks` should be loaded.
 
@@ -76,65 +67,4 @@ class ImageStackLoader(Protocol[P]):
 
     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
+    ) -> Sequence[GenericImageStack]: ...

careamics/dataset_ng/patch_extractor/patch_extractor.py

@@ -1,17 +1,18 @@
 from collections.abc import Sequence
+from typing import Generic
 
 from numpy.typing import NDArray
 
-from .image_stack import ImageStack
+from .image_stack import GenericImageStack
 
 
-class PatchExtractor:
+class PatchExtractor(Generic[GenericImageStack]):
     """
     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 __init__(self, image_stacks: Sequence[GenericImageStack]):
+        self.image_stacks: list[GenericImageStack] = list(image_stacks)
 
     def extract_patch(
         self,

careamics/dataset_ng/patch_extractor/patch_extractor_factory.py

@@ -1,29 +1,29 @@
 from collections.abc import Sequence
 from pathlib import Path
-from typing import Any, Literal, Optional, Union, overload
+from typing import Any
 
 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 (
+    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 +31,78 @@ 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.
 
-# TIFF and ZARR case
-@overload
-def create_patch_extractor(
+    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)
+
+
+# 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 files that match our supported types.
+    Create a patch extractor from a sequence of OME ZARR files.
 
-    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).
+    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`.
+    axes: str
+        The original axes of the data, must be a subset of "STCZYX".
 
     Returns
     -------
     PatchExtractor
     """
+    # 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)
 
 
 # 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 +110,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 +121,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 +158,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 +171,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.