ngio 0.4.0a3__py3-none-any.whl → 0.4.0b1__py3-none-any.whl

ngio/images/_ome_zarr_container.py

@@ -4,6 +4,7 @@ import warnings
 from collections.abc import Sequence
 
 import numpy as np
+from zarr.types import DIMENSION_SEPARATOR
 
 from ngio.images._create import create_empty_image_container
 from ngio.images._image import Image, ImagesContainer
@@ -408,6 +409,8 @@ class OmeZarrContainer:
         name: str | None = None,
         chunks: Sequence[int] | None = None,
         dtype: str | None = None,
+        dimension_separator: DIMENSION_SEPARATOR | None = None,
+        compressor=None,
         copy_labels: bool = False,
         copy_tables: bool = False,
         overwrite: bool = False,
@@ -425,6 +428,11 @@ class OmeZarrContainer:
             chunks (Sequence[int] | None): The chunk shape of the new image.
             dtype (str | None): The data type of the new image.
             name (str | None): The name of the new image.
+            dimension_separator (DIMENSION_SEPARATOR | None): The dimension
+                separator to use. If None, the dimension separator of the
+                reference image will be used.
+            compressor: The compressor to use. If None, the compressor of the
+                reference image will be used.
             copy_labels (bool): Whether to copy the labels from the reference image.
             copy_tables (bool): Whether to copy the tables from the reference image.
             overwrite (bool): Whether to overwrite an existing image.
@@ -443,6 +451,8 @@ class OmeZarrContainer:
             name=name,
             chunks=chunks,
             dtype=dtype,
+            dimension_separator=dimension_separator,
+            compressor=compressor,
             overwrite=overwrite,
         )
 
@@ -683,6 +693,8 @@ class OmeZarrContainer:
         axes_names: Sequence[str] | None = None,
         chunks: Sequence[int] | None = None,
         dtype: str = "uint32",
+        dimension_separator: DIMENSION_SEPARATOR | None = None,
+        compressor=None,
         overwrite: bool = False,
     ) -> "Label":
         """Create an empty OME-Zarr label from a reference image.
@@ -699,6 +711,11 @@ class OmeZarrContainer:
                 For labels, the channel axis is not allowed.
             chunks (Sequence[int] | None): The chunk shape of the new image.
             dtype (str): The data type of the new label.
+            dimension_separator (DIMENSION_SEPARATOR | None): The dimension
+                separator to use. If None, the dimension separator of the
+                reference image will be used.
+            compressor: The compressor to use. If None, the compressor of the
+                reference image will be used.
             overwrite (bool): Whether to overwrite an existing image.
 
         Returns:
@@ -715,6 +732,8 @@ class OmeZarrContainer:
             axes_names=axes_names,
             chunks=chunks,
             dtype=dtype,
+            dimension_separator=dimension_separator,
+            compressor=compressor,
             overwrite=overwrite,
         )
 
@@ -819,6 +838,8 @@ def create_empty_ome_zarr(
     name: str | None = None,
     chunks: Sequence[int] | None = None,
     dtype: str = "uint16",
+    dimension_separator: DIMENSION_SEPARATOR = "/",
+    compressor="default",
     channel_labels: list[str] | None = None,
     channel_wavelengths: list[str] | None = None,
     channel_colors: Sequence[str] | None = None,
@@ -851,6 +872,9 @@ def create_empty_ome_zarr(
         chunks (Sequence[int] | None, optional): The chunk shape. If None the shape
             is used. Defaults to None.
         dtype (str, optional): The data type of the image. Defaults to "uint16".
+        dimension_separator (DIMENSION_SEPARATOR): The dimension
+            separator to use. Defaults to "/".
+        compressor: The compressor to use. Defaults to "default".
         channel_labels (list[str] | None, optional): The labels of the channels.
             Defaults to None.
         channel_wavelengths (list[str] | None, optional): The wavelengths of the
@@ -879,6 +903,8 @@ def create_empty_ome_zarr(
         name=name,
         chunks=chunks,
         dtype=dtype,
+        dimension_separator=dimension_separator,
+        compressor=compressor,
         overwrite=overwrite,
         version=version,
     )
@@ -913,6 +939,8 @@ def create_ome_zarr_from_array(
     channel_active: Sequence[bool] | None = None,
     name: str | None = None,
     chunks: Sequence[int] | None = None,
+    dimension_separator: DIMENSION_SEPARATOR = "/",
+    compressor: str = "default",
     overwrite: bool = False,
     version: NgffVersions = DefaultNgffVersion,
 ) -> OmeZarrContainer:
@@ -950,6 +978,9 @@ def create_ome_zarr_from_array(
             Defaults to None.
         channel_active (Sequence[bool] | None, optional): Whether the channels are
             active. Defaults to None.
+        dimension_separator (DIMENSION_SEPARATOR): The separator to use for
+            dimensions. Defaults to "/".
+        compressor: The compressor to use. Defaults to "default".
         overwrite (bool, optional): Whether to overwrite an existing image.
             Defaults to True.
         version (str, optional): The version of the OME-Zarr specification.
@@ -971,6 +1002,8 @@ def create_ome_zarr_from_array(
         chunks=chunks,
         dtype=str(array.dtype),
         overwrite=overwrite,
+        dimension_separator=dimension_separator,
+        compressor=compressor,
         version=version,
     )
 

ngio/io_pipes/__init__.py

@@ -0,0 +1,75 @@
+"""I/O pipes for reading and writing data from zarr to numpy and dask arrays.
+
+There are 3 main types of I/O pipes:
+- Standard I/O pipes: NumpyGetter, NumpySetter, DaskGetter, DaskSetter:
+    These pipes read and write data from simple integer indexing and slicing.
+- ROI I/O pipes: NumpyRoiGetter, NumpyRoiSetter, DaskRoiGetter, DaskRoiSetter:
+    These pipes read and write data from a region of interest (ROI) defined in physical
+    coordinates.
+- Masked I/O pipes: NumpyGetterMasked, NumpySetterMasked, DaskGetterMasked,
+    DaskSetterMasked: These pipes like the ROI pipes read and write data
+    from a region of interest (ROI). However they also load a boolean mask
+    from a label zarr array to mask the data being read or written.
+
+All the io pipes are structured in the same way.
+
+When reading data the order of operations is:
+- Step 1: Slice the zarr array to load only the data needed into memory.
+- Step 2: Apply axes operations to reorder, squeeze or expand the axes.
+    To match the user desired axes order.
+- Step 3: Apply any additional transforms to the data.
+
+When writing data the order of operations is the reverse.
+
+The Transforms must implement the TransformProtocol.
+They should be stateless and only depend on the input array and the slicing
+and axes ops. This allows them to be easily reused between different I/O pipes.
+
+"""
+
+from ngio.io_pipes._io_pipes import (
+    DaskGetter,
+    DaskSetter,
+    DataGetter,
+    DataSetter,
+    NumpyGetter,
+    NumpySetter,
+)
+from ngio.io_pipes._io_pipes_masked import (
+    DaskGetterMasked,
+    DaskSetterMasked,
+    NumpyGetterMasked,
+    NumpySetterMasked,
+)
+from ngio.io_pipes._io_pipes_roi import (
+    DaskRoiGetter,
+    DaskRoiSetter,
+    NumpyRoiGetter,
+    NumpyRoiSetter,
+)
+from ngio.io_pipes._match_shape import dask_match_shape, numpy_match_shape
+from ngio.io_pipes._ops_slices import SlicingInputType, SlicingOps, SlicingType
+from ngio.io_pipes._ops_transforms import TransformProtocol
+
+__all__ = [
+    "DaskGetter",
+    "DaskGetterMasked",
+    "DaskRoiGetter",
+    "DaskRoiSetter",
+    "DaskSetter",
+    "DaskSetterMasked",
+    "DataGetter",
+    "DataSetter",
+    "NumpyGetter",
+    "NumpyGetterMasked",
+    "NumpyRoiGetter",
+    "NumpyRoiSetter",
+    "NumpySetter",
+    "NumpySetterMasked",
+    "SlicingInputType",
+    "SlicingOps",
+    "SlicingType",
+    "TransformProtocol",
+    "dask_match_shape",
+    "numpy_match_shape",
+]

ngio/io_pipes/_io_pipes.py

@@ -0,0 +1,361 @@
+from abc import ABC, abstractmethod
+from collections.abc import Sequence
+from typing import Generic, TypeVar
+
+import numpy as np
+import zarr
+from dask.array import Array as DaskArray
+
+from ngio.common._dimensions import Dimensions
+from ngio.common._roi import Roi, RoiPixels
+from ngio.io_pipes._ops_axes import (
+    AxesOps,
+    build_axes_ops,
+    get_as_dask_axes_ops,
+    get_as_numpy_axes_ops,
+    set_as_dask_axes_ops,
+    set_as_numpy_axes_ops,
+)
+from ngio.io_pipes._ops_slices import (
+    SlicingInputType,
+    SlicingOps,
+    build_slicing_ops,
+    get_slice_as_dask,
+    get_slice_as_numpy,
+    set_slice_as_dask,
+    set_slice_as_numpy,
+)
+from ngio.io_pipes._ops_transforms import (
+    TransformProtocol,
+    get_as_dask_transform,
+    get_as_numpy_transform,
+    set_as_dask_transform,
+    set_as_numpy_transform,
+)
+
+
+def setup_io_pipe(
+    *,
+    dimensions: Dimensions,
+    slicing_dict: dict[str, SlicingInputType] | None = None,
+    axes_order: Sequence[str] | None = None,
+    remove_channel_selection: bool = False,
+) -> tuple[SlicingOps, AxesOps]:
+    """Setup the slicing tuple and axes ops for an IO pipe."""
+    slicing_ops = build_slicing_ops(
+        dimensions=dimensions,
+        slicing_dict=slicing_dict,
+        remove_channel_selection=remove_channel_selection,
+    )
+
+    axes_ops = build_axes_ops(
+        dimensions=dimensions,
+        input_axes=slicing_ops.slice_axes,
+        axes_order=axes_order,
+    )
+    return slicing_ops, axes_ops
+
+
+##############################################################
+#
+# "From Disk" Pipes
+#
+##############################################################
+
+T = TypeVar("T")
+
+
+class DataGetter(ABC, Generic[T]):
+    def __init__(
+        self,
+        zarr_array: zarr.Array,
+        slicing_ops: SlicingOps,
+        axes_ops: AxesOps,
+        transforms: Sequence[TransformProtocol] | None = None,
+        roi: Roi | RoiPixels | None = None,
+    ) -> None:
+        self._zarr_array = zarr_array
+        self._slicing_ops = slicing_ops
+        self._axes_ops = axes_ops
+        self._transforms = transforms
+        self._roi = roi
+
+    def __repr__(self) -> str:
+        name = self.__class__.__name__
+        return (
+            f"{name}(zarr_array={self._zarr_array}, "
+            f"slicing_ops={self._slicing_ops}, "
+            f"axes_ops={self._axes_ops}, "
+            f"transforms={self._transforms})"
+        )
+
+    @property
+    def zarr_array(self) -> zarr.Array:
+        return self._zarr_array
+
+    @property
+    def slicing_ops(self) -> SlicingOps:
+        return self._slicing_ops
+
+    @property
+    def axes_ops(self) -> AxesOps:
+        return self._axes_ops
+
+    @property
+    def transforms(self) -> Sequence[TransformProtocol] | None:
+        return self._transforms
+
+    @property
+    def roi(self) -> Roi | RoiPixels:
+        if self._roi is None:
+            name = self.__class__.__name__
+            raise ValueError(f"No ROI defined for {name}.")
+        return self._roi
+
+    def __call__(self) -> T:
+        return self.get()
+
+    @abstractmethod
+    def get(self) -> T:
+        pass
+
+
+class DataSetter(ABC, Generic[T]):
+    def __init__(
+        self,
+        zarr_array: zarr.Array,
+        slicing_ops: SlicingOps,
+        axes_ops: AxesOps,
+        transforms: Sequence[TransformProtocol] | None = None,
+        roi: Roi | RoiPixels | None = None,
+    ) -> None:
+        self._zarr_array = zarr_array
+        self._slicing_ops = slicing_ops
+        self._axes_ops = axes_ops
+        self._transforms = transforms
+        self._roi = roi
+
+    def __repr__(self) -> str:
+        name = self.__class__.__name__
+        return (
+            f"{name}(zarr_array={self._zarr_array}, "
+            f"slicing_ops={self._slicing_ops}, "
+            f"axes_ops={self._axes_ops}, "
+            f"transforms={self._transforms})"
+        )
+
+    @property
+    def zarr_array(self) -> zarr.Array:
+        return self._zarr_array
+
+    @property
+    def slicing_ops(self) -> SlicingOps:
+        return self._slicing_ops
+
+    @property
+    def axes_ops(self) -> AxesOps:
+        return self._axes_ops
+
+    @property
+    def transforms(self) -> Sequence[TransformProtocol] | None:
+        return self._transforms
+
+    @property
+    def roi(self) -> Roi | RoiPixels:
+        if self._roi is None:
+            name = self.__class__.__name__
+            raise ValueError(f"No ROI defined for {name}.")
+        return self._roi
+
+    def __call__(self, patch: T) -> None:
+        return self.set(patch)
+
+    @abstractmethod
+    def set(self, patch: T) -> None:
+        pass
+
+
+class NumpyGetter(DataGetter[np.ndarray]):
+    def __init__(
+        self,
+        *,
+        zarr_array: zarr.Array,
+        dimensions: Dimensions,
+        axes_order: Sequence[str] | None = None,
+        transforms: Sequence[TransformProtocol] | None = None,
+        slicing_dict: dict[str, SlicingInputType] | None = None,
+        remove_channel_selection: bool = False,
+        roi: Roi | RoiPixels | None = None,
+    ) -> None:
+        """Build a pipe to get a numpy or dask array from a zarr array."""
+        slicing_ops, axes_ops = setup_io_pipe(
+            dimensions=dimensions,
+            slicing_dict=slicing_dict,
+            axes_order=axes_order,
+            remove_channel_selection=remove_channel_selection,
+        )
+        super().__init__(
+            zarr_array=zarr_array,
+            slicing_ops=slicing_ops,
+            axes_ops=axes_ops,
+            transforms=transforms,
+            roi=roi,
+        )
+
+    def get(self) -> np.ndarray:
+        """Get a numpy array from the zarr array with ops."""
+        array = get_slice_as_numpy(self._zarr_array, slicing_ops=self._slicing_ops)
+        array = get_as_numpy_axes_ops(array, axes_ops=self._axes_ops)
+        array = get_as_numpy_transform(
+            array,
+            slicing_ops=self._slicing_ops,
+            axes_ops=self._axes_ops,
+            transforms=self._transforms,
+        )
+        return array
+
+
+class DaskGetter(DataGetter[DaskArray]):
+    def __init__(
+        self,
+        *,
+        zarr_array: zarr.Array,
+        dimensions: Dimensions,
+        axes_order: Sequence[str] | None = None,
+        transforms: Sequence[TransformProtocol] | None = None,
+        slicing_dict: dict[str, SlicingInputType] | None = None,
+        remove_channel_selection: bool = False,
+        roi: Roi | RoiPixels | None = None,
+    ) -> None:
+        """Build a pipe to get a numpy or dask array from a zarr array."""
+        slicing_ops, axes_ops = setup_io_pipe(
+            dimensions=dimensions,
+            slicing_dict=slicing_dict,
+            axes_order=axes_order,
+            remove_channel_selection=remove_channel_selection,
+        )
+        super().__init__(
+            zarr_array=zarr_array,
+            slicing_ops=slicing_ops,
+            axes_ops=axes_ops,
+            transforms=transforms,
+            roi=roi,
+        )
+
+    def get(self) -> DaskArray:
+        """Get a dask array from the zarr array with ops.
+
+        The order of operations is:
+        * get slice will load the data from the zarr array
+        * get axes ops will reorder, squeeze or expand the axes
+        * get transform will apply any additional transforms
+
+        """
+        array = get_slice_as_dask(self._zarr_array, slicing_ops=self._slicing_ops)
+        array = get_as_dask_axes_ops(array, axes_ops=self._axes_ops)
+        array = get_as_dask_transform(
+            array,
+            slicing_ops=self._slicing_ops,
+            axes_ops=self._axes_ops,
+            transforms=self._transforms,
+        )
+        return array
+
+
+##############################################################
+#
+# "To Disk" Pipes
+#
+##############################################################
+
+
+class NumpySetter(DataSetter[np.ndarray]):
+    def __init__(
+        self,
+        *,
+        zarr_array: zarr.Array,
+        dimensions: Dimensions,
+        axes_order: Sequence[str] | None = None,
+        transforms: Sequence[TransformProtocol] | None = None,
+        slicing_dict: dict[str, SlicingInputType] | None = None,
+        remove_channel_selection: bool = False,
+        roi: Roi | RoiPixels | None = None,
+    ) -> None:
+        """Build a pipe to get a numpy or dask array from a zarr array."""
+        slicing_ops, axes_ops = setup_io_pipe(
+            dimensions=dimensions,
+            slicing_dict=slicing_dict,
+            axes_order=axes_order,
+            remove_channel_selection=remove_channel_selection,
+        )
+        super().__init__(
+            zarr_array=zarr_array,
+            slicing_ops=slicing_ops,
+            axes_ops=axes_ops,
+            transforms=transforms,
+            roi=roi,
+        )
+
+    def set(self, patch: np.ndarray) -> None:
+        """Get a numpy array from the zarr array with ops."""
+        patch = set_as_numpy_transform(
+            array=patch,
+            slicing_ops=self._slicing_ops,
+            axes_ops=self._axes_ops,
+            transforms=self._transforms,
+        )
+        patch = set_as_numpy_axes_ops(
+            array=patch,
+            axes_ops=self._axes_ops,
+        )
+        set_slice_as_numpy(
+            zarr_array=self._zarr_array,
+            patch=patch,
+            slicing_ops=self._slicing_ops,
+        )
+
+
+class DaskSetter(DataSetter[DaskArray]):
+    def __init__(
+        self,
+        *,
+        zarr_array: zarr.Array,
+        dimensions: Dimensions,
+        axes_order: Sequence[str] | None = None,
+        transforms: Sequence[TransformProtocol] | None = None,
+        slicing_dict: dict[str, SlicingInputType] | None = None,
+        remove_channel_selection: bool = False,
+        roi: Roi | RoiPixels | None = None,
+    ) -> None:
+        """Build a pipe to get a numpy or dask array from a zarr array."""
+        slicing_ops, axes_ops = setup_io_pipe(
+            dimensions=dimensions,
+            slicing_dict=slicing_dict,
+            axes_order=axes_order,
+            remove_channel_selection=remove_channel_selection,
+        )
+        super().__init__(
+            zarr_array=zarr_array,
+            slicing_ops=slicing_ops,
+            axes_ops=axes_ops,
+            transforms=transforms,
+            roi=roi,
+        )
+
+    def set(self, patch: DaskArray) -> None:
+        """Get a dask array from the zarr array with ops."""
+        patch = set_as_dask_transform(
+            array=patch,
+            slicing_ops=self._slicing_ops,
+            axes_ops=self._axes_ops,
+            transforms=self._transforms,
+        )
+        patch = set_as_dask_axes_ops(
+            array=patch,
+            axes_ops=self._axes_ops,
+        )
+        set_slice_as_dask(
+            zarr_array=self._zarr_array,
+            patch=patch,
+            slicing_ops=self._slicing_ops,
+        )