ngio 0.1.5__py3-none-any.whl → 0.2.0__py3-none-any.whl

ngio/common/_masking_roi.py

@@ -0,0 +1,158 @@
+"""Utilities to build masking regions of interest (ROIs)."""
+
+import itertools
+
+import dask
+import dask.array as da
+import dask.delayed
+import numpy as np
+import scipy.ndimage as ndi
+
+from ngio.common._roi import Roi, RoiPixels
+from ngio.ome_zarr_meta import PixelSize
+from ngio.utils import NgioValueError
+
+
+def _compute_offsets(chunks):
+    """Given a chunks tuple, compute cumulative offsets for each axis.
+
+    Returns a list where each element is a list of offsets for that dimension.
+    """
+    offsets = []
+    for dim_chunks in chunks:
+        dim_offsets = [0]
+        for size in dim_chunks:
+            dim_offsets.append(dim_offsets[-1] + size)
+        offsets.append(dim_offsets)
+    return offsets
+
+
+def _adjust_slices(slices, offset):
+    """Adjust slices to global coordinates using the provided offset."""
+    adjusted_slices = {}
+    for label, s in slices.items():
+        adjusted = tuple(
+            slice(s_dim.start + off, s_dim.stop + off)
+            for s_dim, off in zip(s, offset, strict=True)
+        )
+        adjusted_slices[label] = adjusted
+    return adjusted_slices
+
+
+@dask.delayed
+def _process_chunk(chunk, offset):
+    """Process a single chunk.
+
+    run ndi.find_objects and adjust the slices
+    to global coordinates using the provided offset.
+    """
+    local_slices = compute_slices(chunk)
+    local_slices = _adjust_slices(local_slices, offset)
+    return local_slices
+
+
+def _merge_slices(
+    slice1: tuple[slice, ...], slice2: tuple[slice, ...]
+) -> tuple[slice, ...]:
+    """Merge two slices."""
+    merged = []
+    for s1, s2 in zip(slice1, slice2, strict=True):
+        start = min(s1.start, s2.start)
+        stop = max(s1.stop, s2.stop)
+        merged.append(slice(start, stop))
+    return tuple(merged)
+
+
+@dask.delayed
+def _collect_slices(
+    local_slices: list[dict[int, tuple[slice, ...]]],
+) -> dict[int, tuple[slice]]:
+    """Collect the slices from the delayed results."""
+    global_slices = {}
+    for result in local_slices:
+        for label, s in result.items():
+            if label in global_slices:
+                global_slices[label] = _merge_slices(global_slices[label], s)
+            else:
+                global_slices[label] = s
+    return global_slices
+
+
+def compute_slices(segmentation: np.ndarray) -> dict[int, tuple[slice, ...]]:
+    """Compute slices for each label in a segmentation.
+
+    Args:
+        segmentation (ndarray): The segmentation array.
+
+    Returns:
+        dict[int, tuple[slice]]: A dictionary with the label as key
+            and the slice as value.
+    """
+    slices = ndi.find_objects(segmentation)
+    slices_dict = {}
+    for label, s in enumerate(slices, start=1):
+        if s is None:
+            continue
+        else:
+            slices_dict[label] = s
+    return slices_dict
+
+
+def lazy_compute_slices(segmentation: da.Array) -> dict[int, tuple[slice, ...]]:
+    """Compute slices for each label in a segmentation."""
+    global_offsets = _compute_offsets(segmentation.chunks)
+    delayed_chunks = segmentation.to_delayed()
+
+    grid_shape = tuple(len(c) for c in segmentation.chunks)
+
+    grid_indices = list(itertools.product(*[range(n) for n in grid_shape]))
+    delayed_results = []
+    for idx, chunk in zip(grid_indices, np.ravel(delayed_chunks), strict=True):
+        offset = tuple(global_offsets[dim][idx[dim]] for dim in range(len(idx)))
+        delayed_result = _process_chunk(chunk, offset)
+        delayed_results.append(delayed_result)
+
+    return _collect_slices(delayed_results).compute()
+
+
+def compute_masking_roi(
+    segmentation: np.ndarray | da.Array, pixel_size: PixelSize
+) -> list[Roi]:
+    """Compute a ROIs for each label in a segmentation.
+
+    This function expects a 2D or 3D segmentation array.
+    And this function expects the axes order to be 'zyx' or 'yx'.
+    Other axes orders are not supported.
+
+    """
+    if segmentation.ndim not in [2, 3]:
+        raise NgioValueError("Only 2D and 3D segmentations are supported.")
+
+    if isinstance(segmentation, da.Array):
+        slices = lazy_compute_slices(segmentation)
+    else:
+        slices = compute_slices(segmentation)
+
+    rois = []
+    for label, slice_ in slices.items():
+        if len(slice_) == 2:
+            min_z, min_y, min_x = 0, slice_[0].start, slice_[1].start
+            max_z, max_y, max_x = 1, slice_[0].stop, slice_[1].stop
+        elif len(slice_) == 3:
+            min_z, min_y, min_x = slice_[0].start, slice_[1].start, slice_[2].start
+            max_z, max_y, max_x = slice_[0].stop, slice_[1].stop, slice_[2].stop
+        else:
+            raise ValueError("Invalid slice length.")
+        roi = RoiPixels(
+            name=str(label),
+            x_length=max_x - min_x,
+            y_length=max_y - min_y,
+            z_length=max_z - min_z,
+            x=min_x,
+            y=min_y,
+            z=min_z,
+        )
+
+        roi = roi.to_roi(pixel_size)
+        rois.append(roi)
+    return rois

ngio/common/_pyramid.py

@@ -0,0 +1,228 @@
+import math
+from collections.abc import Collection
+from typing import Literal
+
+import dask.array as da
+import numpy as np
+import zarr
+
+from ngio.common._zoom import _zoom_inputs_check, dask_zoom, numpy_zoom
+from ngio.utils import (
+    AccessModeLiteral,
+    NgioValueError,
+    StoreOrGroup,
+    open_group_wrapper,
+)
+
+
+def _on_disk_numpy_zoom(
+    source: zarr.Array,
+    target: zarr.Array,
+    order: Literal[0, 1, 2] = 1,
+) -> None:
+    target[...] = numpy_zoom(source[...], target_shape=target.shape, order=order)
+
+
+def _on_disk_dask_zoom(
+    source: zarr.Array,
+    target: zarr.Array,
+    order: Literal[0, 1, 2] = 1,
+) -> None:
+    source_array = da.from_zarr(source)
+    target_array = dask_zoom(source_array, target_shape=target.shape, order=order)
+
+    target_array = target_array.rechunk(target.chunks)
+    target_array.compute_chunk_sizes()
+    target_array.to_zarr(target)
+
+
+def _on_disk_coarsen(
+    source: zarr.Array,
+    target: zarr.Array,
+    _order: Literal[0, 1] = 1,
+    aggregation_function: np.ufunc | None = None,
+) -> None:
+    """Apply a coarsening operation from a source zarr array to a target zarr array.
+
+    Args:
+        source (zarr.Array): The source array to coarsen.
+        target (zarr.Array): The target array to save the coarsened result to.
+        _order (Literal[0, 1]): The order of interpolation is not really implemented
+            for coarsening, but it is kept for compatibility with the zoom function.
+            _order=1 -> linear interpolation ~ np.mean
+            _order=0 -> nearest interpolation ~ np.max
+        aggregation_function (np.ufunc): The aggregation function to use.
+    """
+    source_array = da.from_zarr(source)
+
+    _scale, _target_shape = _zoom_inputs_check(
+        source_array=source_array, scale=None, target_shape=target.shape
+    )
+
+    assert _target_shape == target.shape, (
+        "Target shape must match the target array shape"
+    )
+
+    if aggregation_function is None:
+        if _order == 1:
+            aggregation_function = np.mean
+        elif _order == 0:
+            aggregation_function = np.max
+        else:
+            raise NgioValueError(
+                f"Aggregation function must be provided for order {_order}"
+            )
+
+    coarsening_setup = {}
+    for i, s in enumerate(_scale):
+        factor = 1 / s
+        # This check is very strict, but it is necessary to avoid
+        # a few pixels shift in the coarsening
+        # We could add a tolerance
+        if factor.is_integer():
+            coarsening_setup[i] = int(factor)
+        else:
+            raise NgioValueError(
+                f"Coarsening factor must be an integer, got {factor} on axis {i}"
+            )
+
+    out_target = da.coarsen(
+        aggregation_function, source_array, coarsening_setup, trim_excess=True
+    )
+    out_target = out_target.rechunk(target.chunks)
+    out_target.to_zarr(target)
+
+
+def on_disk_zoom(
+    source: zarr.Array,
+    target: zarr.Array,
+    order: Literal[0, 1, 2] = 1,
+    mode: Literal["dask", "numpy", "coarsen"] = "dask",
+) -> None:
+    """Apply a zoom operation from a source zarr array to a target zarr array.
+
+    Args:
+        source (zarr.Array): The source array to zoom.
+        target (zarr.Array): The target array to save the zoomed result to.
+        order (Literal[0, 1, 2]): The order of interpolation. Defaults to 1.
+        mode (Literal["dask", "numpy", "coarsen"]): The mode to use. Defaults to "dask".
+    """
+    if not isinstance(source, zarr.Array):
+        raise NgioValueError("source must be a zarr array")
+
+    if not isinstance(target, zarr.Array):
+        raise NgioValueError("target must be a zarr array")
+
+    if source.dtype != target.dtype:
+        raise NgioValueError("source and target must have the same dtype")
+
+    match mode:
+        case "numpy":
+            return _on_disk_numpy_zoom(source, target, order)
+        case "dask":
+            return _on_disk_dask_zoom(source, target, order)
+        case "coarsen":
+            return _on_disk_coarsen(
+                source,
+                target,
+            )
+        case _:
+            raise NgioValueError("mode must be either 'dask', 'numpy' or 'coarsen'")
+
+
+def _find_closest_arrays(
+    processed: list[zarr.Array], to_be_processed: list[zarr.Array]
+) -> tuple[int, int]:
+    dist_matrix = np.zeros((len(processed), len(to_be_processed)))
+    for i, arr_to_proc in enumerate(to_be_processed):
+        for j, proc_arr in enumerate(processed):
+            dist_matrix[j, i] = np.sqrt(
+                np.sum(
+                    [
+                        (s1 - s2) ** 2
+                        for s1, s2 in zip(
+                            arr_to_proc.shape, proc_arr.shape, strict=False
+                        )
+                    ]
+                )
+            )
+
+    return np.unravel_index(dist_matrix.argmin(), dist_matrix.shape)
+
+
+def consolidate_pyramid(
+    source: zarr.Array,
+    targets: list[zarr.Array],
+    order: Literal[0, 1, 2] = 1,
+    mode: Literal["dask", "numpy", "coarsen"] = "dask",
+) -> None:
+    """Consolidate the Zarr array."""
+    processed = [source]
+    to_be_processed = targets
+
+    while to_be_processed:
+        source_id, target_id = _find_closest_arrays(processed, to_be_processed)
+
+        source_image = processed[source_id]
+        target_image = to_be_processed.pop(target_id)
+
+        on_disk_zoom(
+            source=source_image,
+            target=target_image,
+            mode=mode,
+            order=order,
+        )
+        processed.append(target_image)
+
+
+def init_empty_pyramid(
+    store: StoreOrGroup,
+    paths: list[str],
+    ref_shape: Collection[int],
+    scaling_factors: Collection[float],
+    chunks: Collection[int] | None = None,
+    dtype: str = "uint16",
+    mode: AccessModeLiteral = "a",
+) -> None:
+    # Return the an Image object
+    if chunks is not None and len(chunks) != len(ref_shape):
+        raise NgioValueError(
+            "The shape and chunks must have the same number of dimensions."
+        )
+
+    if len(ref_shape) != len(scaling_factors):
+        raise NgioValueError(
+            "The shape and scaling factor must have the same number of dimensions."
+        )
+
+    root_group, _ = open_group_wrapper(store, mode=mode)
+    for path in paths:
+        if any(s < 1 for s in ref_shape):
+            raise NgioValueError(
+                "Level shape must be at least 1 on all dimensions. "
+                f"Calculated shape: {ref_shape} at level {path}."
+            )
+        new_arr = root_group.zeros(
+            name=path,
+            shape=ref_shape,
+            dtype=dtype,
+            chunks=chunks,
+            dimension_separator="/",
+            overwrite=True,
+        )
+
+        # Todo redo this with when a proper build of pyramid is implemented
+        _shape = []
+        for s, sc in zip(ref_shape, scaling_factors, strict=True):
+            if math.floor(s / sc) % 2 == 0:
+                _shape.append(math.floor(s / sc))
+            else:
+                _shape.append(math.ceil(s / sc))
+        ref_shape = _shape
+
+        if chunks is None:
+            chunks = new_arr.chunks
+            if chunks is None:
+                raise NgioValueError("Something went wrong with the chunks")
+        chunks = [min(c, s) for c, s in zip(chunks, ref_shape, strict=True)]
+    return None

ngio/common/_roi.py

@@ -0,0 +1,165 @@
+"""Region of interest (ROI) metadata.
+
+These are the interfaces bwteen the ROI tables / masking ROI tables and
+    the ImageLikeHandler.
+"""
+
+from collections.abc import Iterable
+
+import numpy as np
+from pydantic import BaseModel, ConfigDict, Field
+
+from ngio.common._dimensions import Dimensions
+from ngio.ome_zarr_meta.ngio_specs import PixelSize, SpaceUnits
+from ngio.utils import NgioValueError
+
+
+def _to_raster(value: float, pixel_size: float, max_shape: int) -> int:
+    """Convert to raster coordinates."""
+    round_value = int(np.round(value / pixel_size))
+    # Ensure the value is within the image shape boundaries
+    return max(0, min(round_value, max_shape))
+
+
+def _to_world(value: int, pixel_size: float) -> float:
+    """Convert to world coordinates."""
+    return value * pixel_size
+
+
+class Roi(BaseModel):
+    """Region of interest (ROI) metadata."""
+
+    name: str
+    x_length: float
+    y_length: float
+    z_length: float = 1.0
+    x: float = 0.0
+    y: float = 0.0
+    z: float = 0.0
+    unit: SpaceUnits = Field(SpaceUnits.micrometer, repr=False)
+
+    model_config = ConfigDict(extra="allow")
+
+    def to_pixel_roi(
+        self, pixel_size: PixelSize, dimensions: Dimensions
+    ) -> "RoiPixels":
+        """Convert to raster coordinates."""
+        dim_x = dimensions.get("x")
+        dim_y = dimensions.get("y")
+        # Will default to 1 if z does not exist
+        dim_z = dimensions.get("z", strict=False)
+
+        return RoiPixels(
+            name=self.name,
+            x=_to_raster(self.x, pixel_size.x, dim_x),
+            y=_to_raster(self.y, pixel_size.y, dim_y),
+            z=_to_raster(self.z, pixel_size.z, dim_z),
+            x_length=_to_raster(self.x_length, pixel_size.x, dim_x),
+            y_length=_to_raster(self.y_length, pixel_size.y, dim_y),
+            z_length=_to_raster(self.z_length, pixel_size.z, dim_z),
+        )
+
+    def zoom(self, zoom_factor: float = 1) -> "Roi":
+        """Zoom the ROI by a factor.
+
+        Args:
+            zoom_factor: The zoom factor. If the zoom factor
+                is less than 1 the ROI will be zoomed in.
+                If the zoom factor is greater than 1 the ROI will be zoomed out.
+                If the zoom factor is 1 the ROI will not be changed.
+        """
+        return zoom_roi(self, zoom_factor)
+
+
+class RoiPixels(BaseModel):
+    """Region of interest (ROI) metadata."""
+
+    name: str
+    x: int
+    y: int
+    z: int
+    x_length: int
+    y_length: int
+    z_length: int
+    model_config = ConfigDict(extra="allow")
+
+    def to_roi(self, pixel_size: PixelSize) -> Roi:
+        """Convert to world coordinates."""
+        return Roi(
+            name=self.name,
+            x=_to_world(self.x, pixel_size.x),
+            y=_to_world(self.y, pixel_size.y),
+            z=_to_world(self.z, pixel_size.z),
+            x_length=_to_world(self.x_length, pixel_size.x),
+            y_length=_to_world(self.y_length, pixel_size.y),
+            z_length=_to_world(self.z_length, pixel_size.z),
+            unit=pixel_size.space_unit,
+        )
+
+    def to_slices(self) -> dict[str, slice]:
+        """Return the slices for the ROI."""
+        return {
+            "x": slice(self.x, self.x + self.x_length),
+            "y": slice(self.y, self.y + self.y_length),
+            "z": slice(self.z, self.z + self.z_length),
+        }
+
+
+def zoom_roi(roi: Roi, zoom_factor: float = 1) -> Roi:
+    """Zoom the ROI by a factor.
+
+    Args:
+        roi: The ROI to zoom.
+        zoom_factor: The zoom factor. If the zoom factor
+            is less than 1 the ROI will be zoomed in.
+            If the zoom factor is greater than 1 the ROI will be zoomed out.
+            If the zoom factor is 1 the ROI will not be changed.
+    """
+    if zoom_factor <= 0:
+        raise ValueError("Zoom factor must be greater than 0.")
+
+    # the zoom factor needs to be rescaled
+    # from the range [-1, inf) to [0, inf)
+    zoom_factor -= 1
+    diff_x = roi.x_length * zoom_factor
+    diff_y = roi.y_length * zoom_factor
+
+    new_x = max(roi.x - diff_x / 2, 0)
+    new_y = max(roi.y - diff_y / 2, 0)
+
+    new_roi = Roi(
+        name=roi.name,
+        x=new_x,
+        y=new_y,
+        z=roi.z,
+        x_length=roi.x_length + diff_x,
+        y_length=roi.y_length + diff_y,
+        z_length=roi.z_length,
+        unit=roi.unit,
+    )
+
+    return new_roi
+
+
+def roi_to_slice_kwargs(
+    roi: Roi,
+    pixel_size: PixelSize,
+    dimensions: Dimensions,
+    **slice_kwargs: slice | int | Iterable[int],
+) -> dict[str, slice | int | Iterable[int]]:
+    """Convert a WorldCooROI to slice_kwargs."""
+    raster_roi = roi.to_pixel_roi(
+        pixel_size=pixel_size, dimensions=dimensions
+    ).to_slices()
+
+    if not dimensions.has_axis(axis_name="z"):
+        raster_roi.pop("z")
+
+    for key in slice_kwargs.keys():
+        if key in raster_roi:
+            raise NgioValueError(
+                f"Key {key} is already in the slice_kwargs. "
+                "Ambiguous which one to use: "
+                f"{key}={slice_kwargs[key]} or roi_{key}={raster_roi[key]}"
+            )
+    return {**raster_roi, **slice_kwargs}

ngio/common/_slicer.py

@@ -0,0 +1,96 @@
+from collections.abc import Iterable
+
+import dask.array as da
+import numpy as np
+import zarr
+
+from ngio.common._dimensions import Dimensions
+from ngio.ome_zarr_meta.ngio_specs import AxesTransformation
+from ngio.utils import NgioValueError
+
+
+def _validate_int(value: int, shape: int) -> int:
+    if not isinstance(value, int):
+        raise NgioValueError(f"Invalid value {value} of type {type(value)}")
+    if value < 0 or value >= shape:
+        raise NgioValueError(
+            f"Invalid value {value}. Index out of bounds for axis of shape {shape}"
+        )
+    return value
+
+
+def _validate_iter_of_ints(value: Iterable[int], shape: int) -> list[int]:
+    if not isinstance(value, list):
+        raise NgioValueError(f"Invalid value {value} of type {type(value)}")
+    value = [_validate_int(v, shape=shape) for v in value]
+    return value
+
+
+def _validate_slice(value: slice, shape: int) -> slice:
+    start = value.start if value.start is not None else 0
+    start = max(start, 0)
+    stop = value.stop if value.stop is not None else shape
+    return slice(start, stop)
+
+
+class SliceTransform(AxesTransformation):
+    slices: tuple[slice | tuple[int, ...], ...]
+
+
+def compute_and_slices(
+    *,
+    dimensions: Dimensions,
+    **slice_kwargs: slice | int | Iterable[int],
+) -> SliceTransform:
+    _slices = {}
+    axes_names = dimensions._axes_mapper.on_disk_axes_names
+    for axis_name, slice_ in slice_kwargs.items():
+        axis = dimensions._axes_mapper.get_axis(axis_name)
+        if axis is None:
+            raise NgioValueError(
+                f"Invalid axis {axis_name}. "
+                f"Not found on the on-disk axes {axes_names}. "
+                "If you want to get/set a singletorn value include "
+                "it in the axes_order parameter."
+            )
+
+        shape = dimensions.get(axis.on_disk_name)
+
+        if isinstance(slice_, int):
+            slice_ = _validate_int(slice_, shape)
+            slice_ = slice(slice_, slice_ + 1)
+
+        elif isinstance(slice_, Iterable):
+            slice_ = _validate_iter_of_ints(slice_, shape)
+            slice_ = tuple(slice_)
+
+        elif isinstance(slice_, slice):
+            slice_ = _validate_slice(slice_, shape)
+
+        elif not isinstance(slice_, slice):
+            raise NgioValueError(
+                f"Invalid slice definition {slice_} of type {type(slice_)}"
+            )
+        _slices[axis.on_disk_name] = slice_
+
+    slices = tuple(_slices.get(axis, slice(None)) for axis in axes_names)
+    return SliceTransform(slices=slices)
+
+
+def numpy_get_slice(array: zarr.Array, slices: SliceTransform) -> np.ndarray:
+    return array[slices.slices]
+
+
+def dask_get_slice(array: zarr.Array, slices: SliceTransform) -> da.Array:
+    da_array = da.from_zarr(array)
+    return da_array[slices.slices]
+
+
+def numpy_set_slice(
+    array: zarr.Array, patch: np.ndarray, slices: SliceTransform
+) -> None:
+    array[slices.slices] = patch
+
+
+def dask_set_slice(array: zarr.Array, patch: da.Array, slices: SliceTransform) -> None:
+    da.to_zarr(arr=patch, url=array, region=slices.slices)