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

ngio/__init__.py

@@ -2,14 +2,42 @@
 
 from importlib.metadata import PackageNotFoundError, version
 
-from ngio.core import Image, Label, NgffImage
-
-__all__ = ["Image", "Label", "NgffImage"]
-
-
 try:
     __version__ = version("ngio")
 except PackageNotFoundError:  # pragma: no cover
     __version__ = "uninstalled"
 __author__ = "Lorenzo Cerrone"
 __email__ = "lorenzo.cerrone@uzh.ch"
+
+from ngio.common import ArrayLike, Dimensions, Roi, RoiPixels
+from ngio.hcs import OmeZarrPlate, create_empty_plate, open_ome_zarr_plate
+from ngio.images import (
+    Image,
+    Label,
+    OmeZarrContainer,
+    create_empty_ome_zarr,
+    create_ome_zarr_from_array,
+    open_image,
+    open_ome_zarr_container,
+)
+from ngio.ome_zarr_meta.ngio_specs import AxesSetup, ImageInWellPath, PixelSize
+
+__all__ = [
+    "ArrayLike",
+    "AxesSetup",
+    "Dimensions",
+    "Image",
+    "ImageInWellPath",
+    "Label",
+    "OmeZarrContainer",
+    "OmeZarrPlate",
+    "PixelSize",
+    "Roi",
+    "RoiPixels",
+    "create_empty_ome_zarr",
+    "create_empty_plate",
+    "create_ome_zarr_from_array",
+    "open_image",
+    "open_ome_zarr_container",
+    "open_ome_zarr_plate",
+]

ngio/common/__init__.py

@@ -0,0 +1,54 @@
+"""Common classes and functions that are used across the package."""
+
+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,
+    transform_numpy_array,
+)
+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 Roi, RoiPixels, roi_to_slice_kwargs
+from ngio.common._slicer import (
+    SliceTransform,
+    compute_and_slices,
+    dask_get_slice,
+    dask_set_slice,
+    numpy_get_slice,
+    numpy_set_slice,
+)
+from ngio.common._zoom import dask_zoom, numpy_zoom
+
+__all__ = [
+    "ArrayLike",
+    "Dimensions",
+    "Roi",
+    "RoiPixels",
+    "SliceTransform",
+    "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",
+    "transform_numpy_array",
+]

ngio/common/_array_pipe.py

@@ -0,0 +1,265 @@
+from collections.abc import Collection, Iterable
+from typing import Literal
+
+import dask
+import dask.delayed
+import numpy as np
+import zarr
+
+from ngio.common._axes_transforms import transform_dask_array, transform_numpy_array
+from ngio.common._common_types import ArrayLike
+from ngio.common._dimensions import Dimensions
+from ngio.common._slicer import (
+    SliceTransform,
+    compute_and_slices,
+    dask_get_slice,
+    dask_set_slice,
+    numpy_get_slice,
+    numpy_set_slice,
+)
+from ngio.ome_zarr_meta.ngio_specs import AxesTransformation
+from ngio.utils import NgioValueError
+
+
+def _compute_from_disk_transforms(
+    *,
+    dimensions: Dimensions,
+    axes_order: Collection[str] | None = None,
+    **slice_kwargs: slice | int | Iterable[int],
+) -> tuple[SliceTransform, tuple[AxesTransformation, ...]]:
+    slices = compute_and_slices(dimensions=dimensions, **slice_kwargs)
+
+    if axes_order is None:
+        return slices, ()
+
+    additional_transformations = dimensions._axes_mapper.to_order(axes_order)
+    return slices, additional_transformations
+
+
+def _compute_to_disk_transforms(
+    *,
+    dimensions: Dimensions,
+    axes_order: Collection[str] | None = None,
+    **slice_kwargs: slice | int | Iterable[int],
+) -> tuple[SliceTransform, tuple[AxesTransformation, ...]]:
+    slices = compute_and_slices(dimensions=dimensions, **slice_kwargs)
+    if axes_order is None:
+        return slices, ()
+
+    additional_transformations = dimensions._axes_mapper.from_order(axes_order)
+    return slices, additional_transformations
+
+
+def _numpy_get_pipe(
+    array: zarr.Array,
+    slices: SliceTransform,
+    transformations: tuple[AxesTransformation, ...],
+) -> np.ndarray:
+    array = numpy_get_slice(array, slices)
+    return transform_numpy_array(array, transformations)
+
+
+def _delayed_numpy_get_pipe(
+    array: zarr.Array,
+    slices: SliceTransform,
+    transformations: tuple[AxesTransformation, ...],
+) -> dask.delayed:
+    array = dask.delayed(numpy_get_slice)(array, slices)
+    return dask.delayed(transform_numpy_array)(array, transformations)
+
+
+def _dask_get_pipe(
+    array: zarr.Array,
+    slices: SliceTransform,
+    transformations: tuple[AxesTransformation, ...],
+) -> dask.array:
+    array = dask_get_slice(array, slices)
+    return transform_dask_array(array, transformations)
+
+
+def _numpy_set_pipe(
+    array: zarr.Array,
+    patch: np.ndarray,
+    slices: SliceTransform,
+    transformations: tuple[AxesTransformation, ...],
+) -> None:
+    patch = transform_numpy_array(patch, transformations)
+    numpy_set_slice(array, patch, slices)
+
+
+def _dask_set_pipe(
+    array: zarr.Array,
+    patch: np.ndarray,
+    slices: SliceTransform,
+    transformations: tuple[AxesTransformation, ...],
+) -> None:
+    patch = transform_dask_array(patch, transformations)
+    dask_set_slice(array, patch, slices)
+
+
+def _delayed_numpy_set_pipe(
+    array: zarr.Array,
+    patch: np.ndarray,
+    slices: SliceTransform,
+    transformations: tuple[AxesTransformation, ...],
+) -> dask.delayed:
+    patch = dask.delayed(transform_numpy_array)(patch, transformations)
+    return dask.delayed(numpy_set_slice)(array, patch, slices)
+
+
+def get_pipe(
+    array: zarr.Array,
+    *,
+    dimensions: Dimensions,
+    axes_order: Collection[str] | None = None,
+    mode: Literal["numpy", "dask", "delayed"] = "numpy",
+    **slice_kwargs: slice | int | Iterable[int],
+):
+    slices, transformations = _compute_from_disk_transforms(
+        dimensions=dimensions, axes_order=axes_order, **slice_kwargs
+    )
+    match mode:
+        case "numpy":
+            return _numpy_get_pipe(array, slices, transformations)
+        case "dask":
+            return _dask_get_pipe(array, slices, transformations)
+
+        case "delayed":
+            return _delayed_numpy_get_pipe(array, slices, transformations)
+
+        case _:
+            raise NgioValueError(
+                f"Unknown get pipe mode {mode}, expected 'numpy', 'dask' or 'delayed'."
+            )
+
+
+def set_pipe(
+    array: zarr.Array,
+    patch: ArrayLike,
+    *,
+    dimensions: Dimensions,
+    axes_order: Collection[str] | None = None,
+    **slice_kwargs: slice | int | Iterable[int],
+):
+    slices, transformations = _compute_to_disk_transforms(
+        dimensions=dimensions, axes_order=axes_order, **slice_kwargs
+    )
+    if isinstance(patch, dask.array.Array):
+        _dask_set_pipe(
+            array=array, patch=patch, slices=slices, transformations=transformations
+        )
+    elif isinstance(patch, np.ndarray):
+        _numpy_set_pipe(
+            array=array, patch=patch, slices=slices, transformations=transformations
+        )
+    elif isinstance(patch, dask.delayed.Delayed):
+        _delayed_numpy_set_pipe(
+            array=array, patch=patch, slices=slices, transformations=transformations
+        )
+    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 not dimensions_label.has_axis("c"):
+        # Remove the 'c' from the slice_kwargs
+        # This will not work if the query uses non-default
+        # axes names for channel
+        slice_kwargs = {k: v for k, v in slice_kwargs.items() if k != "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

@@ -0,0 +1,64 @@
+from typing import TypeVar
+
+import dask.array as da
+import numpy as np
+
+from ngio.ome_zarr_meta.ngio_specs._axes import (
+    AxesExpand,
+    AxesSqueeze,
+    AxesTransformation,
+    AxesTranspose,
+)
+from ngio.utils import NgioValueError
+
+T = TypeVar("T")
+
+
+def transform_list(
+    input_list: list[T], default: T, operations: tuple[AxesTransformation, ...]
+) -> list[T]:
+    if isinstance(input_list, tuple):
+        input_list = list(input_list)
+
+    for operation in operations:
+        if isinstance(operation, AxesTranspose):
+            input_list = [input_list[i] for i in operation.axes]
+
+        if isinstance(operation, AxesExpand):
+            for ax in operation.axes:
+                input_list.insert(ax, default)
+        elif isinstance(operation, AxesSqueeze):
+            for offset, ax in enumerate(operation.axes):
+                input_list.pop(ax - offset)
+
+    return input_list
+
+
+def transform_numpy_array(
+    array: np.ndarray, operations: tuple[AxesTransformation, ...]
+) -> np.ndarray:
+    for operation in operations:
+        if isinstance(operation, AxesTranspose):
+            array = np.transpose(array, operation.axes)
+        elif isinstance(operation, AxesExpand):
+            array = np.expand_dims(array, axis=operation.axes)
+        elif isinstance(operation, AxesSqueeze):
+            array = np.squeeze(array, axis=operation.axes)
+        else:
+            raise NgioValueError(f"Unknown operation {operation}")
+    return array
+
+
+def transform_dask_array(
+    array: da.Array, operations: tuple[AxesTransformation, ...]
+) -> da.Array:
+    for operation in operations:
+        if isinstance(operation, AxesTranspose):
+            array = da.transpose(array, axes=operation.axes)
+        elif isinstance(operation, AxesExpand):
+            array = da.expand_dims(array, axis=operation.axes)
+        elif isinstance(operation, AxesSqueeze):
+            array = da.squeeze(array, axis=operation.axes)
+        else:
+            raise NgioValueError(f"Unknown operation {operation}")
+    return array

ngio/common/_common_types.py

@@ -0,0 +1,5 @@
+import numpy as np
+import zarr
+from dask import array as da
+
+ArrayLike = np.ndarray | da.core.Array | zarr.Array  # type: ignore

ngio/common/_dimensions.py

@@ -0,0 +1,120 @@
+"""Dimension metadata.
+
+This is not related to the NGFF metadata,
+but it is based on the actual metadata of the image data.
+"""
+
+from collections.abc import Collection
+
+from ngio.common._axes_transforms import transform_list
+from ngio.ome_zarr_meta import AxesMapper
+from ngio.utils import NgioValidationError, NgioValueError
+
+
+class Dimensions:
+    """Dimension metadata."""
+
+    def __init__(
+        self,
+        shape: tuple[int, ...],
+        axes_mapper: AxesMapper,
+    ) -> None:
+        """Create a Dimension object from a Zarr array.
+
+        Args:
+            shape: The shape of the Zarr array.
+            axes_mapper: The axes mapper object.
+        """
+        self._shape = shape
+        self._axes_mapper = axes_mapper
+
+        if len(self._shape) != len(self._axes_mapper.on_disk_axes):
+            raise NgioValidationError(
+                "The number of dimensions must match the number of axes. "
+                f"Expected Axis {self._axes_mapper.on_disk_axes_names} but got shape "
+                f"{self._shape}."
+            )
+
+    def __str__(self) -> str:
+        """Return the string representation of the object."""
+        dims = ", ".join(
+            f"{ax.on_disk_name}: {s}"
+            for ax, s in zip(self._axes_mapper.on_disk_axes, self._shape, strict=True)
+        )
+        return f"Dimensions({dims})"
+
+    def get(self, axis_name: str, strict: bool = True) -> int:
+        """Return the dimension of the given axis name.
+
+        Args:
+            axis_name: The name of the axis (either canonical or non-canonical).
+            strict: If True, raise an error if the axis does not exist.
+        """
+        index = self._axes_mapper.get_index(axis_name)
+        if index is None and strict:
+            raise NgioValueError(f"Axis {axis_name} does not exist.")
+        elif index is None:
+            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)
+        return tuple(transform_list(list(self._shape), 1, transforms))
+
+    def get_canonical_shape(self) -> tuple[int, ...]:
+        """Return the shape in the canonical order."""
+        transforms = self._axes_mapper.to_canonical()
+        return tuple(transform_list(list(self._shape), 1, transforms))
+
+    def __repr__(self) -> str:
+        """Return the string representation of the object."""
+        return str(self)
+
+    @property
+    def on_disk_shape(self) -> tuple[int, ...]:
+        """Return the shape as a tuple."""
+        return tuple(self._shape)
+
+    @property
+    def is_time_series(self) -> bool:
+        """Return whether the data is a time series."""
+        if self.get("t", strict=False) == 1:
+            return False
+        return True
+
+    @property
+    def is_2d(self) -> bool:
+        """Return whether the data is 2D."""
+        if self.get("z", strict=False) != 1:
+            return False
+        return True
+
+    @property
+    def is_2d_time_series(self) -> bool:
+        """Return whether the data is a 2D time series."""
+        return self.is_2d and self.is_time_series
+
+    @property
+    def is_3d(self) -> bool:
+        """Return whether the data is 3D."""
+        return not self.is_2d
+
+    @property
+    def is_3d_time_series(self) -> bool:
+        """Return whether the data is a 3D time series."""
+        return self.is_3d and self.is_time_series
+
+    @property
+    def is_multi_channels(self) -> bool:
+        """Return whether the data has multiple channels."""
+        if self.get("c", strict=False) == 1:
+            return False
+        return True