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

careamics/dataset_ng/factory.py

@@ -0,0 +1,451 @@
+from collections.abc import Sequence
+from enum import Enum
+from pathlib import Path
+from typing import Any, Optional
+
+from numpy.typing import NDArray
+from typing_extensions import ParamSpec
+
+from careamics.config.data.ng_data_model import NGDataConfig
+from careamics.config.support import SupportedData
+from careamics.dataset_ng.patch_extractor import ImageStackLoader, PatchExtractor
+from careamics.dataset_ng.patch_extractor.image_stack import (
+    CziImageStack,
+    GenericImageStack,
+    ImageStack,
+    InMemoryImageStack,
+    ZarrImageStack,
+)
+from careamics.dataset_ng.patch_extractor.patch_extractor_factory import (
+    create_array_extractor,
+    create_custom_file_extractor,
+    create_custom_image_stack_extractor,
+    create_czi_extractor,
+    create_ome_zarr_extractor,
+    create_tiff_extractor,
+)
+from careamics.file_io.read import ReadFunc
+
+from .dataset import CareamicsDataset, Mode
+
+P = ParamSpec("P")
+
+
+# Enum class used to determine which loading functions should be used
+class DatasetType(Enum):
+    """Labels for the dataset based on the underlying data and how it is loaded."""
+
+    ARRAY = "array"
+    IN_MEM_TIFF = "in_mem_tiff"
+    LAZY_TIFF = "lazy_tiff"
+    IN_MEM_CUSTOM_FILE = "in_mem_custom_file"
+    OME_ZARR = "ome_zarr"
+    CZI = "czi"
+    CUSTOM_IMAGE_STACK = "custom_image_stack"
+
+
+# bit of a mess of if-else statements
+def determine_dataset_type(
+    data_type: SupportedData,
+    in_memory: bool,
+    read_func: Optional[ReadFunc] = None,
+    image_stack_loader: Optional[ImageStackLoader] = None,
+) -> DatasetType:
+    """Determine what the dataset type should be based on the input arguments.
+
+    Parameters
+    ----------
+    data_type : SupportedData
+        The underlying datatype.
+    in_memory : bool
+        Whether all the data should be loaded into memory. This is argument is ignored
+        unless the `data_type` is "tiff" or "custom".
+    read_func : ReadFunc, optional
+        A function that can be used to load custom data. This argument is
+        ignored unless the `data_type` is "custom".
+    image_stack_loader : ImageStackLoader, optional
+        A function for custom image stack loading. This argument is ignored unless the
+        `data_type` is "custom".
+
+    Returns
+    -------
+    DatasetType
+        The Dataset type.
+
+    Raises
+    ------
+    NotImplementedError
+        For lazy-loading (`in_memory=False`) of a custom file type.
+    ValueError
+        If the `data_type` is "custom" but both `read_func` and `image_stack_loader` are
+        None.
+    ValueError
+        If the `data_type` is unrecognized.
+    """
+    if data_type == SupportedData.ARRAY:
+        # TODO: ignoring in_memory arg, error if False?
+        return DatasetType.ARRAY
+    elif data_type == SupportedData.TIFF:
+        if in_memory:
+            return DatasetType.IN_MEM_TIFF
+        else:
+            return DatasetType.LAZY_TIFF
+    elif data_type == SupportedData.CZI:
+        return DatasetType.CZI
+    elif data_type == SupportedData.CUSTOM:
+        if read_func is not None:
+            if in_memory:
+                return DatasetType.IN_MEM_CUSTOM_FILE
+            else:
+                raise NotImplementedError(
+                    "Lazy loading has not been implemented for custom file types yet."
+                )
+        elif image_stack_loader is not None:
+            # TODO: ignoring im_memory arg
+            return DatasetType.CUSTOM_IMAGE_STACK
+        else:
+            raise ValueError(
+                "Found `data_type='custom'` but no `read_func` or `image_stack_loader` "
+                "has been provided."
+            )
+    # TODO: ZARR
+    else:
+        raise ValueError(f"Unrecognized `data_type`, '{data_type}'.")
+
+
+# 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,
+    mode: Mode,
+    inputs: Any,
+    targets: Any,
+    in_memory: bool,
+    read_func: Optional[ReadFunc] = None,
+    read_kwargs: Optional[dict[str, Any]] = None,
+    image_stack_loader: Optional[ImageStackLoader] = None,
+    image_stack_loader_kwargs: Optional[dict[str, Any]] = None,
+) -> CareamicsDataset[ImageStack]:
+    """
+    Convenience function to create the CAREamicsDataset.
+
+    Parameters
+    ----------
+    config : DataConfig or InferenceConfig
+        The data configuration.
+    mode : Mode
+        Whether to create the dataset in "training", "validation" or "predicting" mode.
+    inputs : Any
+        The input sources to the dataset.
+    targets : Any, optional
+        The target sources to the dataset.
+    in_memory : bool
+        Whether all the data should be loaded into memory. This is argument is ignored
+        unless the `data_type` in `config` is "tiff" or "custom".
+    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`.
+
+    Returns
+    -------
+    CareamicsDataset[ImageStack]
+        The CAREamicsDataset.
+
+    Raises
+    ------
+    ValueError
+        For an unrecognized `data_type` in the `config`.
+    """
+    data_type = SupportedData(config.data_type)
+    dataset_type = determine_dataset_type(
+        data_type, in_memory, read_func, image_stack_loader
+    )
+    if dataset_type == DatasetType.ARRAY:
+        return create_array_dataset(config, mode, inputs, targets)
+    elif dataset_type == DatasetType.IN_MEM_TIFF:
+        return create_tiff_dataset(config, mode, inputs, targets)
+    # TODO: Lazy tiff
+    elif dataset_type == DatasetType.CZI:
+        return create_czi_dataset(config, mode, inputs, targets)
+    elif dataset_type == DatasetType.IN_MEM_CUSTOM_FILE:
+        if read_kwargs is None:
+            read_kwargs = {}
+        assert read_func is not None  # should be true from `determine_dataset_type`
+        return create_custom_file_dataset(
+            config, mode, inputs, targets, read_func=read_func, read_kwargs=read_kwargs
+        )
+    elif dataset_type == DatasetType.CUSTOM_IMAGE_STACK:
+        if image_stack_loader_kwargs is None:
+            image_stack_loader_kwargs = {}
+        assert image_stack_loader is not None  # should be true
+        return create_custom_image_stack_dataset(
+            config,
+            mode,
+            inputs,
+            targets,
+            image_stack_loader,
+            **image_stack_loader_kwargs,
+        )
+    else:
+        raise ValueError(f"Unrecognized dataset type, {dataset_type}.")
+
+
+def create_array_dataset(
+    config: NGDataConfig,
+    mode: Mode,
+    inputs: Sequence[NDArray[Any]],
+    targets: Optional[Sequence[NDArray[Any]]],
+) -> CareamicsDataset[InMemoryImageStack]:
+    """
+    Create a CAREamicsDataset from array data.
+
+    Parameters
+    ----------
+    config : DataConfig or InferenceConfig
+        The data configuration.
+    mode : Mode
+        Whether to create the dataset in "training", "validation" or "predicting" mode.
+    inputs : Any
+        The input sources to the dataset.
+    targets : Any, optional
+        The target sources to the dataset.
+
+    Returns
+    -------
+    CareamicsDataset[InMemoryImageStack]
+        A CAREamicsDataset.
+    """
+    input_extractor = create_array_extractor(source=inputs, axes=config.axes)
+    target_extractor: Optional[PatchExtractor[InMemoryImageStack]]
+    if targets is not None:
+        target_extractor = create_array_extractor(source=targets, axes=config.axes)
+    else:
+        target_extractor = None
+    return CareamicsDataset(config, mode, input_extractor, target_extractor)
+
+
+def create_tiff_dataset(
+    config: NGDataConfig,
+    mode: Mode,
+    inputs: Sequence[Path],
+    targets: Optional[Sequence[Path]],
+) -> CareamicsDataset[InMemoryImageStack]:
+    """
+    Create a CAREamicsDataset from tiff files that will be all loaded into memory.
+
+    Parameters
+    ----------
+    config : DataConfig or InferenceConfig
+        The data configuration.
+    mode : Mode
+        Whether to create the dataset in "training", "validation" or "predicting" mode.
+    inputs : Any
+        The input sources to the dataset.
+    targets : Any, optional
+        The target sources to the dataset.
+
+    Returns
+    -------
+    CareamicsDataset[InMemoryImageStack]
+        A CAREamicsDataset.
+    """
+    input_extractor = create_tiff_extractor(
+        source=inputs,
+        axes=config.axes,
+    )
+    target_extractor: Optional[PatchExtractor[InMemoryImageStack]]
+    if targets is not None:
+        target_extractor = create_tiff_extractor(source=targets, axes=config.axes)
+    else:
+        target_extractor = None
+    dataset = CareamicsDataset(config, mode, input_extractor, target_extractor)
+    return dataset
+
+
+def create_czi_dataset(
+    config: NGDataConfig,
+    mode: Mode,
+    inputs: Sequence[Path],
+    targets: Optional[Sequence[Path]],
+) -> CareamicsDataset[CziImageStack]:
+    """
+    Create a dataset from CZI files.
+
+    Parameters
+    ----------
+    config : DataConfig or InferenceConfig
+        The data configuration.
+    mode : Mode
+        Whether to create the dataset in "training", "validation" or "predicting" mode.
+    inputs : Any
+        The input sources to the dataset.
+    targets : Any, optional
+        The target sources to the dataset.
+
+    Returns
+    -------
+    CareamicsDataset[CziImageStack]
+        A CAREamicsDataset.
+    """
+
+    input_extractor = create_czi_extractor(source=inputs, axes=config.axes)
+    target_extractor: Optional[PatchExtractor[CziImageStack]]
+    if targets is not None:
+        target_extractor = create_czi_extractor(source=targets, axes=config.axes)
+    else:
+        target_extractor = None
+    dataset = CareamicsDataset(config, mode, input_extractor, target_extractor)
+    return dataset
+
+
+def create_ome_zarr_dataset(
+    config: NGDataConfig,
+    mode: Mode,
+    inputs: Sequence[Path],
+    targets: Optional[Sequence[Path]],
+) -> CareamicsDataset[ZarrImageStack]:
+    """
+    Create a dataset from OME ZARR files.
+
+    Parameters
+    ----------
+    config : DataConfig or InferenceConfig
+        The data configuration.
+    mode : Mode
+        Whether to create the dataset in "training", "validation" or "predicting" mode.
+    inputs : Any
+        The input sources to the dataset.
+    targets : Any, optional
+        The target sources to the dataset.
+
+    Returns
+    -------
+    CareamicsDataset[ZarrImageStack]
+        A CAREamicsDataset.
+    """
+
+    input_extractor = create_ome_zarr_extractor(source=inputs, axes=config.axes)
+    target_extractor: Optional[PatchExtractor[ZarrImageStack]]
+    if targets is not None:
+        target_extractor = create_ome_zarr_extractor(source=targets, axes=config.axes)
+    else:
+        target_extractor = None
+    dataset = CareamicsDataset(config, mode, input_extractor, target_extractor)
+    return dataset
+
+
+def create_custom_file_dataset(
+    config: NGDataConfig,
+    mode: Mode,
+    inputs: Sequence[Path],
+    targets: Optional[Sequence[Path]],
+    *,
+    read_func: ReadFunc,
+    read_kwargs: dict[str, Any],
+) -> CareamicsDataset[InMemoryImageStack]:
+    """
+    Create a CAREamicsDataset from custom files that will be all loaded into memory.
+
+    Parameters
+    ----------
+    config : DataConfig or InferenceConfig
+        The data configuration.
+    mode : Mode
+        Whether to create the dataset in "training", "validation" or "predicting" mode.
+    inputs : Any
+        The input sources to the dataset.
+    targets : Any, optional
+        The target sources to the dataset.
+    read_func : Optional[ReadFunc], optional
+        A function that can that can be used to load custom data. This argument is
+        ignored unless the `data_type` is "custom".
+    image_stack_loader : Optional[ImageStackLoader], optional
+        A function for custom image stack loading. This argument is ignored unless the
+        `data_type` is "custom".
+
+    Returns
+    -------
+    CareamicsDataset[InMemoryImageStack]
+        A CAREamicsDataset.
+    """
+    input_extractor = create_custom_file_extractor(
+        source=inputs, axes=config.axes, read_func=read_func, read_kwargs=read_kwargs
+    )
+    target_extractor: Optional[PatchExtractor[InMemoryImageStack]]
+    if targets is not None:
+        target_extractor = create_custom_file_extractor(
+            source=targets,
+            axes=config.axes,
+            read_func=read_func,
+            read_kwargs=read_kwargs,
+        )
+    else:
+        target_extractor = None
+    dataset = CareamicsDataset(config, mode, input_extractor, target_extractor)
+    return dataset
+
+
+def create_custom_image_stack_dataset(
+    config: NGDataConfig,
+    mode: Mode,
+    inputs: Any,
+    targets: Optional[Any],
+    image_stack_loader: ImageStackLoader[P, GenericImageStack],
+    *args: P.args,
+    **kwargs: P.kwargs,
+) -> CareamicsDataset[GenericImageStack]:
+    """
+    Create a CAREamicsDataset from a custom `ImageStack` class.
+
+    The custom `ImageStack` class can be loaded using the `image_stack_loader` function.
+
+    Parameters
+    ----------
+    config : DataConfig or InferenceConfig
+        The data configuration.
+    mode : Mode
+        Whether to create the dataset in "training", "validation" or "predicting" mode.
+    inputs : Any
+        The input sources to the dataset.
+    targets : Any, optional
+        The target sources to the dataset.
+    image_stack_loader : ImageStackLoader
+        A function for custom image stack loading. This argument is ignored unless the
+        `data_type` is "custom".
+    *args : Any
+        Positional arguments to pass to the `image_stack_loader`.
+    **kwargs : Any
+        Key-word arguments to pass to the `image_stack_loader`.
+
+    Returns
+    -------
+    CareamicsDataset[GenericImageStack]
+        A CAREamicsDataset.
+    """
+    input_extractor = create_custom_image_stack_extractor(
+        inputs,
+        config.axes,
+        image_stack_loader,
+        *args,
+        **kwargs,
+    )
+    target_extractor: Optional[PatchExtractor[GenericImageStack]]
+    if targets is not None:
+        target_extractor = create_custom_image_stack_extractor(
+            targets,
+            config.axes,
+            image_stack_loader,
+            *args,
+            **kwargs,
+        )
+    else:
+        target_extractor = None
+    dataset = CareamicsDataset(config, mode, input_extractor, target_extractor)
+    return dataset

careamics/dataset_ng/legacy_interoperability.py

@@ -0,0 +1,170 @@
+"""
+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, strict=False)
+    ]
+
+
+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,
+)
 
 
 # %%
@@ -33,7 +35,7 @@ def create_zarr_array(file_path: Path, data_path: str, data: NDArray):
 
 
 def create_zarr(file_path: Path, data_paths: Sequence[str], data: Sequence[NDArray]):
-    for data_path, array in zip(data_paths, data):
+    for data_path, array in zip(data_paths, data, strict=False):
         create_zarr_array(file_path=file_path, data_path=data_path, data=array)
 
 
@@ -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,12 @@
 __all__ = [
+    "CziImageStack",
+    "GenericImageStack",
     "ImageStack",
     "InMemoryImageStack",
     "ZarrImageStack",
 ]
 
-from .image_stack_protocol import ImageStack
+from .czi_image_stack import CziImageStack
+from .image_stack_protocol import GenericImageStack, ImageStack
 from .in_memory_image_stack import InMemoryImageStack
 from .zarr_image_stack import ZarrImageStack