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

ngio/experimental/iterators/_image_processing.py

@@ -0,0 +1,169 @@
+from collections.abc import Callable, Generator, 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_numpy_getter,
+    build_roi_numpy_setter,
+)
+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.utils._errors import NgioValidationError
+
+
+class ImageProcessingIterator(AbstractIteratorBuilder):
+    """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().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
+
+        # 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}."
+            )
+
+    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):
+        return build_roi_numpy_getter(
+            zarr_array=self._input.zarr_array,
+            dimensions=self._input.dimensions,
+            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(
+            zarr_array=self._output.zarr_array,
+            dimensions=self._output.dimensions,
+            axes_order=self._axes_order,
+            transforms=self._output_transforms,
+            pixel_size=self._output.pixel_size,
+            roi=roi,
+            slicing_dict=self._output_slicing_kwargs,
+        )
+
+    def build_dask_getter(self, roi: Roi):
+        return build_roi_dask_getter(
+            zarr_array=self._input.zarr_array,
+            dimensions=self._input.dimensions,
+            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(
+            zarr_array=self._output.zarr_array,
+            dimensions=self._output.dimensions,
+            axes_order=self._axes_order,
+            transforms=self._output_transforms,
+            pixel_size=self._output.pixel_size,
+            roi=roi,
+            slicing_dict=self._output_slicing_kwargs,
+        )
+
+    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/_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 = "",
+) -> 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=f"{base_name}({t}, {z}, {y}, {x})",
+                        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_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")
+
+    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,
+    )

ngio/experimental/iterators/_segmentation.py

@@ -0,0 +1,282 @@
+from collections.abc import Callable, Generator, 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.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.images._masked_image import MaskedImage
+from ngio.utils._errors import NgioValidationError
+
+
+class SegmentationIterator(AbstractIteratorBuilder):
+    """Base class for iterators over ROIs."""
+
+    def __init__(
+        self,
+        input_image: Image,
+        output_label: Label,
+        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_label (Label): The label image where the ROIs will be written.
+            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.
+            output_transforms (Sequence[TransformProtocol] | None): Optional
+                transforms to apply to the output label.
+        """
+        self._input = input_image
+        self._output = output_label
+        self._ref_image = input_image
+        self._rois = input_image.build_image_roi_table().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._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}."
+            )
+
+    def get_init_kwargs(self) -> dict:
+        """Return the initialization arguments for the iterator."""
+        return {
+            "input_image": self._input,
+            "output_label": self._output,
+            "channel_selection": self._channel_selection,
+            "axes_order": self._axes_order,
+            "input_transforms": self._input_transforms,
+            "output_transforms": self._output_transforms,
+        }
+
+    def build_numpy_getter(self, roi: Roi):
+        return build_roi_numpy_getter(
+            zarr_array=self._input.zarr_array,
+            dimensions=self._input.dimensions,
+            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(
+            zarr_array=self._output.zarr_array,
+            dimensions=self._output.dimensions,
+            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(
+            zarr_array=self._input.zarr_array,
+            dimensions=self._input.dimensions,
+            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(
+            zarr_array=self._output.zarr_array,
+            dimensions=self._output.dimensions,
+            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."""
+
+    def __init__(
+        self,
+        input_image: MaskedImage,
+        output_label: Label,
+        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 (MaskedImage): The input image to be used as input for the
+                segmentation.
+            output_label (Label): The label image where the ROIs will be written.
+            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.
+            output_transforms (Sequence[TransformProtocol] | None): Optional
+                transforms to apply to the output label.
+        """
+        self._input = input_image
+        self._output = output_label
+
+        self._ref_image = input_image
+        self._set_rois(input_image._masking_roi_table.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._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}."
+            )
+
+    def get_init_kwargs(self) -> dict:
+        """Return the initialization arguments for the iterator."""
+        return {
+            "input_image": self._input,
+            "output_label": self._output,
+            "channel_selection": self._channel_selection,
+            "axes_order": self._axes_order,
+            "input_transforms": self._input_transforms,
+            "output_transforms": self._output_transforms,
+        }
+
+    def build_numpy_getter(self, roi: Roi):
+        return build_roi_masked_numpy_getter(
+            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_numpy_setter(self, roi: Roi):
+        return build_roi_masked_numpy_setter(
+            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(
+            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(
+            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 post_consolidate(self):
+        self._output.consolidate()