ngio 0.2.0a2__py3-none-any.whl → 0.2.0b1__py3-none-any.whl

ngio/__init__.py

@@ -10,7 +10,7 @@ __author__ = "Lorenzo Cerrone"
 __email__ = "lorenzo.cerrone@uzh.ch"
 
 from ngio.common import ArrayLike, Dimensions
-from ngio.hcs import OmeZarrPlate, OmeZarrWell, open_omezarr_plate, open_omezarr_well
+from ngio.hcs import OmeZarrPlate, create_empty_plate, open_omezarr_plate
 from ngio.images import (
     Image,
     Label,
@@ -20,22 +20,22 @@ from ngio.images import (
     open_image,
     open_omezarr_container,
 )
-from ngio.ome_zarr_meta.ngio_specs import AxesSetup, PixelSize
+from ngio.ome_zarr_meta.ngio_specs import AxesSetup, ImageInWellPath, PixelSize
 
 __all__ = [
     "ArrayLike",
     "AxesSetup",
     "Dimensions",
     "Image",
+    "ImageInWellPath",
     "Label",
     "OmeZarrContainer",
     "OmeZarrPlate",
-    "OmeZarrWell",
     "PixelSize",
     "create_empty_omezarr",
+    "create_empty_plate",
     "create_omezarr_from_array",
     "open_image",
     "open_omezarr_container",
     "open_omezarr_plate",
-    "open_omezarr_well",
 ]

ngio/common/__init__.py

@@ -1,6 +1,11 @@
 """Common classes and functions that are used across the package."""
 
-from ngio.common._array_pipe import get_pipe, set_pipe
+from ngio.common._array_pipe import (
+    get_masked_pipe,
+    get_pipe,
+    set_masked_pipe,
+    set_pipe,
+)
 from ngio.common._axes_transforms import (
     transform_dask_array,
     transform_list,
@@ -8,8 +13,9 @@ from ngio.common._axes_transforms import (
 )
 from ngio.common._common_types import ArrayLike
 from ngio.common._dimensions import Dimensions
+from ngio.common._masking_roi import compute_masking_roi
 from ngio.common._pyramid import consolidate_pyramid, init_empty_pyramid, on_disk_zoom
-from ngio.common._roi import RasterCooROI, WorldCooROI
+from ngio.common._roi import RasterCooROI, WorldCooROI, roi_to_slice_kwargs
 from ngio.common._slicer import (
     SliceTransform,
     compute_and_slices,
@@ -27,16 +33,20 @@ __all__ = [
     "SliceTransform",
     "WorldCooROI",
     "compute_and_slices",
+    "compute_masking_roi",
     "consolidate_pyramid",
     "dask_get_slice",
     "dask_set_slice",
     "dask_zoom",
+    "get_masked_pipe",
     "get_pipe",
     "init_empty_pyramid",
     "numpy_get_slice",
     "numpy_set_slice",
     "numpy_zoom",
     "on_disk_zoom",
+    "roi_to_slice_kwargs",
+    "set_masked_pipe",
     "set_pipe",
     "transform_dask_array",
     "transform_list",

ngio/common/_array_pipe.py

@@ -158,3 +158,109 @@ def set_pipe(
         )
     else:
         raise NgioValueError("Unknown patch type, expected numpy, dask or delayed.")
+
+
+def _mask_pipe_common(
+    array: zarr.Array,
+    label_array: zarr.Array,
+    label: int,
+    *,
+    dimensions_array: Dimensions,
+    dimensions_label: Dimensions,
+    axes_order: Collection[str] | None = None,
+    mode: Literal["numpy", "dask", "delayed"] = "numpy",
+    **slice_kwargs: slice | int | Iterable[int],
+):
+    array_patch = get_pipe(
+        array,
+        dimensions=dimensions_array,
+        axes_order=axes_order,
+        mode=mode,
+        **slice_kwargs,
+    )
+
+    if "c" in slice_kwargs.keys():
+        # This makes the strong assumption that the
+        # user is passing the channel axis as "c"
+        # This will fail if the channel axis is queried
+        # with a different on-disk name
+        slice_kwargs.pop("c")
+
+    label_patch = get_pipe(
+        label_array,
+        dimensions=dimensions_label,
+        axes_order=axes_order,
+        mode=mode,
+        **slice_kwargs,
+    )
+
+    if isinstance(array_patch, np.ndarray):
+        label_patch = np.broadcast_to(label_patch, array_patch.shape)
+    elif isinstance(array_patch, dask.array.Array):
+        label_patch = dask.array.broadcast_to(label_patch, array_patch.shape)
+    else:
+        raise NgioValueError(f"Mode {mode} not yet supported for masked array.")
+
+    mask = label_patch == label
+    return array_patch, mask
+
+
+def get_masked_pipe(
+    array: zarr.Array,
+    label_array: zarr.Array,
+    label: int,
+    *,
+    dimensions_array: Dimensions,
+    dimensions_label: Dimensions,
+    axes_order: Collection[str] | None = None,
+    mode: Literal["numpy", "dask", "delayed"] = "numpy",
+    **slice_kwargs: slice | int | Iterable[int],
+):
+    array_patch, mask = _mask_pipe_common(
+        array=array,
+        label_array=label_array,
+        label=label,
+        dimensions_array=dimensions_array,
+        dimensions_label=dimensions_label,
+        axes_order=axes_order,
+        mode=mode,
+        **slice_kwargs,
+    )
+    array_patch[~mask] = 0
+    return array_patch
+
+
+def set_masked_pipe(
+    array: zarr.Array,
+    label_array: zarr.Array,
+    label: int,
+    patch: ArrayLike,
+    *,
+    dimensions_array: Dimensions,
+    dimensions_label: Dimensions,
+    axes_order: Collection[str] | None = None,
+    **slice_kwargs: slice | int | Iterable[int],
+):
+    if isinstance(patch, dask.array.Array):
+        mode = "dask"
+    elif isinstance(patch, np.ndarray):
+        mode = "numpy"
+    else:
+        raise NgioValueError(
+            "Mode not yet supported for masked array. Expected a numpy or dask array."
+        )
+
+    array_patch, mask = _mask_pipe_common(
+        array=array,
+        label_array=label_array,
+        label=label,
+        dimensions_array=dimensions_array,
+        dimensions_label=dimensions_label,
+        axes_order=axes_order,
+        mode=mode,
+        **slice_kwargs,
+    )
+    patch = np.where(mask, patch, array_patch)
+    set_pipe(
+        array, patch, dimensions=dimensions_array, axes_order=axes_order, **slice_kwargs
+    )

ngio/common/_axes_transforms.py

@@ -9,6 +9,7 @@ from ngio.ome_zarr_meta.ngio_specs._axes import (
     AxesTransformation,
     AxesTranspose,
 )
+from ngio.utils import NgioValueError
 
 T = TypeVar("T")
 
@@ -44,7 +45,7 @@ def transform_numpy_array(
         elif isinstance(operation, AxesSqueeze):
             array = np.squeeze(array, axis=operation.axes)
         else:
-            raise ValueError(f"Unknown operation {operation}")
+            raise NgioValueError(f"Unknown operation {operation}")
     return array
 
 
@@ -59,5 +60,5 @@ def transform_dask_array(
         elif isinstance(operation, AxesSqueeze):
             array = da.squeeze(array, axis=operation.axes)
         else:
-            raise ValueError(f"Unknown operation {operation}")
+            raise NgioValueError(f"Unknown operation {operation}")
     return array

ngio/common/_dimensions.py

@@ -57,6 +57,13 @@ class Dimensions:
             return 1
         return self._shape[index]
 
+    def has_axis(self, axis_name: str) -> bool:
+        """Return whether the axis exists."""
+        index = self._axes_mapper.get_axis(axis_name)
+        if index is None:
+            return False
+        return True
+
     def get_shape(self, axes_order: Collection[str]) -> tuple[int, ...]:
         """Return the shape in the given axes order."""
         transforms = self._axes_mapper.to_order(axes_order)

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 RasterCooROI, WorldCooROI
+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[WorldCooROI]:
+    """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 = RasterCooROI(
+            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_world_coo_roi(pixel_size)
+        rois.append(roi)
+    return rois

ngio/common/_pyramid.py

@@ -7,7 +7,12 @@ import numpy as np
 import zarr
 
 from ngio.common._zoom import _zoom_inputs_check, dask_zoom, numpy_zoom
-from ngio.utils import AccessModeLiteral, StoreOrGroup, open_group_wrapper
+from ngio.utils import (
+    AccessModeLiteral,
+    NgioValueError,
+    StoreOrGroup,
+    open_group_wrapper,
+)
 
 
 def _on_disk_numpy_zoom(
@@ -64,7 +69,7 @@ def _on_disk_coarsen(
         elif _order == 0:
             aggregation_function = np.max
         else:
-            raise ValueError(
+            raise NgioValueError(
                 f"Aggregation function must be provided for order {_order}"
             )
 
@@ -77,7 +82,7 @@ def _on_disk_coarsen(
         if factor.is_integer():
             coarsening_setup[i] = int(factor)
         else:
-            raise ValueError(
+            raise NgioValueError(
                 f"Coarsening factor must be an integer, got {factor} on axis {i}"
             )
 
@@ -103,13 +108,13 @@ def on_disk_zoom(
         mode (Literal["dask", "numpy", "coarsen"]): The mode to use. Defaults to "dask".
     """
     if not isinstance(source, zarr.Array):
-        raise ValueError("source must be a zarr array")
+        raise NgioValueError("source must be a zarr array")
 
     if not isinstance(target, zarr.Array):
-        raise ValueError("target must be a zarr array")
+        raise NgioValueError("target must be a zarr array")
 
     if source.dtype != target.dtype:
-        raise ValueError("source and target must have the same dtype")
+        raise NgioValueError("source and target must have the same dtype")
 
     match mode:
         case "numpy":
@@ -122,7 +127,7 @@ def on_disk_zoom(
                 target,
             )
         case _:
-            raise ValueError("mode must be either 'dask', 'numpy' or 'coarsen'")
+            raise NgioValueError("mode must be either 'dask', 'numpy' or 'coarsen'")
 
 
 def _find_closest_arrays(
@@ -181,19 +186,19 @@ def init_empty_pyramid(
 ) -> None:
     # Return the an Image object
     if chunks is not None and len(chunks) != len(ref_shape):
-        raise ValueError(
+        raise NgioValueError(
             "The shape and chunks must have the same number of dimensions."
         )
 
     if len(ref_shape) != len(scaling_factors):
-        raise ValueError(
+        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 ValueError(
+            raise NgioValueError(
                 "Level shape must be at least 1 on all dimensions. "
                 f"Calculated shape: {ref_shape} at level {path}."
             )
@@ -218,6 +223,6 @@ def init_empty_pyramid(
         if chunks is None:
             chunks = new_arr.chunks
             if chunks is None:
-                raise ValueError("Something went wrong with the chunks")
+                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

@@ -4,11 +4,14 @@ 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:
@@ -56,6 +59,17 @@ class WorldCooROI(BaseModel):
             z_length=_to_raster(self.z_length, pixel_size.z, dim_z),
         )
 
+    def zoom(self, zoom_factor: float = 1) -> "WorldCooROI":
+        """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 RasterCooROI(BaseModel):
     """Region of interest (ROI) metadata."""
@@ -89,3 +103,63 @@ class RasterCooROI(BaseModel):
             "y": slice(self.y, self.y + self.y_length),
             "z": slice(self.z, self.z + self.z_length),
         }
+
+
+def zoom_roi(roi: WorldCooROI, zoom_factor: float = 1) -> WorldCooROI:
+    """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 = WorldCooROI(
+        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: WorldCooROI,
+    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_raster_coo(
+        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

@@ -1,4 +1,3 @@
-# %%
 from collections.abc import Iterable
 
 import dask.array as da
@@ -69,7 +68,7 @@ def compute_and_slices(
             slice_ = _validate_slice(slice_, shape)
 
         elif not isinstance(slice_, slice):
-            raise ValueError(
+            raise NgioValueError(
                 f"Invalid slice definition {slice_} of type {type(slice_)}"
             )
         _slices[axis.on_disk_name] = slice_

ngio/common/_zoom.py

@@ -5,6 +5,8 @@ import dask.array as da
 import numpy as np
 from scipy.ndimage import zoom as scipy_zoom
 
+from ngio.utils import NgioValueError
+
 
 def _stacked_zoom(x, zoom_y, zoom_x, order=1, mode="grid-constant", grid_mode=True):
     *rest, yshape, xshape = x.shape
@@ -53,15 +55,15 @@ def _zoom_inputs_check(
     target_shape: tuple[int, ...] | None = None,
 ) -> tuple[np.ndarray, tuple[int, ...]]:
     if scale is None and target_shape is None:
-        raise ValueError("Either scale or target_shape must be provided")
+        raise NgioValueError("Either scale or target_shape must be provided")
 
     if scale is not None and target_shape is not None:
-        raise ValueError("Only one of scale or target_shape must be provided")
+        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 ValueError(
+            raise NgioValueError(
                 "Target shape must have the "
                 "same number of dimensions as "
                 "the source array"

ngio/hcs/__init__.py

@@ -1,60 +1,5 @@
 """OME-Zarr HCS objects models."""
 
-from typing import TYPE_CHECKING
+from ngio.hcs.plate import OmeZarrPlate, create_empty_plate, open_omezarr_plate
 
-if TYPE_CHECKING:
-    from ngio.images import OmeZarrContainer
-
-
-class OmeZarrPlate:
-    """Placeholder for the OME-Zarr image object."""
-
-    def __init__(self, *args, **kwargs):
-        """Initialize the OME-Zarr plate."""
-        raise NotImplementedError
-
-    def wells(self) -> list[str]:
-        """Return the wells."""
-        raise NotImplementedError
-
-    def columns(self) -> list[str]:
-        """Return the number of columns."""
-        raise NotImplementedError
-
-    def rows(self) -> list[str]:
-        """Return the number of rows."""
-        raise NotImplementedError
-
-    def get_omezarr_well(self, *args, **kwargs) -> "OmeZarrWell":
-        """Return the OME-Zarr well."""
-        raise NotImplementedError
-
-    def get_omezarr_image(self, *args, **kwargs) -> "OmeZarrContainer":
-        """Return the OME-Zarr image."""
-        raise NotImplementedError
-
-
-class OmeZarrWell:
-    """Placeholder for the Image object."""
-
-    def __init__(self, *args, **kwargs):
-        """Initialize the OME-Zarr well."""
-        raise NotImplementedError
-
-    def acquisitions(self) -> list[str]:
-        """Return the acquisition."""
-        raise NotImplementedError
-
-    def get_ome_zarr_image(self, *args, **kwargs) -> "OmeZarrContainer":
-        """Return the OME-Zarr image."""
-        raise NotImplementedError
-
-
-def open_omezarr_plate(*args, **kwargs):
-    """Open an OME-Zarr plate."""
-    return OmeZarrPlate(*args, **kwargs)
-
-
-def open_omezarr_well(*args, **kwargs):
-    """Open an OME-Zarr well."""
-    return OmeZarrWell(*args, **kwargs)
+__all__ = ["OmeZarrPlate", "create_empty_plate", "open_omezarr_plate"]