ngio 0.5.0b6__py3-none-any.whl

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,126 @@
+from ngio import Roi
+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 = Roi.from_values(
+                        name=base_name,
+                        slices={
+                            "x": (x, size_x),
+                            "y": (y, size_y),
+                            "z": (z, size_z),
+                            "t": (t, size_t),
+                        },
+                        space="pixel",
+                    )
+                    new_rois.append(roi.to_world(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,
+    )

ngio/experimental/iterators/_segmentation.py

@@ -0,0 +1,235 @@
+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, Label
+from ngio.images._image import (
+    ChannelSlicingInputType,
+    add_channel_selection_to_slicing_dict,
+)
+from ngio.images._masked_image import MaskedImage
+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[np.ndarray, da.Array]):
+    """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(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._output_transforms = output_transforms
+
+        self._input.require_dimensions_match(self._output, allow_singleton=False)
+
+    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) -> 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,
+            remove_channel_selection=True,
+        )
+
+    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,
+            remove_channel_selection=True,
+        )
+
+    def post_consolidate(self):
+        self._output.consolidate()
+
+
+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 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,
+            axes_order=self._axes_order,
+            transforms=self._input_transforms,
+            slicing_dict=self._input_slicing_kwargs,
+        )
+
+    def build_numpy_setter(self, roi: Roi):
+        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,
+            axes_order=self._axes_order,
+            transforms=self._output_transforms,
+            remove_channel_selection=True,
+        )
+
+    def build_dask_getter(self, roi: Roi):
+        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,
+            axes_order=self._axes_order,
+            transforms=self._input_transforms,
+            slicing_dict=self._input_slicing_kwargs,
+        )
+
+    def build_dask_setter(self, roi: Roi):
+        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,
+            axes_order=self._axes_order,
+            transforms=self._output_transforms,
+            remove_channel_selection=True,
+        )
+
+    def post_consolidate(self):
+        self._output.consolidate()

ngio/hcs/__init__.py

@@ -0,0 +1,19 @@
+"""OME-Zarr HCS objects models."""
+
+from ngio.hcs._plate import (
+    OmeZarrPlate,
+    OmeZarrWell,
+    create_empty_plate,
+    create_empty_well,
+    open_ome_zarr_plate,
+    open_ome_zarr_well,
+)
+
+__all__ = [
+    "OmeZarrPlate",
+    "OmeZarrWell",
+    "create_empty_plate",
+    "create_empty_well",
+    "open_ome_zarr_plate",
+    "open_ome_zarr_well",
+]