ngio 0.5.0b6__py3-none-any.whl

ngio/common/_roi.py

@@ -0,0 +1,315 @@
+"""Region of interest (ROI) metadata.
+
+These are the interfaces between the ROI tables / masking ROI tables and
+    the ImageLikeHandler.
+"""
+
+from collections.abc import Callable, Mapping
+from typing import Literal, Self
+
+from pydantic import BaseModel, ConfigDict, Field, field_validator
+
+from ngio.ome_zarr_meta import PixelSize
+from ngio.utils import NgioValueError
+
+
+def world_to_pixel(value: float, pixel_size: float, eps: float = 1e-6) -> float:
+    raster_value = value / pixel_size
+
+    # If the value is very close to an integer, round it
+    # This ensures that we don't have floating point precision issues
+    # When loading ROIs that were originally defined in pixel coordinates
+    _rounded = round(raster_value)
+    if abs(_rounded - raster_value) < eps:
+        return _rounded
+    return raster_value
+
+
+def pixel_to_world(value: float, pixel_size: float) -> float:
+    return value * pixel_size
+
+
+def _join_roi_names(name1: str | None, name2: str | None) -> str | None:
+    if name1 is not None and name2 is not None:
+        if name1 == name2:
+            return name1
+        return f"{name1}:{name2}"
+    return name1 or name2
+
+
+def _join_roi_labels(label1: int | None, label2: int | None) -> int | None:
+    if label1 is not None and label2 is not None:
+        if label1 == label2:
+            return label1
+        raise NgioValueError("Cannot join ROIs with different labels")
+    return label1 or label2
+
+
+class RoiSlice(BaseModel):
+    axis_name: str
+    start: float | None = Field(default=None)
+    length: float | None = Field(default=None, ge=0)
+
+    model_config = ConfigDict(extra="forbid")
+
+    @classmethod
+    def _from_slice(
+        cls,
+        axis_name: str,
+        selection: slice,
+    ) -> "RoiSlice":
+        start = selection.start
+        length = (
+            None
+            if selection.stop is None or selection.start is None
+            else selection.stop - selection.start
+        )
+        return cls(axis_name=axis_name, start=start, length=length)
+
+    @classmethod
+    def from_value(
+        cls,
+        axis_name: str,
+        value: float | tuple[float | None, float | None] | slice,
+    ) -> "RoiSlice":
+        if isinstance(value, slice):
+            return cls._from_slice(axis_name=axis_name, selection=value)
+        elif isinstance(value, tuple):
+            return cls(axis_name=axis_name, start=value[0], length=value[1])
+        elif isinstance(value, int | float):
+            return cls(axis_name=axis_name, start=value, length=1)
+        else:
+            raise TypeError(f"Unsupported type for slice value: {type(value)}")
+
+    def __repr__(self) -> str:
+        return f"{self.axis_name}: {self.start}->{self.end}"
+
+    @property
+    def end(self) -> float | None:
+        if self.start is None or self.length is None:
+            return None
+        return self.start + self.length
+
+    def to_slice(self) -> slice:
+        return slice(self.start, self.end)
+
+    def _is_compatible(self, other: "RoiSlice", msg: str) -> None:
+        if self.axis_name != other.axis_name:
+            raise NgioValueError(
+                f"{msg}: Cannot operate on RoiSlices with different axis names"
+            )
+
+    def union(self, other: "RoiSlice") -> "RoiSlice":
+        self._is_compatible(other, "RoiSlice union failed")
+        start = min(self.start or 0, other.start or 0)
+        end = max(self.end or float("inf"), other.end or float("inf"))
+        length = end - start if end > start else 0
+        if length == float("inf"):
+            length = None
+        return RoiSlice(axis_name=self.axis_name, start=start, length=length)
+
+    def intersection(self, other: "RoiSlice") -> "RoiSlice | None":
+        self._is_compatible(other, "RoiSlice intersection failed")
+        start = max(self.start or 0, other.start or 0)
+        end = min(self.end or float("inf"), other.end or float("inf"))
+        if end <= start:
+            # No intersection
+            return None
+        length = end - start
+        if length == float("inf"):
+            length = None
+        return RoiSlice(axis_name=self.axis_name, start=start, length=length)
+
+    def to_world(self, pixel_size: float) -> "RoiSlice":
+        start = (
+            pixel_to_world(self.start, pixel_size) if self.start is not None else None
+        )
+        length = (
+            pixel_to_world(self.length, pixel_size) if self.length is not None else None
+        )
+        return RoiSlice(axis_name=self.axis_name, start=start, length=length)
+
+    def to_pixel(self, pixel_size: float) -> "RoiSlice":
+        start = (
+            world_to_pixel(self.start, pixel_size) if self.start is not None else None
+        )
+        length = (
+            world_to_pixel(self.length, pixel_size) if self.length is not None else None
+        )
+        return RoiSlice(axis_name=self.axis_name, start=start, length=length)
+
+    def zoom(self, zoom_factor: float = 1.0) -> "RoiSlice":
+        if zoom_factor <= 0:
+            raise NgioValueError("Zoom factor must be greater than 0")
+        zoom_factor -= 1.0
+        if self.length is None:
+            return self
+
+        diff_length = self.length * zoom_factor
+        length = self.length + diff_length
+        start = max((self.start or 0) - (diff_length / 2), 0)
+        return RoiSlice(axis_name=self.axis_name, start=start, length=length)
+
+
+class Roi(BaseModel):
+    name: str | None
+    slices: list[RoiSlice] = Field(min_length=2)
+    label: int | None = Field(default=None, ge=0)
+    space: Literal["world", "pixel"] = "world"
+
+    model_config = ConfigDict(extra="allow")
+
+    @field_validator("slices")
+    @classmethod
+    def validate_no_duplicate_axes(cls, v: list[RoiSlice]) -> list[RoiSlice]:
+        axis_names = [s.axis_name for s in v]
+        if len(axis_names) != len(set(axis_names)):
+            raise NgioValueError("Roi slices must have unique axis names")
+        return v
+
+    def _nice_repr__(self) -> str:
+        slices_repr = ", ".join(repr(s) for s in self.slices)
+        if self.label is None:
+            label_str = ""
+        else:
+            label_str = f", label={self.label}"
+
+        if self.name is None:
+            name_str = ""
+        else:
+            name_str = f"name={self.name}, "
+        return f"Roi({name_str}{slices_repr}{label_str}, space={self.space})"
+
+    @classmethod
+    def from_values(
+        cls,
+        slices: Mapping[str, float | tuple[float | None, float | None] | slice],
+        name: str | None,
+        label: int | None = None,
+        space: Literal["world", "pixel"] = "world",
+        **kwargs,
+    ) -> Self:
+        _slices = []
+        for axis, _slice in slices.items():
+            _slices.append(RoiSlice.from_value(axis_name=axis, value=_slice))
+        return cls.model_construct(
+            name=name, slices=_slices, label=label, space=space, **kwargs
+        )
+
+    def get(self, axis_name: str) -> RoiSlice | None:
+        for roi_slice in self.slices:
+            if roi_slice.axis_name == axis_name:
+                return roi_slice
+        return None
+
+    def get_name(self) -> str:
+        if self.name is not None:
+            return self.name
+        if self.label is not None:
+            return str(self.label)
+        return self._nice_repr__()
+
+    @staticmethod
+    def _apply_sym_ops(
+        self_slices: list[RoiSlice],
+        other_slices: list[RoiSlice],
+        op: Callable[[RoiSlice, RoiSlice], RoiSlice | None],
+    ) -> list[RoiSlice] | None:
+        self_axis_dict = {s.axis_name: s for s in self_slices}
+        other_axis_dict = {s.axis_name: s for s in other_slices}
+        common_axis_names = self_axis_dict.keys() | other_axis_dict.keys()
+        new_slices = []
+        for axis_name in common_axis_names:
+            slice_a = self_axis_dict.get(axis_name)
+            slice_b = other_axis_dict.get(axis_name)
+            if slice_a is not None and slice_b is not None:
+                result = op(slice_a, slice_b)
+                if result is None:
+                    return None
+                new_slices.append(result)
+            elif slice_a is not None:
+                new_slices.append(slice_a)
+            elif slice_b is not None:
+                new_slices.append(slice_b)
+        return new_slices
+
+    def intersection(self, other: Self) -> Self | None:
+        if self.space != other.space:
+            raise NgioValueError(
+                "Roi intersection failed: One ROI is in pixel space and the "
+                "other in world space"
+            )
+
+        out_slices = self._apply_sym_ops(
+            self.slices, other.slices, op=lambda a, b: a.intersection(b)
+        )
+        if out_slices is None:
+            return None
+
+        name = _join_roi_names(self.name, other.name)
+        label = _join_roi_labels(self.label, other.label)
+        return self.model_copy(
+            update={"name": name, "slices": out_slices, "label": label}
+        )
+
+    def union(self, other: Self) -> Self:
+        if self.space != other.space:
+            raise NgioValueError(
+                "Roi union failed: One ROI is in pixel space and the "
+                "other in world space"
+            )
+
+        out_slices = self._apply_sym_ops(
+            self.slices, other.slices, op=lambda a, b: a.union(b)
+        )
+        if out_slices is None:
+            raise NgioValueError("Roi union failed: could not compute union")
+
+        name = _join_roi_names(self.name, other.name)
+        label = _join_roi_labels(self.label, other.label)
+        return self.model_copy(
+            update={"name": name, "slices": out_slices, "label": label}
+        )
+
+    def zoom(
+        self, zoom_factor: float = 1.0, axes: tuple[str, ...] = ("x", "y")
+    ) -> Self:
+        new_slices = []
+        for roi_slice in self.slices:
+            if roi_slice.axis_name in axes:
+                new_slices.append(roi_slice.zoom(zoom_factor=zoom_factor))
+            else:
+                new_slices.append(roi_slice)
+        return self.model_copy(update={"slices": new_slices})
+
+    def to_world(self, pixel_size: PixelSize | None = None) -> Self:
+        if self.space == "world":
+            return self.model_copy()
+        if pixel_size is None:
+            raise NgioValueError(
+                "Pixel sizes must be provided to convert ROI from pixel to world"
+            )
+        new_slices = []
+        for roi_slice in self.slices:
+            pixel_size_ = pixel_size.get(roi_slice.axis_name, default=1.0)
+            new_slices.append(roi_slice.to_world(pixel_size=pixel_size_))
+        return self.model_copy(update={"slices": new_slices, "space": "world"})
+
+    def to_pixel(self, pixel_size: PixelSize | None = None) -> Self:
+        if self.space == "pixel":
+            return self.model_copy()
+
+        if pixel_size is None:
+            raise NgioValueError(
+                "Pixel sizes must be provided to convert ROI from world to pixel"
+            )
+
+        new_slices = []
+        for roi_slice in self.slices:
+            pixel_size_ = pixel_size.get(roi_slice.axis_name, default=1.0)
+            new_slices.append(roi_slice.to_pixel(pixel_size=pixel_size_))
+        return self.model_copy(update={"slices": new_slices, "space": "pixel"})
+
+    def to_slicing_dict(self, pixel_size: PixelSize | None = None) -> dict[str, slice]:
+        roi = self.to_pixel(pixel_size=pixel_size)
+        return {roi_slice.axis_name: roi_slice.to_slice() for roi_slice in roi.slices}

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

@@ -0,0 +1,188 @@
+from functools import partial
+from typing import Literal
+
+import dask.array as da
+import numpy as np
+from scipy.ndimage import zoom as scipy_zoom
+
+from ngio.utils import NgioValueError
+
+InterpolationOrder = Literal["nearest", "linear", "cubic"]
+
+
+def order_to_int(order: InterpolationOrder | Literal[0, 1, 2]) -> Literal[0, 1, 2]:
+    if order == "nearest" or order == 0:
+        return 0
+    elif order == "linear" or order == 1:
+        return 1
+    elif order == "cubic" or order == 2:
+        return 2
+    else:
+        raise NgioValueError(f"Invalid order: {order}")
+
+
+def _stacked_zoom(x, zoom_y, zoom_x, order=1, mode="grid-constant", grid_mode=True):
+    *rest, yshape, xshape = x.shape
+    x = x.reshape(-1, yshape, xshape)
+    scale_xy = (zoom_y, zoom_x)
+    _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:])
+
+
+def fast_zoom(x, zoom, order=1, mode="grid-constant", grid_mode=True, auto_stack=True):
+    """Fast zoom function.
+
+    Scipy zoom function that can handle singleton dimensions
+    but the performance degrades with the number of dimensions.
+
+    This function has two small optimizations:
+     - it removes singleton dimensions before calling zoom
+     - if it detects that the zoom is only on the last two dimensions
+         it stacks the first dimensions to call zoom only on the last two.
+    """
+    mask = np.isclose(x.shape, 1)
+    # Always keep the last two dimensions
+    # To avoid issues with singleton x or y dimensions
+    mask[-1] = False
+    mask[-2] = False
+    zoom = np.array(zoom)
+    singletons = tuple(np.where(mask)[0])
+    xs = np.squeeze(x, axis=singletons)
+    new_zoom = zoom[~mask]
+
+    *zoom_rest, zoom_y, zoom_x = new_zoom
+    if auto_stack and np.allclose(zoom_rest, 1):
+        xs = _stacked_zoom(
+            xs, zoom_y, zoom_x, order=order, mode=mode, grid_mode=grid_mode
+        )
+    else:
+        xs = scipy_zoom(xs, new_zoom, order=order, mode=mode, grid_mode=grid_mode)
+    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 | float, ...] | None = None,
+    target_shape: tuple[int, ...] | None = None,
+) -> tuple[np.ndarray, tuple[int, ...]]:
+    if scale is None and target_shape is None:
+        raise NgioValueError("Either scale or target_shape must be provided")
+
+    if scale is not None and target_shape is not None:
+        raise NgioValueError("Only one of scale or target_shape must be provided")
+
+    if scale is None:
+        assert target_shape is not None, "Target shape must be provided"
+        if len(target_shape) != source_array.ndim:
+            raise NgioValueError(
+                "Target shape must have the "
+                "same number of dimensions as "
+                "the source array"
+            )
+        _scale = np.array(target_shape) / np.array(source_array.shape)
+        _target_shape = target_shape
+    else:
+        _scale = np.array(scale)
+        _target_shape = tuple(map(int, np.round(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[float | int, ...] | None = None,
+    target_shape: tuple[int, ...] | None = None,
+    order: InterpolationOrder = "linear",
+) -> da.Array:
+    """Dask implementation of zooming an array.
+
+    Only one of scale or target_shape must be provided.
+
+    Args:
+        source_array (da.Array): The source array to zoom.
+        scale (tuple[int, ...] | None): The scale factor to zoom by.
+        target_shape (tuple[int, ...], None): The target shape to zoom to.
+        order (Literal["nearest", "linear", "cubic"]): The order of interpolation.
+            Defaults to "linear".
+
+    Returns:
+        da.Array: The zoomed array.
+    """
+    # This function follow the implementation from:
+    # https://github.com/ome/ome-zarr-py/blob/master/ome_zarr/dask_utils.py
+    # The module was contributed by Andreas Eisenbarth @aeisenbarth
+    # See https://github.com/toloudis/ome-zarr-py/pull/
+    _scale, _target_shape = _zoom_inputs_check(
+        source_array=source_array, scale=scale, target_shape=target_shape
+    )
+
+    # Rechunk to better match the scaling operation
+    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 (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))
+
+    zoom_wrapper = partial(
+        fast_zoom,
+        zoom=_scale,
+        order=order_to_int(order),
+        mode="grid-constant",
+        grid_mode=True,
+    )
+
+    out_array = da.map_blocks(
+        zoom_wrapper, source_array, chunks=block_output_shape, dtype=source_array.dtype
+    )
+
+    # Slice and rechunk to target
+    slices = tuple(slice(0, ts, 1) for ts in _target_shape)
+    out_array = out_array[slices]
+    return out_array
+
+
+def numpy_zoom(
+    source_array: np.ndarray,
+    scale: tuple[int | float, ...] | None = None,
+    target_shape: tuple[int, ...] | None = None,
+    order: InterpolationOrder = "linear",
+) -> np.ndarray:
+    """Numpy implementation of zooming an array.
+
+    Only one of scale or target_shape must be provided.
+
+    Args:
+        source_array (np.ndarray): The source array to zoom.
+        scale (tuple[int, ...] | None): The scale factor to zoom by.
+        target_shape (tuple[int, ...], None): The target shape to zoom to.
+        order (Literal[0, 1, 2]): The order of interpolation. Defaults to 1.
+
+    Returns:
+        np.ndarray: The zoomed array
+    """
+    _scale, _ = _zoom_inputs_check(
+        source_array=source_array, scale=scale, target_shape=target_shape
+    )
+
+    out_array = fast_zoom(
+        source_array,
+        zoom=_scale,
+        order=order_to_int(order),
+        mode="grid-constant",
+        grid_mode=True,
+    )
+    assert isinstance(out_array, np.ndarray)
+    return out_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,15 @@
+"""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,
+)
+
+__all__ = [
+    "FeatureExtractorIterator",
+    "ImageProcessingIterator",
+    "MaskedSegmentationIterator",
+    "SegmentationIterator",
+]