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

ngio/experimental/iterators/_segmentation.py

@@ -1,4 +1,4 @@
-from collections.abc import Callable, Generator, Sequence
+from collections.abc import Sequence
 
 import dask.array as da
 import numpy as np
@@ -16,17 +16,16 @@ from ngio.io_pipes import (
     DaskRoiGetter,
     DaskRoiSetter,
     DaskSetterMasked,
-    DataGetter,
-    DataSetter,
     NumpyGetterMasked,
     NumpyRoiGetter,
     NumpyRoiSetter,
     NumpySetterMasked,
     TransformProtocol,
 )
+from ngio.io_pipes._io_pipes_types import DataGetterProtocol, DataSetterProtocol
 
 
-class SegmentationIterator(AbstractIteratorBuilder):
+class SegmentationIterator(AbstractIteratorBuilder[np.ndarray, da.Array]):
     """Base class for iterators over ROIs."""
 
     def __init__(
@@ -80,7 +79,7 @@ class SegmentationIterator(AbstractIteratorBuilder):
             "output_transforms": self._output_transforms,
         }
 
-    def build_numpy_getter(self, roi: Roi) -> DataGetter[np.ndarray]:
+    def build_numpy_getter(self, roi: Roi) -> DataGetterProtocol[np.ndarray]:
         return NumpyRoiGetter(
             zarr_array=self._input.zarr_array,
             dimensions=self._input.dimensions,
@@ -90,7 +89,7 @@ class SegmentationIterator(AbstractIteratorBuilder):
             slicing_dict=self._input_slicing_kwargs,
         )
 
-    def build_numpy_setter(self, roi: Roi) -> DataSetter[np.ndarray]:
+    def build_numpy_setter(self, roi: Roi) -> DataSetterProtocol[np.ndarray]:
         return NumpyRoiSetter(
             zarr_array=self._output.zarr_array,
             dimensions=self._output.dimensions,
@@ -100,7 +99,7 @@ class SegmentationIterator(AbstractIteratorBuilder):
             remove_channel_selection=True,
         )
 
-    def build_dask_getter(self, roi: Roi) -> DataGetter[da.Array]:
+    def build_dask_getter(self, roi: Roi) -> DataGetterProtocol[da.Array]:
         return DaskRoiGetter(
             zarr_array=self._input.zarr_array,
             dimensions=self._input.dimensions,
@@ -110,7 +109,7 @@ class SegmentationIterator(AbstractIteratorBuilder):
             slicing_dict=self._input_slicing_kwargs,
         )
 
-    def build_dask_setter(self, roi: Roi) -> DataSetter[da.Array]:
+    def build_dask_setter(self, roi: Roi) -> DataSetterProtocol[da.Array]:
         return DaskRoiSetter(
             zarr_array=self._output.zarr_array,
             dimensions=self._output.dimensions,
@@ -123,36 +122,6 @@ class SegmentationIterator(AbstractIteratorBuilder):
     def post_consolidate(self):
         self._output.consolidate()
 
-    def iter_as_numpy(
-        self,
-    ) -> Generator[tuple[np.ndarray, Callable[[np.ndarray], None]]]:
-        """Create an iterator over the pixels of the ROIs as Dask arrays.
-
-        Returns:
-            Generator[tuple[da.Array, DaskWriter]]: An iterator the input
-                image as Dask arrays and a writer to write the output
-                to the label image.
-        """
-        return super().iter_as_numpy()
-
-    def map_as_numpy(self, func: Callable[[np.ndarray], np.ndarray]) -> None:
-        """Apply a transformation function to the ROI pixels."""
-        return super().map_as_numpy(func)
-
-    def iter_as_dask(self) -> Generator[tuple[da.Array, Callable[[da.Array], None]]]:
-        """Create an iterator over the pixels of the ROIs as Dask arrays.
-
-        Returns:
-            Generator[tuple[da.Array, DaskWriter]]: An iterator the input
-                image as Dask arrays and a writer to write the output
-                to the label image.
-        """
-        return super().iter_as_dask()
-
-    def map_as_dask(self, func: Callable[[da.Array], da.Array]) -> None:
-        """Apply a transformation function to the ROI pixels."""
-        return super().map_as_dask(func)
-
 
 class MaskedSegmentationIterator(SegmentationIterator):
     """Base class for iterators over ROIs."""

ngio/images/_abstract_image.py

@@ -90,7 +90,11 @@ class AbstractImage(Generic[_image_handler]):
     @property
     def dimensions(self) -> Dimensions:
         """Return the dimensions of the image."""
-        return Dimensions(shape=self.zarr_array.shape, dataset=self.dataset)
+        return Dimensions(
+            shape=self.zarr_array.shape,
+            chunks=self.zarr_array.chunks,
+            dataset=self.dataset,
+        )
 
     @property
     def pixel_size(self) -> PixelSize:
@@ -169,7 +173,7 @@ class AbstractImage(Generic[_image_handler]):
 
     def has_axis(self, axis: str) -> bool:
         """Return True if the image has the given axis."""
-        return self.dimensions.has_axis(axis)
+        return self.axes_handler.has_axis(axis)
 
     def _get_as_numpy(
         self,
@@ -461,7 +465,7 @@ class AbstractImage(Generic[_image_handler]):
         return roi_px.to_roi(pixel_size=self.pixel_size)
 
     def build_image_roi_table(self, name: str | None = "image") -> RoiTable:
-        """Build the ROI table for an image."""
+        """Build the ROI table containing the ROI covering the entire image."""
         return RoiTable(rois=[self.roi(name=name)])
 
     def require_dimensions_match(
@@ -483,6 +487,25 @@ class AbstractImage(Generic[_image_handler]):
             other.dimensions, allow_singleton=allow_singleton
         )
 
+    def check_if_dimensions_match(
+        self,
+        other: "AbstractImage",
+        allow_singleton: bool = False,
+    ) -> bool:
+        """Check if two images have matching spatial dimensions.
+
+        Args:
+            other: The other image to compare to.
+            allow_singleton: If True, allow singleton dimensions to be
+                compatible with non-singleton dimensions.
+
+        Returns:
+            bool: True if the images have matching dimensions, False otherwise.
+        """
+        return self.dimensions.check_if_dimensions_match(
+            other.dimensions, allow_singleton=allow_singleton
+        )
+
     def require_axes_match(
         self,
         other: "AbstractImage",
@@ -497,7 +520,22 @@ class AbstractImage(Generic[_image_handler]):
         """
         self.dimensions.require_axes_match(other.dimensions)
 
-    def require_can_be_rescaled(
+    def check_if_axes_match(
+        self,
+        other: "AbstractImage",
+    ) -> bool:
+        """Check if two images have compatible axes.
+
+        Args:
+            other: The other image to compare to.
+
+        Returns:
+            bool: True if the images have compatible axes, False otherwise.
+
+        """
+        return self.dimensions.check_if_axes_match(other.dimensions)
+
+    def require_rescalable(
         self,
         other: "AbstractImage",
     ) -> None:
@@ -512,7 +550,24 @@ class AbstractImage(Generic[_image_handler]):
         Raises:
             NgioValueError: If the images cannot be scaled to each other.
         """
-        self.dimensions.require_can_be_rescaled(other.dimensions)
+        self.dimensions.require_rescalable(other.dimensions)
+
+    def check_if_rescalable(
+        self,
+        other: "AbstractImage",
+    ) -> bool:
+        """Check if two images can be rescaled to each other.
+
+        For this to be true, the images must have the same axes, and
+        the pixel sizes must be compatible (i.e. one can be scaled to the other).
+
+        Args:
+            other: The other image to compare to.
+
+        Returns:
+            bool: True if the images can be rescaled to each other, False otherwise.
+        """
+        return self.dimensions.check_if_rescalable(other.dimensions)
 
 
 def consolidate_image(

ngio/images/_image.py

@@ -641,6 +641,8 @@ class ImagesContainer:
             name=name,
             chunks=chunks,
             dtype=dtype,
+            dimension_separator=dimension_separator,
+            compressor=compressor,
             overwrite=overwrite,
         )
 

ngio/images/_label.py

@@ -216,6 +216,8 @@ class LabelsContainer:
             axes_names=axes_names,
             chunks=chunks,
             dtype=dtype,
+            dimension_separator=dimension_separator,
+            compressor=compressor,
             overwrite=overwrite,
         )
 

ngio/images/_masked_image.py

@@ -217,24 +217,29 @@ class MaskedImage(Image):
         **slicing_kwargs: SlicingInputType,
     ) -> np.ndarray | da.Array:
         """Return the masked array for a given label."""
-        slicing_kwargs = add_channel_selection_to_slicing_dict(
-            image=self, channel_selection=channel_selection, slicing_dict=slicing_kwargs
-        )
+        if mode == "numpy":
+            return self.get_roi_masked_as_numpy(
+                label=label,
+                channel_selection=channel_selection,
+                zoom_factor=zoom_factor,
+                axes_order=axes_order,
+                transforms=transforms,
+                allow_rescaling=allow_rescaling,
+                **slicing_kwargs,
+            )
 
-        roi = self._masking_roi_table.get_label(label)
-        roi = roi.zoom(zoom_factor)
-        masked_getter = NumpyGetterMasked(
-            roi=roi,
-            zarr_array=self.zarr_array,
-            label_zarr_array=self._label.zarr_array,
-            dimensions=self.dimensions,
-            label_dimensions=self._label.dimensions,
-            axes_order=axes_order,
-            transforms=transforms,
-            slicing_dict=slicing_kwargs,
-            allow_rescaling=allow_rescaling,
-        )
-        return masked_getter()
+        elif mode == "dask":
+            return self.get_roi_masked_as_dask(
+                label=label,
+                channel_selection=channel_selection,
+                zoom_factor=zoom_factor,
+                axes_order=axes_order,
+                transforms=transforms,
+                allow_rescaling=allow_rescaling,
+                **slicing_kwargs,
+            )
+        else:
+            raise ValueError(f"Unknown mode: {mode}")
 
     def set_roi_masked(
         self,

ngio/io_pipes/__init__.py

@@ -1,4 +1,31 @@
-"""I/O pipes for reading and writing data from zarr to numpy and dask arrays."""
+"""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,
@@ -20,9 +47,8 @@ from ngio.io_pipes._io_pipes_roi import (
     NumpyRoiGetter,
     NumpyRoiSetter,
 )
-from ngio.io_pipes._io_pipes_utils import SlicingInputType
 from ngio.io_pipes._match_shape import dask_match_shape, numpy_match_shape
-from ngio.io_pipes._ops_slices import SlicingOps, SlicingType
+from ngio.io_pipes._ops_slices import SlicingInputType, SlicingOps, SlicingType
 from ngio.io_pipes._ops_transforms import TransformProtocol
 
 __all__ = [

ngio/io_pipes/_io_pipes.py

@@ -6,16 +6,20 @@ import numpy as np
 import zarr
 from dask.array import Array as DaskArray
 
-from ngio.common import Dimensions
-from ngio.io_pipes._io_pipes_utils import SlicingInputType, setup_io_pipe
+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,
@@ -28,7 +32,29 @@ from ngio.io_pipes._ops_transforms import (
     set_as_dask_transform,
     set_as_numpy_transform,
 )
-from ngio.ome_zarr_meta.ngio_specs._axes import AxesOps
+
+
+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
+
 
 ##############################################################
 #
@@ -36,21 +62,32 @@ from ngio.ome_zarr_meta.ngio_specs._axes import AxesOps
 #
 ##############################################################
 
-ArrayType = TypeVar("ArrayType", np.ndarray, DaskArray)
+T = TypeVar("T")
 
 
-class DataGetter(ABC, Generic[ArrayType]):
+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:
@@ -64,34 +101,48 @@ class DataGetter(ABC, Generic[ArrayType]):
     def axes_ops(self) -> AxesOps:
         return self._axes_ops
 
-    @property
-    def in_memory_axes(self) -> tuple[str, ...]:
-        return self._axes_ops.in_memory_axes
-
     @property
     def transforms(self) -> Sequence[TransformProtocol] | None:
         return self._transforms
 
-    def __call__(self) -> ArrayType:
+    @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) -> ArrayType:
+    def get(self) -> T:
         pass
 
 
-class DataSetter(ABC, Generic[ArrayType]):
+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:
@@ -105,31 +156,36 @@ class DataSetter(ABC, Generic[ArrayType]):
     def axes_ops(self) -> AxesOps:
         return self._axes_ops
 
-    @property
-    def in_memory_axes(self) -> tuple[str, ...]:
-        return self._axes_ops.in_memory_axes
-
     @property
     def transforms(self) -> Sequence[TransformProtocol] | None:
         return self._transforms
 
-    def __call__(self, patch: ArrayType) -> None:
+    @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: ArrayType) -> None:
+    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(
@@ -143,9 +199,11 @@ class NumpyGetter(DataGetter[np.ndarray]):
             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(
@@ -160,12 +218,14 @@ class NumpyGetter(DataGetter[np.ndarray]):
 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(
@@ -179,9 +239,18 @@ class DaskGetter(DataGetter[DaskArray]):
             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(
@@ -203,12 +272,14 @@ class DaskGetter(DataGetter[DaskArray]):
 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(
@@ -222,6 +293,7 @@ class NumpySetter(DataSetter[np.ndarray]):
             slicing_ops=slicing_ops,
             axes_ops=axes_ops,
             transforms=transforms,
+            roi=roi,
         )
 
     def set(self, patch: np.ndarray) -> None:
@@ -246,12 +318,14 @@ class NumpySetter(DataSetter[np.ndarray]):
 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(
@@ -265,6 +339,7 @@ class DaskSetter(DataSetter[DaskArray]):
             slicing_ops=slicing_ops,
             axes_ops=axes_ops,
             transforms=transforms,
+            roi=roi,
         )
 
     def set(self, patch: DaskArray) -> None: