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

ngio/common/_synt_images_utils.py

@@ -0,0 +1,101 @@
+from math import ceil
+
+import numpy as np
+
+
+def _center_crop(arr: np.ndarray, target: int, axis: int) -> np.ndarray:
+    # Center-crop the array `arr` along dimension `axis` to size `target`.
+    # This assumes target < arr.shape[axis].
+    n = arr.shape[axis]
+    start = (n - target) // 2
+    end = start + target
+    slc = [slice(None)] * arr.ndim
+    slc[axis] = slice(start, end)
+    return arr[tuple(slc)]
+
+
+def _tile_to(
+    arr: np.ndarray, target: int, axis: int, label_mode: bool = False
+) -> np.ndarray:
+    # Tile the array `arr` along dimension `axis` to size `target`.
+    # This assumes target > arr.shape[axis].
+    n = arr.shape[axis]
+    reps = ceil(target / n)
+
+    tiles = []
+    flip = False
+    max_label = 0
+    for _ in range(reps):
+        if flip:
+            t_arr = np.flip(arr, axis=axis)
+        else:
+            t_arr = 1 * arr
+        if label_mode:
+            # Remove duplicate labels
+            t_arr = np.where(t_arr > 0, t_arr + max_label, 0)
+            max_label = t_arr.max()
+        tiles.append(t_arr)
+        flip = not flip
+
+    tiled = np.concatenate(tiles, axis=axis)
+
+    slc = [slice(None)] * arr.ndim
+    slc[axis] = slice(0, target)
+    return tiled[tuple(slc)]
+
+
+def _fit_to_shape_2d(
+    src: np.ndarray, out_shape: tuple[int, int], label_mode: bool = False
+) -> np.ndarray:
+    """Fit a 2D array to a target shape by center-cropping or tiling as necessary."""
+    out_r, out_c = out_shape
+    arr = src
+    if out_r < arr.shape[0]:
+        arr = _center_crop(arr, out_r, axis=0)
+    else:
+        arr = _tile_to(arr, out_r, axis=0, label_mode=label_mode)
+
+    if out_c < arr.shape[1]:
+        arr = _center_crop(arr, out_c, axis=1)
+    else:
+        arr = _tile_to(arr, out_c, axis=1, label_mode=label_mode)
+    return arr
+
+
+def fit_to_shape(
+    arr: np.ndarray, out_shape: tuple[int, ...], ensure_unique_info: bool = False
+) -> np.ndarray:
+    """Fit a 2D array to a target shape.
+
+    The x,y dimensions of `arr` are fitted to the last two dimensions of
+    `out_shape` by center-cropping or tiling as necessary.
+    The other dimensions are broadcasted as necessary.
+
+    WARNING: This does not zoom the image, it only crops or tiles it.
+
+    Args:
+        arr (np.ndarray): The input 2D array.
+        out_shape (tuple[int, ...]): The target shape. Must have at least 2
+            and at most 5 dimensions.
+        ensure_unique_info (bool, optional): If True, assumes that `arr` is a label
+            image and ensures that labels do not overlap when tiling. Defaults to False.
+
+    Returns:
+        np.ndarray: The fitted array with shape `out_shape`.
+    """
+    if len(out_shape) < 2:
+        raise ValueError("`out_shape` must contain at least 2 dimensions.")
+
+    if len(out_shape) > 5:
+        raise ValueError("`out_shape` must contain at most 5 dimensions.")
+
+    if any(d <= 0 for d in out_shape):
+        raise ValueError("`out_shape` must contain positive integers.")
+
+    if arr.ndim != 2:
+        raise ValueError("`arr` must be a 2D array.")
+
+    *_, sy, sx = out_shape
+    arr = _fit_to_shape_2d(arr, out_shape=(sy, sx), label_mode=ensure_unique_info)
+    arr = np.broadcast_to(arr, out_shape)
+    return arr

ngio/common/_zoom.py

@@ -12,12 +12,11 @@ def _stacked_zoom(x, zoom_y, zoom_x, order=1, mode="grid-constant", grid_mode=Tr
     *rest, yshape, xshape = x.shape
     x = x.reshape(-1, yshape, xshape)
     scale_xy = (zoom_y, zoom_x)
-    x_out = np.stack(
-        [
-            scipy_zoom(x[i], scale_xy, order=order, mode=mode, grid_mode=True)
-            for i in range(x.shape[0])
-        ]
-    )
+    _x_out = [
+        scipy_zoom(x[i], scale_xy, order=order, mode=mode, grid_mode=grid_mode)
+        for i in range(x.shape[0])
+    ]
+    x_out = np.stack(_x_out)  # type: ignore (scipy_zoom returns np.ndarray, but type is not inferred correctly)
     return x_out.reshape(*rest, *x_out.shape[1:])
 
 
@@ -45,13 +44,13 @@ def fast_zoom(x, zoom, order=1, mode="grid-constant", grid_mode=True, auto_stack
         )
     else:
         xs = scipy_zoom(xs, new_zoom, order=order, mode=mode, grid_mode=grid_mode)
-    x = np.expand_dims(xs, axis=singletons)
+    x = np.expand_dims(xs, axis=singletons)  # type: ignore (scipy_zoom returns np.ndarray, but type is not inferred correctly)
     return x
 
 
 def _zoom_inputs_check(
     source_array: np.ndarray | da.Array,
-    scale: tuple[int, ...] | None = None,
+    scale: tuple[int | float, ...] | None = None,
     target_shape: tuple[int, ...] | None = None,
 ) -> tuple[np.ndarray, tuple[int, ...]]:
     if scale is None and target_shape is None:
@@ -74,12 +73,18 @@ def _zoom_inputs_check(
         _scale = np.array(scale)
         _target_shape = tuple(np.array(source_array.shape) * scale)
 
+    if len(_scale) != source_array.ndim:
+        raise NgioValueError(
+            f"Cannot scale array of shape {source_array.shape} with factors {_scale}."
+            " Target shape must have the same number of dimensions as the source array."
+        )
+
     return _scale, _target_shape
 
 
 def dask_zoom(
     source_array: da.Array,
-    scale: tuple[int, ...] | None = None,
+    scale: tuple[float | int, ...] | None = None,
     target_shape: tuple[int, ...] | None = None,
     order: Literal[0, 1, 2] = 1,
 ) -> da.Array:
@@ -106,10 +111,10 @@ def dask_zoom(
     )
 
     # Rechunk to better match the scaling operation
-    source_chunks = np.array(source_array.chunksize)
+    source_chunks = np.array(source_array.chunksize)  # type: ignore (da.Array.chunksize is a tuple of ints)
     better_source_chunks = np.maximum(1, np.round(source_chunks * _scale) / _scale)
     better_source_chunks = better_source_chunks.astype(int)
-    source_array = source_array.rechunk(better_source_chunks)  # type: ignore
+    source_array = source_array.rechunk(better_source_chunks)  # type: ignore (better_source_chunks is a valid input for rechunk)
 
     # Calculate the block output shape
     block_output_shape = tuple(np.ceil(better_source_chunks * _scale).astype(int))
@@ -130,7 +135,7 @@ def dask_zoom(
 
 def numpy_zoom(
     source_array: np.ndarray,
-    scale: tuple[int, ...] | None = None,
+    scale: tuple[int | float, ...] | None = None,
     target_shape: tuple[int, ...] | None = None,
     order: Literal[0, 1, 2] = 1,
 ) -> np.ndarray:

ngio/common/transforms/__init__.py

@@ -0,0 +1,5 @@
+"""Concrete IO transformations."""
+
+from ngio.common.transforms._zoom import ZoomTransform
+
+__all__ = ["ZoomTransform"]

ngio/common/transforms/_label.py

@@ -0,0 +1,12 @@
+"""Transforms and pre-post-processors for label images."""
+
+# def make_unique_label_np(x: np.ndarray, p: int, n: int) -> np.ndarray:
+#    """Make a unique label for the patch."""
+#    x = np.where(x > 0, (1 + p - n) + x * n, 0)
+#    return x
+#
+#
+# def make_unique_label_da(x: da.Array, p: int, n: int) -> da.Array:
+#    """Make a unique label for the patch."""
+#    x = da.where(x > 0, (1 + p - n) + x * n, 0)
+#    return x

ngio/common/transforms/_zoom.py

@@ -0,0 +1,109 @@
+from collections.abc import Sequence
+from typing import Literal
+
+import dask.array as da
+import numpy as np
+
+from ngio.common._array_io_utils import apply_sequence_axes_ops
+from ngio.common._dimensions import Dimensions
+from ngio.common._zoom import dask_zoom, numpy_zoom
+from ngio.ome_zarr_meta.ngio_specs import SlicingOps
+
+
+class ZoomTransform:
+    def __init__(self, scale: Sequence[float], order: Literal[0, 1, 2]):
+        self._scale = tuple(scale)
+        self._order: Literal[0, 1, 2] = order
+
+    @property
+    def scale(self) -> tuple[float, ...]:
+        return self._scale
+
+    @property
+    def inv_scale(self) -> tuple[float, ...]:
+        return tuple([1 / s for s in self._scale])
+
+    @classmethod
+    def from_dimensions(
+        cls,
+        original_dimension: Dimensions,
+        target_dimension: Dimensions,
+        order: Literal[0, 1, 2],
+    ):
+        scale = []
+        for o_ax_name in original_dimension.axes_mapper.axes_names:
+            t_ax = target_dimension.axes_mapper.get_axis(name=o_ax_name)
+            if t_ax is None:
+                _scale = 1
+            else:
+                t_shape = target_dimension.get(o_ax_name)
+                o_shape = original_dimension.get(o_ax_name)
+                assert t_shape is not None and o_shape is not None
+                _scale = t_shape / o_shape
+            scale.append(_scale)
+
+        return cls(scale, order)
+
+    def apply_numpy_transform(
+        self, array: np.ndarray, slicing_ops: SlicingOps
+    ) -> np.ndarray:
+        """Apply the scaling transformation to a numpy array."""
+        scale = tuple(
+            apply_sequence_axes_ops(
+                self.scale,
+                default=1,
+                squeeze_axes=slicing_ops.squeeze_axes,
+                transpose_axes=slicing_ops.transpose_axes,
+                expand_axes=slicing_ops.expand_axes,
+            )
+        )
+        array = numpy_zoom(source_array=array, scale=scale, order=self._order)
+        return array
+
+    def apply_dask_transform(
+        self, array: da.Array, slicing_ops: SlicingOps
+    ) -> da.Array:
+        """Apply the scaling transformation to a dask array."""
+        scale = tuple(
+            apply_sequence_axes_ops(
+                self.scale,
+                default=1,
+                squeeze_axes=slicing_ops.squeeze_axes,
+                transpose_axes=slicing_ops.transpose_axes,
+                expand_axes=slicing_ops.expand_axes,
+            )
+        )
+        array = dask_zoom(source_array=array, scale=scale, order=self._order)
+        return array
+
+    def apply_inverse_numpy_transform(
+        self, array: np.ndarray, slicing_ops: SlicingOps
+    ) -> np.ndarray:
+        """Apply the inverse scaling transformation to a numpy array."""
+        scale = tuple(
+            apply_sequence_axes_ops(
+                self.inv_scale,
+                default=1,
+                squeeze_axes=slicing_ops.squeeze_axes,
+                transpose_axes=slicing_ops.transpose_axes,
+                expand_axes=slicing_ops.expand_axes,
+            )
+        )
+        array = numpy_zoom(source_array=array, scale=scale, order=self._order)
+        return array
+
+    def apply_inverse_dask_transform(
+        self, array: da.Array, slicing_ops: SlicingOps
+    ) -> da.Array:
+        """Apply the inverse scaling transformation to a dask array."""
+        scale = tuple(
+            apply_sequence_axes_ops(
+                self.inv_scale,
+                default=1,
+                squeeze_axes=slicing_ops.squeeze_axes,
+                transpose_axes=slicing_ops.transpose_axes,
+                expand_axes=slicing_ops.expand_axes,
+            )
+        )
+        array = dask_zoom(source_array=array, scale=scale, order=self._order)
+        return array

ngio/experimental/__init__.py

@@ -0,0 +1,5 @@
+"""This module provides experimental features.
+
+Use with caution as these features may change or be removed in future releases
+without notice.
+"""

ngio/experimental/iterators/__init__.py

@@ -0,0 +1,17 @@
+"""This file is part of NGIO, a library for working with OME-Zarr data."""
+
+from ngio.experimental.iterators._feature import FeatureExtractorIterator
+from ngio.experimental.iterators._image_processing import ImageProcessingIterator
+from ngio.experimental.iterators._segmentation import (
+    MaskedSegmentationIterator,
+    SegmentationIterator,
+)
+
+# from ngio.experimental.iterators._builder import IteratorBuilder
+
+__all__ = [
+    "FeatureExtractorIterator",
+    "ImageProcessingIterator",
+    "MaskedSegmentationIterator",
+    "SegmentationIterator",
+]

ngio/experimental/iterators/_abstract_iterator.py

@@ -0,0 +1,170 @@
+from abc import ABC, abstractmethod
+from collections.abc import Callable
+from typing import Self
+
+from ngio import Roi
+from ngio.experimental.iterators._rois_utils import (
+    by_chunks,
+    by_yx,
+    by_zyx,
+    grid,
+    rois_product,
+)
+from ngio.images._abstract_image import AbstractImage
+from ngio.tables import GenericRoiTable
+
+
+class AbstractIteratorBuilder(ABC):
+    """Base class for building iterators over ROIs."""
+
+    _rois: list[Roi]
+    _ref_image: AbstractImage
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(regions={len(self._rois)})"
+
+    @abstractmethod
+    def get_init_kwargs(self) -> dict:
+        """Return the initialization arguments for the iterator."""
+        pass
+
+    @property
+    def rois(self) -> list[Roi]:
+        """Get the list of ROIs for the iterator."""
+        return self._rois
+
+    def _set_rois(self, rois: list[Roi]) -> None:
+        """Set the list of ROIs for the iterator."""
+        self._rois = rois
+
+    @property
+    def ref_image(self) -> AbstractImage:
+        """Get the reference image for the iterator."""
+        return self._ref_image
+
+    def _new_from_rois(self, rois: list[Roi]) -> Self:
+        """Create a new instance of the iterator with a different set of ROIs."""
+        init_kwargs = self.get_init_kwargs()
+        new_instance = self.__class__(**init_kwargs)
+        new_instance._set_rois(rois)
+        return new_instance
+
+    def grid(
+        self,
+        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 = "",
+    ) -> Self:
+        """Create a grid of ROIs based on the input image dimensions."""
+        rois = grid(
+            rois=self.rois,
+            ref_image=self.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,
+            base_name=base_name,
+        )
+        return self._new_from_rois(rois)
+
+    def by_yx(self) -> Self:
+        """Return a new iterator that iterates over ROIs by YX coordinates."""
+        rois = by_yx(self.rois, self.ref_image)
+        return self._new_from_rois(rois)
+
+    def by_zyx(self, strict: bool = True) -> Self:
+        """Return a new iterator that iterates over ROIs by ZYX coordinates.
+
+        Args:
+            strict (bool): If True, only iterate over ZYX if a Z axis
+                is present and not of size 1.
+
+        """
+        rois = by_zyx(self.rois, self.ref_image, strict=strict)
+        return self._new_from_rois(rois)
+
+    def by_chunks(self, overlap_xy: int = 0, overlap_z: int = 0) -> Self:
+        """Return a new iterator that iterates over ROIs by chunks.
+
+        Args:
+            overlap_xy (int): Overlap in XY dimensions.
+            overlap_z (int): Overlap in Z dimension.
+
+        Returns:
+            SegmentationIterator: A new iterator with chunked ROIs.
+        """
+        rois = by_chunks(
+            self.rois, self.ref_image, overlap_xy=overlap_xy, overlap_z=overlap_z
+        )
+        return self._new_from_rois(rois)
+
+    def product(self, other: list[Roi] | GenericRoiTable) -> Self:
+        """Cartesian product of the current ROIs with an arbitrary list of ROIs."""
+        if isinstance(other, GenericRoiTable):
+            other = other.rois()
+        rois = rois_product(self.rois, other)
+        return self._new_from_rois(rois)
+
+    @abstractmethod
+    def build_numpy_getter(self, roi: Roi):
+        """Build a getter function for the given ROI."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def build_numpy_setter(self, roi: Roi):
+        """Build a setter function for the given ROI."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def build_dask_getter(self, roi: Roi):
+        """Build a Dask reader function for the given ROI."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def build_dask_setter(self, roi: Roi):
+        """Build a Dask setter function for the given ROI."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def post_consolidate(self) -> None:
+        """Post-process the consolidated data."""
+        raise NotImplementedError
+
+    def iter_as_numpy(self):
+        """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)
+        self.post_consolidate()
+
+    def iter_as_dask(self):
+        """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)
+
+    def map_as_numpy(self, func: Callable) -> 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)
+        self.post_consolidate()
+
+    def map_as_dask(self, func: Callable) -> 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)
+        self.post_consolidate()

ngio/experimental/iterators/_feature.py

@@ -0,0 +1,151 @@
+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_numpy_getter,
+)
+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,
+)
+
+
+class FeatureExtractorIterator(AbstractIteratorBuilder):
+    """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().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
+
+    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):
+        data_getter = 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,
+        )
+        label_getter = build_roi_numpy_getter(
+            zarr_array=self._input_label.zarr_array,
+            dimensions=self._input_label.dimensions,
+            axes_order=self._axes_order,
+            transforms=self._label_transforms,
+            pixel_size=self._input_label.pixel_size,
+            roi=roi,
+            remove_channel_selection=True,
+        )
+        return lambda: (data_getter(), label_getter(), roi)
+
+    def build_numpy_setter(self, roi: Roi):
+        return None
+
+    def build_dask_getter(self, roi: Roi):
+        data_getter = 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,
+        )
+        label_getter = build_roi_dask_getter(
+            zarr_array=self._input_label.zarr_array,
+            dimensions=self._input_label.dimensions,
+            axes_order=self._axes_order,
+            transforms=self._label_transforms,
+            pixel_size=self._input_label.pixel_size,
+            roi=roi,
+            remove_channel_selection=True,
+        )
+        return lambda: (data_getter(), label_getter(), roi)
+
+    def build_dask_setter(self, roi: Roi):
+        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 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.")