ngio 0.3.5__py3-none-any.whl → 0.4.0__py3-none-any.whl

ngio/experimental/iterators/_feature.py

@@ -0,0 +1,189 @@
+from collections.abc import Sequence
+from typing import TypeAlias
+
+import dask.array as da
+import numpy as np
+
+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,
+    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[NumpyPipeType, DaskPipeType]):
+    """Base class for iterators over ROIs."""
+
+    def __init__(
+        self,
+        input_image: Image,
+        input_label: Label,
+        channel_selection: ChannelSlicingInputType = None,
+        axes_order: Sequence[str] | None = None,
+        input_transforms: Sequence[TransformProtocol] | None = None,
+        label_transforms: Sequence[TransformProtocol] | None = None,
+    ) -> None:
+        """Initialize the iterator with a ROI table and input/output images.
+
+        Args:
+            input_image (Image): The input image to be used as input for the
+                segmentation.
+            input_label (Label): The input label with the segmentation masks.
+            channel_selection (ChannelSlicingInputType): Optional
+                selection of channels to use for the segmentation.
+            axes_order (Sequence[str] | None): Optional axes order for the
+                segmentation.
+            input_transforms (Sequence[TransformProtocol] | None): Optional
+                transforms to apply to the input image.
+            label_transforms (Sequence[TransformProtocol] | None): Optional
+                transforms to apply to the output label.
+        """
+        self._input = input_image
+        self._input_label = input_label
+        self._ref_image = input_image
+        self._rois = input_image.build_image_roi_table(name=None).rois()
+
+        # Set iteration parameters
+        self._input_slicing_kwargs = add_channel_selection_to_slicing_dict(
+            image=self._input, channel_selection=channel_selection, slicing_dict={}
+        )
+        self._channel_selection = channel_selection
+        self._axes_order = axes_order
+        self._input_transforms = input_transforms
+        self._label_transforms = label_transforms
+
+        self._input.require_axes_match(self._input_label)
+        self._input.require_rescalable(self._input_label)
+
+    def get_init_kwargs(self) -> dict:
+        """Return the initialization arguments for the iterator."""
+        return {
+            "input_image": self._input,
+            "input_label": self._input_label,
+            "channel_selection": self._channel_selection,
+            "axes_order": self._axes_order,
+            "input_transforms": self._input_transforms,
+            "label_transforms": self._label_transforms,
+        }
+
+    def build_numpy_getter(self, roi: Roi) -> NumpyFeatureGetter:
+        data_getter = NumpyRoiGetter(
+            zarr_array=self._input.zarr_array,
+            dimensions=self._input.dimensions,
+            axes_order=self._axes_order,
+            transforms=self._input_transforms,
+            roi=roi,
+            slicing_dict=self._input_slicing_kwargs,
+        )
+        label_getter = NumpyRoiGetter(
+            zarr_array=self._input_label.zarr_array,
+            dimensions=self._input_label.dimensions,
+            axes_order=self._axes_order,
+            transforms=self._label_transforms,
+            roi=roi,
+            remove_channel_selection=True,
+        )
+        return NumpyFeatureGetter(data_getter, label_getter)
+
+    def build_dask_getter(self, roi: Roi) -> DaskFeatureGetter:
+        data_getter = DaskRoiGetter(
+            zarr_array=self._input.zarr_array,
+            dimensions=self._input.dimensions,
+            axes_order=self._axes_order,
+            transforms=self._input_transforms,
+            roi=roi,
+            slicing_dict=self._input_slicing_kwargs,
+        )
+        label_getter = DaskRoiGetter(
+            zarr_array=self._input_label.zarr_array,
+            dimensions=self._input_label.dimensions,
+            axes_order=self._axes_order,
+            transforms=self._label_transforms,
+            roi=roi,
+            remove_channel_selection=True,
+        )
+        return DaskFeatureGetter(data_getter, label_getter)
+
+    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):  # type: ignore[override]
+        """Create an iterator over the pixels of the ROIs."""
+        return self.iter(lazy=False, data_mode="numpy", iterator_mode="readonly")
+
+    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

@@ -0,0 +1,130 @@
+from collections.abc import Sequence
+
+import dask.array as da
+import numpy as np
+
+from ngio.common import Roi
+from ngio.experimental.iterators._abstract_iterator import AbstractIteratorBuilder
+from ngio.images import Image
+from ngio.images._image import (
+    ChannelSlicingInputType,
+    add_channel_selection_to_slicing_dict,
+)
+from ngio.io_pipes import (
+    DaskRoiGetter,
+    DaskRoiSetter,
+    NumpyRoiGetter,
+    NumpyRoiSetter,
+    TransformProtocol,
+)
+from ngio.io_pipes._io_pipes_types import DataGetterProtocol, DataSetterProtocol
+
+
+class ImageProcessingIterator(AbstractIteratorBuilder[np.ndarray, da.Array]):
+    """Base class for iterators over ROIs."""
+
+    def __init__(
+        self,
+        input_image: Image,
+        output_image: Image,
+        input_channel_selection: ChannelSlicingInputType = None,
+        output_channel_selection: ChannelSlicingInputType = None,
+        axes_order: Sequence[str] | None = None,
+        input_transforms: Sequence[TransformProtocol] | None = None,
+        output_transforms: Sequence[TransformProtocol] | None = None,
+    ) -> None:
+        """Initialize the iterator with a ROI table and input/output images.
+
+        Args:
+            input_image (Image): The input image to be used as input for the
+                segmentation.
+            output_image (Image): The image where the ROIs will be written.
+            input_channel_selection (ChannelSlicingInputType): Optional
+                selection of channels to use for the input image.
+            output_channel_selection (ChannelSlicingInputType): Optional
+                selection of channels to use for the output image.
+            axes_order (Sequence[str] | None): Optional axes order for the
+                segmentation.
+            input_transforms (Sequence[TransformProtocol] | None): Optional
+                transforms to apply to the input image.
+            output_transforms (Sequence[TransformProtocol] | None): Optional
+                transforms to apply to the output label.
+        """
+        self._input = input_image
+        self._output = output_image
+        self._ref_image = input_image
+        self._rois = input_image.build_image_roi_table(name=None).rois()
+
+        # Set iteration parameters
+        self._input_slicing_kwargs = add_channel_selection_to_slicing_dict(
+            image=self._input,
+            channel_selection=input_channel_selection,
+            slicing_dict={},
+        )
+        self._output_slicing_kwargs = add_channel_selection_to_slicing_dict(
+            image=self._output,
+            channel_selection=output_channel_selection,
+            slicing_dict={},
+        )
+        self._input_channel_selection = input_channel_selection
+        self._output_channel_selection = output_channel_selection
+        self._axes_order = axes_order
+        self._input_transforms = input_transforms
+        self._output_transforms = output_transforms
+
+        self._input.require_dimensions_match(self._output, allow_singleton=True)
+
+    def get_init_kwargs(self) -> dict:
+        """Return the initialization arguments for the iterator."""
+        return {
+            "input_image": self._input,
+            "output_image": self._output,
+            "input_channel_selection": self._input_channel_selection,
+            "output_channel_selection": self._output_channel_selection,
+            "axes_order": self._axes_order,
+            "input_transforms": self._input_transforms,
+            "output_transforms": self._output_transforms,
+        }
+
+    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,
+            slicing_dict=self._input_slicing_kwargs,
+        )
+
+    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,
+            slicing_dict=self._output_slicing_kwargs,
+        )
+
+    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,
+            slicing_dict=self._input_slicing_kwargs,
+        )
+
+    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,
+            slicing_dict=self._output_slicing_kwargs,
+        )
+
+    def post_consolidate(self):
+        self._output.consolidate()

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

@@ -0,0 +1,127 @@
+from ngio import Roi, RoiPixels
+from ngio.images._abstract_image import AbstractImage
+
+
+def rois_product(rois_a: list[Roi], rois_b: list[Roi]) -> list[Roi]:
+    """Compute the product of two sets of ROIs."""
+    rois_product = []
+    for roi_a in rois_a:
+        for roi_b in rois_b:
+            intersection = roi_a.intersection(roi_b)
+            if intersection:
+                rois_product.append(intersection)
+    return rois_product
+
+
+def grid(
+    rois: list[Roi],
+    ref_image: AbstractImage,
+    size_x: int | None = None,
+    size_y: int | None = None,
+    size_z: int | None = None,
+    size_t: int | None = None,
+    stride_x: int | None = None,
+    stride_y: int | None = None,
+    stride_z: int | None = None,
+    stride_t: int | None = None,
+    base_name: str | None = None,
+) -> list[Roi]:
+    """This method is a placeholder for creating a regular grid of ROIs."""
+    t_dim = ref_image.dimensions.get("t", default=1)
+    z_dim = ref_image.dimensions.get("z", default=1)
+    y_dim = ref_image.dimensions.get("y", default=1)
+    x_dim = ref_image.dimensions.get("x", default=1)
+
+    size_t = size_t if size_t is not None else t_dim
+    size_z = size_z if size_z is not None else z_dim
+    size_y = size_y if size_y is not None else y_dim
+    size_x = size_x if size_x is not None else x_dim
+
+    stride_t = stride_t if stride_t is not None else size_t
+    stride_z = stride_z if stride_z is not None else size_z
+    stride_y = stride_y if stride_y is not None else size_y
+    stride_x = stride_x if stride_x is not None else size_x
+
+    # Here we would create a grid of ROIs based on the specified parameters.
+    new_rois = []
+    for t in range(0, t_dim, stride_t):
+        for z in range(0, z_dim, stride_z):
+            for y in range(0, y_dim, stride_y):
+                for x in range(0, x_dim, stride_x):
+                    roi = RoiPixels(
+                        name=base_name,
+                        x=x,
+                        y=y,
+                        z=z,
+                        t=t,
+                        x_length=size_x,
+                        y_length=size_y,
+                        z_length=size_z,
+                        t_length=size_t,
+                    )
+                    new_rois.append(roi.to_roi(pixel_size=ref_image.pixel_size))
+
+    return rois_product(rois, new_rois)
+
+
+def by_yx(rois: list[Roi], ref_image: AbstractImage) -> list[Roi]:
+    """Return a new iterator that iterates over ROIs by YX coordinates."""
+    return grid(
+        rois=rois,
+        ref_image=ref_image,
+        size_z=1,
+        stride_z=1,
+        size_t=1,
+        stride_t=1,
+    )
+
+
+def by_zyx(rois: list[Roi], ref_image: AbstractImage, strict: bool = True) -> list[Roi]:
+    """Return a new iterator that iterates over ROIs by ZYX coordinates."""
+    if strict and not ref_image.is_3d:
+        raise ValueError(
+            "Reference Input image must be 3D to iterate by ZXY coordinates. "
+            f"Current dimensions: {ref_image.dimensions}"
+        )
+    return grid(
+        rois=rois,
+        ref_image=ref_image,
+        size_t=1,
+        stride_t=1,
+    )
+
+
+def by_chunks(
+    rois: list[Roi],
+    ref_image: AbstractImage,
+    overlap_xy: int = 0,
+    overlap_z: int = 0,
+    overlap_t: int = 0,
+) -> list[Roi]:
+    """This method is a placeholder for chunked processing."""
+    chunk_size = ref_image.chunks
+    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
+    size_z = chunk_size[z_axis] if z_axis is not None else None
+    size_t = chunk_size[t_axis] if t_axis is not None else None
+    stride_x = size_x - overlap_xy if size_x is not None else None
+    stride_y = size_y - overlap_xy if size_y is not None else None
+    stride_z = size_z - overlap_z if size_z is not None else None
+    stride_t = size_t - overlap_t if size_t is not None else None
+    return grid(
+        rois=rois,
+        ref_image=ref_image,
+        size_x=size_x,
+        size_y=size_y,
+        size_z=size_z,
+        size_t=size_t,
+        stride_x=stride_x,
+        stride_y=stride_y,
+        stride_z=stride_z,
+        stride_t=stride_t,
+    )