ngio 0.4.0a4__py3-none-any.whl → 0.4.1__py3-none-any.whl

ngio/experimental/iterators/_abstract_iterator.py

@@ -1,8 +1,9 @@
 from abc import ABC, abstractmethod
-from collections.abc import Callable
-from typing import Self
+from collections.abc import Callable, Generator
+from typing import Generic, Literal, Self, TypeVar, overload
 
 from ngio import Roi
+from ngio.experimental.iterators._mappers import BasicMapper, MapperProtocol
 from ngio.experimental.iterators._rois_utils import (
     by_chunks,
     by_yx,
@@ -11,10 +12,16 @@ from ngio.experimental.iterators._rois_utils import (
     rois_product,
 )
 from ngio.images._abstract_image import AbstractImage
+from ngio.io_pipes._io_pipes_types import DataGetterProtocol, DataSetterProtocol
+from ngio.io_pipes._ops_slices_utils import check_if_regions_overlap
 from ngio.tables import GenericRoiTable
+from ngio.utils import NgioValueError
 
+NumpyPipeType = TypeVar("NumpyPipeType")
+DaskPipeType = TypeVar("DaskPipeType")
 
-class AbstractIteratorBuilder(ABC):
+
+class AbstractIteratorBuilder(ABC, Generic[NumpyPipeType, DaskPipeType]):
     """Base class for building iterators over ROIs."""
 
     _rois: list[Roi]
@@ -25,7 +32,11 @@ class AbstractIteratorBuilder(ABC):
 
     @abstractmethod
     def get_init_kwargs(self) -> dict:
-        """Return the initialization arguments for the iterator."""
+        """Return the initialization arguments for the iterator.
+
+        This is used to clone the iterator with the same parameters
+        after every "product" operation.
+        """
         pass
 
     @property
@@ -116,22 +127,22 @@ class AbstractIteratorBuilder(ABC):
         return self._new_from_rois(rois)
 
     @abstractmethod
-    def build_numpy_getter(self, roi: Roi):
+    def build_numpy_getter(self, roi: Roi) -> DataGetterProtocol[NumpyPipeType]:
         """Build a getter function for the given ROI."""
         raise NotImplementedError
 
     @abstractmethod
-    def build_numpy_setter(self, roi: Roi):
+    def build_numpy_setter(self, roi: Roi) -> DataSetterProtocol[NumpyPipeType] | None:
         """Build a setter function for the given ROI."""
         raise NotImplementedError
 
     @abstractmethod
-    def build_dask_getter(self, roi: Roi):
+    def build_dask_getter(self, roi: Roi) -> DataGetterProtocol[DaskPipeType]:
         """Build a Dask reader function for the given ROI."""
         raise NotImplementedError
 
     @abstractmethod
-    def build_dask_setter(self, roi: Roi):
+    def build_dask_setter(self, roi: Roi) -> DataSetterProtocol[DaskPipeType] | None:
         """Build a Dask setter function for the given ROI."""
         raise NotImplementedError
 
@@ -140,31 +151,240 @@ class AbstractIteratorBuilder(ABC):
         """Post-process the consolidated data."""
         raise NotImplementedError
 
-    def iter_as_numpy(self):
+    def _numpy_getters_generator(self) -> Generator[DataGetterProtocol[NumpyPipeType]]:
+        """Return a list of numpy getter functions for all ROIs."""
+        yield from (self.build_numpy_getter(roi) for roi in self.rois)
+
+    def _dask_getters_generator(self) -> Generator[DataGetterProtocol[DaskPipeType]]:
+        """Return a list of dask getter functions for all ROIs."""
+        yield from (self.build_dask_getter(roi) for roi in self.rois)
+
+    def _numpy_setters_generator(
+        self,
+    ) -> Generator[DataSetterProtocol[NumpyPipeType] | None]:
+        """Return a list of numpy setter functions for all ROIs."""
+        yield from (self.build_numpy_setter(roi) for roi in self.rois)
+
+    def _dask_setters_generator(
+        self,
+    ) -> Generator[DataSetterProtocol[DaskPipeType] | None]:
+        """Return a list of dask setter functions for all ROIs."""
+        yield from (self.build_dask_setter(roi) for roi in self.rois)
+
+    def _read_and_write_generator(
+        self,
+        getters: Generator[
+            DataGetterProtocol[NumpyPipeType] | DataGetterProtocol[DaskPipeType]
+        ],
+        setters: Generator[
+            DataSetterProtocol[NumpyPipeType] | DataSetterProtocol[DaskPipeType] | None
+        ],
+    ) -> Generator[
+        tuple[
+            DataGetterProtocol[NumpyPipeType] | DataGetterProtocol[DaskPipeType],
+            DataSetterProtocol[NumpyPipeType] | DataSetterProtocol[DaskPipeType],
+        ]
+    ]:
         """Create an iterator over the pixels of the ROIs."""
-        for roi in self.rois:
-            data = self.build_numpy_getter(roi)()
-            yield data, self.build_numpy_setter(roi)
+        for getter, setter in zip(getters, setters, strict=True):
+            if setter is None:
+                name = self.__class__.__name__
+                raise NgioValueError(f"Iterator is read-only: {name}")
+            yield getter, setter
         self.post_consolidate()
 
-    def iter_as_dask(self):
+    @overload
+    def iter(
+        self,
+        lazy: Literal[True],
+        data_mode: Literal["numpy"],
+        iterator_mode: Literal["readwrite"],
+    ) -> Generator[
+        tuple[DataGetterProtocol[NumpyPipeType], DataSetterProtocol[NumpyPipeType]]
+    ]: ...
+
+    @overload
+    def iter(
+        self,
+        lazy: Literal[True],
+        data_mode: Literal["numpy"],
+        iterator_mode: Literal["readonly"] = ...,
+    ) -> Generator[DataGetterProtocol[NumpyPipeType]]: ...
+
+    @overload
+    def iter(
+        self,
+        lazy: Literal[True],
+        data_mode: Literal["dask"],
+        iterator_mode: Literal["readwrite"],
+    ) -> Generator[
+        tuple[DataGetterProtocol[DaskPipeType], DataSetterProtocol[DaskPipeType]]
+    ]: ...
+
+    @overload
+    def iter(
+        self,
+        lazy: Literal[True],
+        data_mode: Literal["dask"],
+        iterator_mode: Literal["readonly"] = ...,
+    ) -> Generator[DataGetterProtocol[DaskPipeType]]: ...
+
+    @overload
+    def iter(
+        self,
+        lazy: Literal[False],
+        data_mode: Literal["numpy"],
+        iterator_mode: Literal["readwrite"],
+    ) -> Generator[tuple[NumpyPipeType, DataSetterProtocol[NumpyPipeType]]]: ...
+
+    @overload
+    def iter(
+        self,
+        lazy: Literal[False],
+        data_mode: Literal["numpy"],
+        iterator_mode: Literal["readonly"] = ...,
+    ) -> Generator[NumpyPipeType]: ...
+
+    @overload
+    def iter(
+        self,
+        lazy: Literal[False],
+        data_mode: Literal["dask"],
+        iterator_mode: Literal["readwrite"],
+    ) -> Generator[tuple[DaskPipeType, DataSetterProtocol[DaskPipeType]]]: ...
+
+    @overload
+    def iter(
+        self,
+        lazy: Literal[False],
+        data_mode: Literal["dask"],
+        iterator_mode: Literal["readonly"] = ...,
+    ) -> Generator[DaskPipeType]: ...
+
+    def iter(
+        self,
+        lazy: bool = False,
+        data_mode: Literal["numpy", "dask"] = "dask",
+        iterator_mode: Literal["readwrite", "readonly"] = "readwrite",
+    ) -> Generator:
         """Create an iterator over the pixels of the ROIs."""
-        for roi in self.rois:
-            data = self.build_dask_getter(roi)()
-            yield data, self.build_dask_setter(roi)
+        if data_mode == "numpy":
+            getters = self._numpy_getters_generator()
+            setters = self._numpy_setters_generator()
+        elif data_mode == "dask":
+            getters = self._dask_getters_generator()
+            setters = self._dask_setters_generator()
+        else:
+            raise NgioValueError(f"Invalid mode: {data_mode}")
+
+        if iterator_mode == "readonly":
+            if lazy:
+                return getters
+            else:
+                return (getter() for getter in getters)
+        if lazy:
+            return self._read_and_write_generator(getters, setters)
+        else:
+            gen = self._read_and_write_generator(getters, setters)
+            return ((getter(), setter) for getter, setter in gen)
 
-    def map_as_numpy(self, func: Callable) -> None:
+    def iter_as_numpy(
+        self,
+    ):
+        """Create an iterator over the pixels of the ROIs."""
+        return self.iter(lazy=False, data_mode="numpy", iterator_mode="readwrite")
+
+    def iter_as_dask(
+        self,
+    ):
+        """Create an iterator over the pixels of the ROIs."""
+        return self.iter(lazy=False, data_mode="dask", iterator_mode="readwrite")
+
+    def map_as_numpy(
+        self,
+        func: Callable[[NumpyPipeType], NumpyPipeType],
+        mapper: MapperProtocol[NumpyPipeType] | None = None,
+    ) -> None:
         """Apply a transformation function to the ROI pixels."""
-        for roi in self.rois:
-            data = self.build_numpy_getter(roi)()
-            data = func(data)
-            self.build_numpy_setter(roi)
+        if mapper is None:
+            _mapper = BasicMapper[NumpyPipeType]()
+        else:
+            _mapper = mapper
+
+        _mapper(
+            func=func,
+            getters=self._numpy_getters_generator(),
+            setters=self._numpy_setters_generator(),
+        )
         self.post_consolidate()
 
-    def map_as_dask(self, func: Callable) -> None:
+    def map_as_dask(
+        self,
+        func: Callable[[DaskPipeType], DaskPipeType],
+        mapper: MapperProtocol[DaskPipeType] | None = None,
+    ) -> None:
         """Apply a transformation function to the ROI pixels."""
-        for roi in self.rois:
-            data = self.build_dask_getter(roi)()
-            data = func(data)
-            self.build_dask_setter(roi)
+        if mapper is None:
+            _mapper = BasicMapper[DaskPipeType]()
+        else:
+            _mapper = mapper
+
+        _mapper(
+            func=func,
+            getters=self._dask_getters_generator(),
+            setters=self._dask_setters_generator(),
+        )
         self.post_consolidate()
+
+    def check_if_regions_overlap(self) -> bool:
+        """Check if any of the ROIs overlap logically.
+
+        If two ROIs cover the same pixel, they are considered to overlap.
+        This does not consider chunking or other storage details.
+
+        Returns:
+            bool: True if any ROIs overlap. False otherwise.
+        """
+        if len(self.rois) < 2:
+            # Less than 2 ROIs cannot overlap
+            return False
+
+        slicing_tuples = (
+            g.slicing_ops.normalized_slicing_tuple
+            for g in self._numpy_getters_generator()
+        )
+        x = check_if_regions_overlap(slicing_tuples)
+        return x
+
+    def require_no_regions_overlap(self) -> None:
+        """Ensure that the Iterator's ROIs do not overlap."""
+        if self.check_if_regions_overlap():
+            raise NgioValueError("Some rois overlap.")
+
+    def check_if_chunks_overlap(self) -> bool:
+        """Check if any of the ROIs overlap in terms of chunks.
+
+        If two ROIs cover the same chunk, they are considered to overlap in chunks.
+        This does not consider pixel-level overlaps.
+
+        Returns:
+            bool: True if any ROIs overlap in chunks, False otherwise.
+        """
+        from ngio.io_pipes._ops_slices_utils import check_if_chunks_overlap
+
+        if len(self.rois) < 2:
+            # Less than 2 ROIs cannot overlap
+            return False
+
+        slicing_tuples = (
+            g.slicing_ops.normalized_slicing_tuple
+            for g in self._numpy_getters_generator()
+        )
+        shape = self.ref_image.shape
+        chunks = self.ref_image.chunks
+        return check_if_chunks_overlap(slicing_tuples, shape, chunks)
+
+    def require_no_chunks_overlap(self) -> None:
+        """Ensure that the ROIs do not overlap in terms of chunks."""
+        if self.check_if_chunks_overlap():
+            raise NgioValueError("Some rois overlap in chunks.")

ngio/experimental/iterators/_feature.py

@@ -1,19 +1,84 @@
-from collections.abc import Callable, Generator, Sequence
+from collections.abc import Sequence
+from typing import TypeAlias
 
 import dask.array as da
 import numpy as np
 
-from ngio.common import Roi
+from ngio.common import Roi, RoiPixels
 from ngio.experimental.iterators._abstract_iterator import AbstractIteratorBuilder
 from ngio.images import Image, Label
 from ngio.images._image import (
     ChannelSlicingInputType,
     add_channel_selection_to_slicing_dict,
 )
-from ngio.io_pipes import DaskRoiGetter, NumpyRoiGetter, TransformProtocol
+from ngio.io_pipes import (
+    DaskRoiGetter,
+    DataGetter,
+    NumpyRoiGetter,
+    TransformProtocol,
+)
+
+NumpyPipeType: TypeAlias = tuple[np.ndarray, np.ndarray, Roi | RoiPixels]
+DaskPipeType: TypeAlias = tuple[da.Array, da.Array, Roi | RoiPixels]
+
+
+class NumpyFeatureGetter(DataGetter[NumpyPipeType]):
+    def __init__(
+        self,
+        image_getter: NumpyRoiGetter,
+        label_getter: NumpyRoiGetter,
+    ) -> None:
+        self._image_getter = image_getter
+        self._label_getter = label_getter
+        super().__init__(
+            zarr_array=self._image_getter.zarr_array,
+            slicing_ops=self._image_getter.slicing_ops,
+            axes_ops=self._image_getter.axes_ops,
+            transforms=self._image_getter.transforms,
+            roi=self._image_getter.roi,
+        )
+
+    def get(self) -> NumpyPipeType:
+        return self._image_getter(), self._label_getter(), self.roi
+
+    @property
+    def image(self) -> np.ndarray:
+        return self._image_getter()
+
+    @property
+    def label(self) -> np.ndarray:
+        return self._label_getter()
+
+
+class DaskFeatureGetter(DataGetter[DaskPipeType]):
+    def __init__(
+        self,
+        image_getter: DaskRoiGetter,
+        label_getter: DaskRoiGetter,
+    ) -> None:
+        self._image_getter = image_getter
+        self._label_getter = label_getter
+        super().__init__(
+            zarr_array=self._image_getter.zarr_array,
+            slicing_ops=self._image_getter.slicing_ops,
+            axes_ops=self._image_getter.axes_ops,
+            transforms=self._image_getter.transforms,
+            roi=self._image_getter.roi,
+        )
+
+    def get(self) -> DaskPipeType:
+        return self._image_getter(), self._label_getter(), self.roi
+
+    @property
+    def image(self) -> da.Array:
+        return self._image_getter()
 
+    @property
+    def label(self) -> da.Array:
+        return self._label_getter()
 
-class FeatureExtractorIterator(AbstractIteratorBuilder):
+
+class FeatureExtractorIterator(AbstractIteratorBuilder[NumpyPipeType, DaskPipeType]):
     """Base class for iterators over ROIs."""
 
     def __init__(
@@ -55,7 +120,7 @@ class FeatureExtractorIterator(AbstractIteratorBuilder):
         self._label_transforms = label_transforms
 
         self._input.require_axes_match(self._input_label)
-        self._input.require_can_be_rescaled(self._input_label)
+        self._input.require_rescalable(self._input_label)
 
     def get_init_kwargs(self) -> dict:
         """Return the initialization arguments for the iterator."""
@@ -68,7 +133,7 @@ class FeatureExtractorIterator(AbstractIteratorBuilder):
             "label_transforms": self._label_transforms,
         }
 
-    def build_numpy_getter(self, roi: Roi):
+    def build_numpy_getter(self, roi: Roi) -> NumpyFeatureGetter:
         data_getter = NumpyRoiGetter(
             zarr_array=self._input.zarr_array,
             dimensions=self._input.dimensions,
@@ -85,12 +150,9 @@ class FeatureExtractorIterator(AbstractIteratorBuilder):
             roi=roi,
             remove_channel_selection=True,
         )
-        return lambda: (data_getter(), label_getter(), roi)
-
-    def build_numpy_setter(self, roi: Roi):
-        return None
+        return NumpyFeatureGetter(data_getter, label_getter)
 
-    def build_dask_getter(self, roi: Roi):
+    def build_dask_getter(self, roi: Roi) -> DaskFeatureGetter:
         data_getter = DaskRoiGetter(
             zarr_array=self._input.zarr_array,
             dimensions=self._input.dimensions,
@@ -107,40 +169,21 @@ class FeatureExtractorIterator(AbstractIteratorBuilder):
             roi=roi,
             remove_channel_selection=True,
         )
-        return lambda: (data_getter(), label_getter(), roi)
+        return DaskFeatureGetter(data_getter, label_getter)
 
-    def build_dask_setter(self, roi: Roi):
+    def build_numpy_setter(self, roi: Roi) -> None:
+        return None
+
+    def build_dask_setter(self, roi: Roi) -> None:
         return None
 
     def post_consolidate(self):
         pass
 
-    def iter_as_numpy(self) -> Generator[tuple[np.ndarray, np.ndarray, Roi]]:  # type: ignore (non compatible override)
-        """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.
-        """
-        for (data, label, roi), _ in super().iter_as_numpy():
-            yield data, label, roi
-
-    def map_as_numpy(self, func: Callable[[np.ndarray], np.ndarray]) -> None:
-        """Apply a transformation function to the ROI pixels."""
-        raise NotImplementedError("Numpy mapping not implemented for this iterator.")
-
-    def iter_as_dask(self) -> Generator[tuple[da.Array, da.Array, Roi]]:  # type: ignore (non compatible override)
-        """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.
-        """
-        for (data, label, roi), _ in super().iter_as_dask():
-            yield data, label, roi
+    def iter_as_numpy(self):  # type: ignore[override]
+        """Create an iterator over the pixels of the ROIs."""
+        return self.iter(lazy=False, data_mode="numpy", iterator_mode="readonly")
 
-    def map_as_dask(self, func: Callable[[da.Array], da.Array]) -> None:
-        """Apply a transformation function to the ROI pixels."""
-        raise NotImplementedError("Dask mapping not implemented for this iterator.")
+    def iter_as_dask(self):  # type: ignore[override]
+        """Create an iterator over the pixels of the ROIs."""
+        return self.iter(lazy=False, data_mode="dask", iterator_mode="readonly")

ngio/experimental/iterators/_image_processing.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
@@ -17,9 +17,10 @@ from ngio.io_pipes import (
     NumpyRoiSetter,
     TransformProtocol,
 )
+from ngio.io_pipes._io_pipes_types import DataGetterProtocol, DataSetterProtocol
 
 
-class ImageProcessingIterator(AbstractIteratorBuilder):
+class ImageProcessingIterator(AbstractIteratorBuilder[np.ndarray, da.Array]):
     """Base class for iterators over ROIs."""
 
     def __init__(
@@ -85,7 +86,7 @@ class ImageProcessingIterator(AbstractIteratorBuilder):
             "output_transforms": self._output_transforms,
         }
 
-    def build_numpy_getter(self, roi: Roi):
+    def build_numpy_getter(self, roi: Roi) -> DataGetterProtocol[np.ndarray]:
         return NumpyRoiGetter(
             zarr_array=self._input.zarr_array,
             dimensions=self._input.dimensions,
@@ -95,7 +96,7 @@ class ImageProcessingIterator(AbstractIteratorBuilder):
             slicing_dict=self._input_slicing_kwargs,
         )
 
-    def build_numpy_setter(self, roi: Roi):
+    def build_numpy_setter(self, roi: Roi) -> DataSetterProtocol[np.ndarray]:
         return NumpyRoiSetter(
             zarr_array=self._output.zarr_array,
             dimensions=self._output.dimensions,
@@ -105,7 +106,7 @@ class ImageProcessingIterator(AbstractIteratorBuilder):
             slicing_dict=self._output_slicing_kwargs,
         )
 
-    def build_dask_getter(self, roi: Roi):
+    def build_dask_getter(self, roi: Roi) -> DataGetterProtocol[da.Array]:
         return DaskRoiGetter(
             zarr_array=self._input.zarr_array,
             dimensions=self._input.dimensions,
@@ -115,7 +116,7 @@ class ImageProcessingIterator(AbstractIteratorBuilder):
             slicing_dict=self._input_slicing_kwargs,
         )
 
-    def build_dask_setter(self, roi: Roi):
+    def build_dask_setter(self, roi: Roi) -> DataSetterProtocol[da.Array]:
         return DaskRoiSetter(
             zarr_array=self._output.zarr_array,
             dimensions=self._output.dimensions,
@@ -127,33 +128,3 @@ class ImageProcessingIterator(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)

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)