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

ngio/experimental/iterators/_mappers.py

@@ -0,0 +1,48 @@
+"""Mappers for iterators.
+
+Mappers are classes that can be passed to the `map` method of iterators to
+transform the items yielded by the iterator.
+
+"""
+
+from collections.abc import Callable, Iterable
+from typing import Generic, Protocol, TypeVar
+
+from ngio.io_pipes._io_pipes_types import DataGetterProtocol, DataSetterProtocol
+from ngio.utils import NgioValueError
+
+T = TypeVar("T")
+
+
+class MapperProtocol(Protocol[T]):
+    """Protocol for mappers."""
+
+    def __call__(
+        self,
+        func: Callable[[T], T],
+        getters: Iterable[DataGetterProtocol[T]],
+        setters: Iterable[DataSetterProtocol[T] | None],
+    ) -> None:
+        """Map an item to another item."""
+        ...
+
+
+class BasicMapper(Generic[T]):
+    """A basic mapper that simply applies a function to the data."""
+
+    def __call__(
+        self,
+        func: Callable[[T], T],
+        getters: Iterable[DataGetterProtocol[T]],
+        setters: Iterable[DataSetterProtocol[T] | None],
+    ) -> None:
+        """Map an item to another item."""
+        for getter, setter in zip(getters, setters, strict=True):
+            data = getter()
+            data = func(data)
+            if setter is None:
+                raise NgioValueError(
+                    "Error in BasicMapper: setter is None, "
+                    "this iterator is read-only, mapping is not possible."
+                )
+            setter(data)

ngio/experimental/iterators/_rois_utils.py

@@ -100,10 +100,10 @@ def by_chunks(
 ) -> list[Roi]:
     """This method is a placeholder for chunked processing."""
     chunk_size = ref_image.chunks
-    t_axis = ref_image.axes_mapper.get_index("t")
-    z_axis = ref_image.axes_mapper.get_index("z")
-    y_axis = ref_image.axes_mapper.get_index("y")
-    x_axis = ref_image.axes_mapper.get_index("x")
+    t_axis = ref_image.axes_handler.get_index("t")
+    z_axis = ref_image.axes_handler.get_index("z")
+    y_axis = ref_image.axes_handler.get_index("y")
+    x_axis = ref_image.axes_handler.get_index("x")
 
     size_x = chunk_size[x_axis] if x_axis is not None else None
     size_y = chunk_size[y_axis] if y_axis is not None else None

ngio/experimental/iterators/_segmentation.py

@@ -1,20 +1,9 @@
-from collections.abc import Callable, Generator, Sequence
+from collections.abc import Sequence
 
 import dask.array as da
 import numpy as np
 
-from ngio.common import (
-    Roi,
-    TransformProtocol,
-    build_roi_dask_getter,
-    build_roi_dask_setter,
-    build_roi_masked_dask_getter,
-    build_roi_masked_dask_setter,
-    build_roi_masked_numpy_getter,
-    build_roi_masked_numpy_setter,
-    build_roi_numpy_getter,
-    build_roi_numpy_setter,
-)
+from ngio.common import Roi
 from ngio.experimental.iterators._abstract_iterator import AbstractIteratorBuilder
 from ngio.images import Image, Label
 from ngio.images._image import (
@@ -22,10 +11,21 @@ from ngio.images._image import (
     add_channel_selection_to_slicing_dict,
 )
 from ngio.images._masked_image import MaskedImage
-from ngio.utils._errors import NgioValidationError
+from ngio.io_pipes import (
+    DaskGetterMasked,
+    DaskRoiGetter,
+    DaskRoiSetter,
+    DaskSetterMasked,
+    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__(
@@ -66,12 +66,7 @@ class SegmentationIterator(AbstractIteratorBuilder):
         self._input_transforms = input_transforms
         self._output_transforms = output_transforms
 
-        # Check compatibility between input and output images
-        if not self._input.dimensions.is_compatible_with(self._output.dimensions):
-            raise NgioValidationError(
-                "Input image and output label have incompatible dimensions. "
-                f"Input: {self._input.dimensions}, Output: {self._output.dimensions}."
-            )
+        self._input.require_dimensions_match(self._output, allow_singleton=False)
 
     def get_init_kwargs(self) -> dict:
         """Return the initialization arguments for the iterator."""
@@ -84,83 +79,49 @@ class SegmentationIterator(AbstractIteratorBuilder):
             "output_transforms": self._output_transforms,
         }
 
-    def build_numpy_getter(self, roi: Roi):
-        return build_roi_numpy_getter(
+    def build_numpy_getter(self, roi: Roi) -> DataGetterProtocol[np.ndarray]:
+        return NumpyRoiGetter(
             zarr_array=self._input.zarr_array,
             dimensions=self._input.dimensions,
+            roi=roi,
             axes_order=self._axes_order,
             transforms=self._input_transforms,
-            pixel_size=self._input.pixel_size,
-            roi=roi,
             slicing_dict=self._input_slicing_kwargs,
         )
 
-    def build_numpy_setter(self, roi: Roi):
-        return build_roi_numpy_setter(
+    def build_numpy_setter(self, roi: Roi) -> DataSetterProtocol[np.ndarray]:
+        return NumpyRoiSetter(
             zarr_array=self._output.zarr_array,
             dimensions=self._output.dimensions,
+            roi=roi,
             axes_order=self._axes_order,
             transforms=self._output_transforms,
-            pixel_size=self._output.pixel_size,
-            roi=roi,
             remove_channel_selection=True,
         )
 
-    def build_dask_getter(self, roi: Roi):
-        return build_roi_dask_getter(
+    def build_dask_getter(self, roi: Roi) -> DataGetterProtocol[da.Array]:
+        return DaskRoiGetter(
             zarr_array=self._input.zarr_array,
             dimensions=self._input.dimensions,
+            roi=roi,
             axes_order=self._axes_order,
             transforms=self._input_transforms,
-            pixel_size=self._input.pixel_size,
-            roi=roi,
             slicing_dict=self._input_slicing_kwargs,
         )
 
-    def build_dask_setter(self, roi: Roi):
-        return build_roi_dask_setter(
+    def build_dask_setter(self, roi: Roi) -> DataSetterProtocol[da.Array]:
+        return DaskRoiSetter(
             zarr_array=self._output.zarr_array,
             dimensions=self._output.dimensions,
+            roi=roi,
             axes_order=self._axes_order,
             transforms=self._output_transforms,
-            pixel_size=self._output.pixel_size,
-            roi=roi,
             remove_channel_selection=True,
         )
 
     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."""
@@ -205,11 +166,11 @@ class MaskedSegmentationIterator(SegmentationIterator):
         self._output_transforms = output_transforms
 
         # Check compatibility between input and output images
-        if not self._input.dimensions.is_compatible_with(self._output.dimensions):
-            raise NgioValidationError(
-                "Input image and output label have incompatible dimensions. "
-                f"Input: {self._input.dimensions}, Output: {self._output.dimensions}."
-            )
+        # if not self._input.dimensions.is_compatible_with(self._output.dimensions):
+        #    raise NgioValidationError(
+        #        "Input image and output label have incompatible dimensions. "
+        #        f"Input: {self._input.dimensions}, Output: {self._output.dimensions}."
+        #    )
 
     def get_init_kwargs(self) -> dict:
         """Return the initialization arguments for the iterator."""
@@ -223,58 +184,50 @@ class MaskedSegmentationIterator(SegmentationIterator):
         }
 
     def build_numpy_getter(self, roi: Roi):
-        return build_roi_masked_numpy_getter(
-            roi=roi,
+        return NumpyGetterMasked(
             zarr_array=self._input.zarr_array,
             dimensions=self._input.dimensions,
+            roi=roi,
             label_zarr_array=self._input._label.zarr_array,
             label_dimensions=self._input._label.dimensions,
-            label_pixel_size=self._input._label.pixel_size,
             axes_order=self._axes_order,
             transforms=self._input_transforms,
-            pixel_size=self._input.pixel_size,
             slicing_dict=self._input_slicing_kwargs,
         )
 
     def build_numpy_setter(self, roi: Roi):
-        return build_roi_masked_numpy_setter(
+        return NumpySetterMasked(
             roi=roi,
             zarr_array=self._output.zarr_array,
             dimensions=self._output.dimensions,
             label_zarr_array=self._input._label.zarr_array,
             label_dimensions=self._input._label.dimensions,
-            label_pixel_size=self._input._label.pixel_size,
             axes_order=self._axes_order,
             transforms=self._output_transforms,
-            pixel_size=self._output.pixel_size,
             remove_channel_selection=True,
         )
 
     def build_dask_getter(self, roi: Roi):
-        return build_roi_masked_dask_getter(
+        return DaskGetterMasked(
             roi=roi,
             zarr_array=self._input.zarr_array,
             dimensions=self._input.dimensions,
             label_zarr_array=self._input._label.zarr_array,
             label_dimensions=self._input._label.dimensions,
-            label_pixel_size=self._input._label.pixel_size,
             axes_order=self._axes_order,
             transforms=self._input_transforms,
-            pixel_size=self._input.pixel_size,
             slicing_dict=self._input_slicing_kwargs,
         )
 
     def build_dask_setter(self, roi: Roi):
-        return build_roi_masked_dask_setter(
+        return DaskSetterMasked(
             roi=roi,
             zarr_array=self._output.zarr_array,
             dimensions=self._output.dimensions,
             label_zarr_array=self._input._label.zarr_array,
             label_dimensions=self._input._label.dimensions,
-            label_pixel_size=self._input._label.pixel_size,
             axes_order=self._axes_order,
             transforms=self._output_transforms,
-            pixel_size=self._output.pixel_size,
             remove_channel_selection=True,
         )