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

ngio/images/abstract_image.py

@@ -0,0 +1,349 @@
+"""Generic class to handle Image-like data in a OME-NGFF file."""
+
+from collections.abc import Collection, Iterable
+from typing import Generic, Literal, TypeVar
+
+import zarr
+
+from ngio.common import (
+    ArrayLike,
+    Dimensions,
+    Roi,
+    RoiPixels,
+    consolidate_pyramid,
+    get_pipe,
+    roi_to_slice_kwargs,
+    set_pipe,
+)
+from ngio.ome_zarr_meta import (
+    AxesMapper,
+    Dataset,
+    ImageMetaHandler,
+    LabelMetaHandler,
+    PixelSize,
+)
+from ngio.tables import RoiTable
+from ngio.utils import NgioFileExistsError, ZarrGroupHandler
+
+_image_handler = TypeVar("_image_handler", ImageMetaHandler, LabelMetaHandler)
+
+
+class AbstractImage(Generic[_image_handler]):
+    """A class to handle a single image (or level) in an OME-Zarr image.
+
+    This class is meant to be subclassed by specific image types.
+    """
+
+    def __init__(
+        self,
+        group_handler: ZarrGroupHandler,
+        path: str,
+        meta_handler: _image_handler,
+    ) -> None:
+        """Initialize the Image at a single level.
+
+        Args:
+            group_handler: The Zarr group handler.
+            path: The path to the image in the ome_zarr file.
+            meta_handler: The image metadata handler.
+
+        """
+        self._path = path
+        self._group_handler = group_handler
+        self._meta_handler = meta_handler
+
+        self._dataset = self._meta_handler.meta.get_dataset(path=path)
+        self._pixel_size = self._dataset.pixel_size
+
+        try:
+            self._zarr_array = self._group_handler.get_array(self._dataset.path)
+        except NgioFileExistsError as e:
+            raise NgioFileExistsError(f"Could not find the dataset at {path}.") from e
+
+        self._dimensions = Dimensions(
+            shape=self._zarr_array.shape, axes_mapper=self._dataset.axes_mapper
+        )
+
+        self._axes_mapper = self._dataset.axes_mapper
+
+    def __repr__(self) -> str:
+        """Return a string representation of the image."""
+        return f"Image(path={self.path}, {self.dimensions})"
+
+    @property
+    def meta_handler(self) -> _image_handler:
+        """Return the metadata."""
+        return self._meta_handler
+
+    @property
+    def zarr_array(self) -> zarr.Array:
+        """Return the Zarr array."""
+        return self._zarr_array
+
+    @property
+    def shape(self) -> tuple[int, ...]:
+        """Return the shape of the image."""
+        return self.zarr_array.shape
+
+    @property
+    def dtype(self) -> str:
+        """Return the dtype of the image."""
+        return str(self.zarr_array.dtype)
+
+    @property
+    def chunks(self) -> tuple[int, ...]:
+        """Return the chunks of the image."""
+        return self.zarr_array.chunks
+
+    @property
+    def dimensions(self) -> Dimensions:
+        """Return the dimensions of the image."""
+        return self._dimensions
+
+    @property
+    def axes_mapper(self) -> AxesMapper:
+        """Return the axes mapper of the image."""
+        return self._axes_mapper
+
+    @property
+    def is_3d(self) -> bool:
+        """Return True if the image is 3D."""
+        return self.dimensions.is_3d
+
+    @property
+    def is_2d(self) -> bool:
+        """Return True if the image is 2D."""
+        return self.dimensions.is_2d
+
+    @property
+    def is_time_series(self) -> bool:
+        """Return True if the image is a time series."""
+        return self.dimensions.is_time_series
+
+    @property
+    def is_2d_time_series(self) -> bool:
+        """Return True if the image is a 2D time series."""
+        return self.dimensions.is_2d_time_series
+
+    @property
+    def is_3d_time_series(self) -> bool:
+        """Return True if the image is a 3D time series."""
+        return self.dimensions.is_3d_time_series
+
+    @property
+    def is_multi_channels(self) -> bool:
+        """Return True if the image is multichannel."""
+        return self.dimensions.is_multi_channels
+
+    @property
+    def pixel_size(self) -> PixelSize:
+        """Return the pixel size of the image."""
+        return self._pixel_size
+
+    @property
+    def dataset(self) -> Dataset:
+        """Return the dataset of the image."""
+        return self._dataset
+
+    @property
+    def path(self) -> str:
+        """Return the path of the image."""
+        return self._dataset.path
+
+    def has_axis(self, axis: str) -> bool:
+        """Return True if the image has the given axis."""
+        return self.dimensions.has_axis(axis)
+
+    def get_array(
+        self,
+        axes_order: Collection[str] | None = None,
+        mode: Literal["numpy", "dask", "delayed"] = "numpy",
+        **slice_kwargs: slice | int | Iterable[int],
+    ) -> ArrayLike:
+        """Get a slice of the image.
+
+        Args:
+            axes_order: The order of the axes to return the array.
+            mode: The mode to return the array.
+            **slice_kwargs: The slices to get the array.
+
+        Returns:
+            The array of the region of interest.
+        """
+        return get_pipe(
+            array=self.zarr_array,
+            dimensions=self.dimensions,
+            axes_order=axes_order,
+            mode=mode,
+            **slice_kwargs,
+        )
+
+    def get_roi(
+        self,
+        roi: Roi,
+        axes_order: Collection[str] | None = None,
+        mode: Literal["numpy", "dask", "delayed"] = "numpy",
+        **slice_kwargs: slice | int | Iterable[int],
+    ) -> ArrayLike:
+        """Get a slice of the image.
+
+        Args:
+            roi: The region of interest to get the array.
+            axes_order: The order of the axes to return the array.
+            mode: The mode to return the array.
+            **slice_kwargs: The slices to get the array.
+
+        Returns:
+            The array of the region of interest.
+        """
+        return get_roi_pipe(
+            image=self, roi=roi, axes_order=axes_order, mode=mode, **slice_kwargs
+        )
+
+    def set_array(
+        self,
+        patch: ArrayLike,
+        axes_order: Collection[str] | None = None,
+        **slice_kwargs: slice | int | Iterable[int],
+    ) -> None:
+        """Set a slice of the image.
+
+        Args:
+            patch: The patch to set.
+            axes_order: The order of the axes to set the patch.
+            **slice_kwargs: The slices to set the patch.
+
+        """
+        set_pipe(
+            array=self.zarr_array,
+            patch=patch,
+            dimensions=self.dimensions,
+            axes_order=axes_order,
+            **slice_kwargs,
+        )
+
+    def set_roi(
+        self,
+        roi: Roi,
+        patch: ArrayLike,
+        axes_order: Collection[str] | None = None,
+        **slice_kwargs: slice | int | Iterable[int],
+    ) -> None:
+        """Set a slice of the image.
+
+        Args:
+            roi: The region of interest to set the patch.
+            patch: The patch to set.
+            axes_order: The order of the axes to set the patch.
+            **slice_kwargs: The slices to set the patch.
+
+        """
+        return set_roi_pipe(
+            image=self, roi=roi, patch=patch, axes_order=axes_order, **slice_kwargs
+        )
+
+    def build_image_roi_table(self, name: str = "image") -> RoiTable:
+        """Build the ROI table for an image."""
+        return build_image_roi_table(image=self, name=name)
+
+
+def consolidate_image(
+    image: AbstractImage,
+    order: Literal[0, 1, 2] = 1,
+    mode: Literal["dask", "numpy", "coarsen"] = "dask",
+) -> None:
+    """Consolidate the image on disk."""
+    target_paths = image._meta_handler.meta.paths
+    targets = [
+        image._group_handler.get_array(path)
+        for path in target_paths
+        if path != image.path
+    ]
+    consolidate_pyramid(
+        source=image.zarr_array, targets=targets, order=order, mode=mode
+    )
+
+
+def get_roi_pipe(
+    image: AbstractImage,
+    roi: Roi,
+    axes_order: Collection[str] | None = None,
+    mode: Literal["numpy", "dask", "delayed"] = "numpy",
+    **slice_kwargs: slice | int | Iterable[int],
+) -> ArrayLike:
+    """Get a slice of the image.
+
+    Args:
+        image: The image to get the ROI.
+        roi: The region of interest to get the array.
+        axes_order: The order of the axes to return the array.
+        mode: The mode to return the array.
+        **slice_kwargs: The slices to get the array.
+
+    Returns:
+        The array of the region of interest.
+    """
+    slice_kwargs = roi_to_slice_kwargs(
+        roi=roi,
+        pixel_size=image.pixel_size,
+        dimensions=image.dimensions,
+        **slice_kwargs,
+    )
+    return get_pipe(
+        array=image.zarr_array,
+        dimensions=image.dimensions,
+        axes_order=axes_order,
+        mode=mode,
+        **slice_kwargs,
+    )
+
+
+def set_roi_pipe(
+    image: AbstractImage,
+    roi: Roi,
+    patch: ArrayLike,
+    axes_order: Collection[str] | None = None,
+    **slice_kwargs: slice | int | Iterable[int],
+) -> None:
+    """Set a slice of the image.
+
+    Args:
+        image: The image to set the ROI.
+        roi: The region of interest to set the patch.
+        patch: The patch to set.
+        axes_order: The order of the axes to set the patch.
+        **slice_kwargs: The slices to set the patch.
+
+    """
+    slice_kwargs = roi_to_slice_kwargs(
+        roi=roi,
+        pixel_size=image.pixel_size,
+        dimensions=image.dimensions,
+        **slice_kwargs,
+    )
+    set_pipe(
+        array=image.zarr_array,
+        patch=patch,
+        dimensions=image.dimensions,
+        axes_order=axes_order,
+        **slice_kwargs,
+    )
+
+
+def build_image_roi_table(image: AbstractImage, name: str = "image") -> RoiTable:
+    """Build the ROI table for an image."""
+    dim_z, dim_y, dim_x = (
+        image.dimensions.get("z", strict=False),
+        image.dimensions.get("y"),
+        image.dimensions.get("x"),
+    )
+    image_roi = RoiPixels(
+        name=name,
+        x=0,
+        y=0,
+        z=0,
+        x_length=dim_x,
+        y_length=dim_y,
+        z_length=dim_z,
+    )
+    return RoiTable(rois=[image_roi.to_roi(pixel_size=image.pixel_size)])

ngio/images/create.py

@@ -0,0 +1,270 @@
+"""Utility functions for working with OME-Zarr images."""
+
+from collections.abc import Collection
+from typing import TypeVar
+
+from ngio.common._pyramid import init_empty_pyramid
+from ngio.ome_zarr_meta import (
+    NgioImageMeta,
+    NgioLabelMeta,
+    PixelSize,
+    get_image_meta_handler,
+    get_label_meta_handler,
+)
+from ngio.ome_zarr_meta.ngio_specs import (
+    SpaceUnits,
+    TimeUnits,
+    canonical_axes_order,
+    canonical_label_axes_order,
+)
+from ngio.utils import NgioValueError, StoreOrGroup, ZarrGroupHandler
+
+_image_or_label_meta = TypeVar("_image_or_label_meta", NgioImageMeta, NgioLabelMeta)
+
+
+def _init_generic_meta(
+    meta_type: type[_image_or_label_meta],
+    pixelsize: float,
+    axes_names: Collection[str],
+    z_spacing: float = 1.0,
+    time_spacing: float = 1.0,
+    levels: int | list[str] = 5,
+    yx_scaling_factor: float | tuple[float, float] = 2.0,
+    z_scaling_factor: float = 1.0,
+    space_unit: SpaceUnits | str | None = None,
+    time_unit: TimeUnits | str | None = None,
+    name: str | None = None,
+    version: str = "0.4",
+) -> tuple[_image_or_label_meta, list[float]]:
+    """Initialize the metadata for an image or label."""
+    scaling_factors = []
+    for ax in axes_names:
+        if ax == "z":
+            scaling_factors.append(z_scaling_factor)
+        elif ax in ["x"]:
+            if isinstance(yx_scaling_factor, tuple):
+                scaling_factors.append(yx_scaling_factor[1])
+            else:
+                scaling_factors.append(yx_scaling_factor)
+        elif ax in ["y"]:
+            if isinstance(yx_scaling_factor, tuple):
+                scaling_factors.append(yx_scaling_factor[0])
+            else:
+                scaling_factors.append(yx_scaling_factor)
+        else:
+            scaling_factors.append(1.0)
+
+    if space_unit is None:
+        space_unit = SpaceUnits.micrometer
+    elif isinstance(space_unit, str):
+        space_unit = SpaceUnits(space_unit)
+    elif not isinstance(space_unit, SpaceUnits):
+        raise NgioValueError(f"space_unit can not be {type(space_unit)}.")
+
+    if time_unit is None:
+        time_unit = TimeUnits.seconds
+    elif isinstance(time_unit, str):
+        time_unit = TimeUnits(time_unit)
+    elif not isinstance(time_unit, TimeUnits):
+        raise NgioValueError(f"time_units can not be {type(time_unit)}.")
+
+    pixel_sizes = PixelSize(
+        x=pixelsize,
+        y=pixelsize,
+        z=z_spacing,
+        t=time_spacing,
+        space_unit=space_unit,
+        time_unit=time_unit,
+    )
+
+    meta = meta_type.default_init(
+        name=name,
+        levels=levels,
+        axes_names=axes_names,
+        pixel_size=pixel_sizes,
+        scaling_factors=scaling_factors,
+        version=version,
+    )
+    return meta, scaling_factors
+
+
+def _create_empty_label(
+    store: StoreOrGroup,
+    shape: Collection[int],
+    pixelsize: float,
+    z_spacing: float = 1.0,
+    time_spacing: float = 1.0,
+    levels: int | list[str] = 5,
+    yx_scaling_factor: float | tuple[float, float] = 2.0,
+    z_scaling_factor: float = 1.0,
+    space_unit: SpaceUnits | str | None = None,
+    time_unit: TimeUnits | str | None = None,
+    axes_names: Collection[str] | None = None,
+    name: str | None = None,
+    chunks: Collection[int] | None = None,
+    dtype: str = "uint16",
+    overwrite: bool = False,
+    version: str = "0.4",
+) -> ZarrGroupHandler:
+    """Create an empty label with the given shape and metadata.
+
+    Args:
+        store (StoreOrGroup): The Zarr store or group to create the image in.
+        shape (Collection[int]): The shape of the image.
+        pixelsize (float): The pixel size in x and y dimensions.
+        z_spacing (float, optional): The spacing between z slices. Defaults to 1.0.
+        time_spacing (float, optional): The spacing between time points.
+            Defaults to 1.0.
+        levels (int | list[str], optional): The number of levels in the pyramid or a
+            list of level names. Defaults to 5.
+        yx_scaling_factor (float, optional): The down-scaling factor in x and y
+            dimensions. Defaults to 2.0.
+        z_scaling_factor (float, optional): The down-scaling factor in z dimension.
+            Defaults to 1.0.
+        space_unit (SpaceUnits | str | None, optional): The unit of space. Defaults to
+            None.
+        time_unit (TimeUnits | str | None, optional): The unit of time. Defaults to
+            None.
+        axes_names (Collection[str] | None, optional): The names of the axes.
+            If None the canonical names are used. Defaults to None.
+        name (str | None, optional): The name of the image. Defaults to None.
+        chunks (Collection[int] | None, optional): The chunk shape. If None the shape
+            is used. Defaults to None.
+        dtype (str, optional): The data type of the image. Defaults to "uint16".
+        overwrite (bool, optional): Whether to overwrite an existing image.
+            Defaults to True.
+        version (str, optional): The version of the OME-Zarr specification.
+            Defaults to "0.4".
+
+    """
+    if axes_names is None:
+        axes_names = canonical_label_axes_order()[-len(shape) :]
+
+    if len(axes_names) != len(shape):
+        raise NgioValueError(
+            f"Number of axes names {axes_names} does not match the number of "
+            f"dimensions {shape}."
+        )
+
+    meta, scaling_factors = _init_generic_meta(
+        meta_type=NgioLabelMeta,
+        pixelsize=pixelsize,
+        z_spacing=z_spacing,
+        time_spacing=time_spacing,
+        levels=levels,
+        yx_scaling_factor=yx_scaling_factor,
+        z_scaling_factor=z_scaling_factor,
+        space_unit=space_unit,
+        time_unit=time_unit,
+        axes_names=axes_names,
+        name=name,
+        version=version,
+    )
+
+    mode = "w" if overwrite else "w-"
+    group_handler = ZarrGroupHandler(store=store, mode=mode, cache=False)
+    image_handler = get_label_meta_handler(version=version, group_handler=group_handler)
+    image_handler.write_meta(meta)
+
+    init_empty_pyramid(
+        store=store,
+        paths=meta.paths,
+        scaling_factors=scaling_factors,
+        ref_shape=shape,
+        chunks=chunks,
+        dtype=dtype,
+        mode="a",
+    )
+    group_handler._mode = "r+"
+    return group_handler
+
+
+def _create_empty_image(
+    store: StoreOrGroup,
+    shape: Collection[int],
+    pixelsize: float,
+    z_spacing: float = 1.0,
+    time_spacing: float = 1.0,
+    levels: int | list[str] = 5,
+    yx_scaling_factor: float | tuple[float, float] = 2,
+    z_scaling_factor: float = 1.0,
+    space_unit: SpaceUnits | str | None = None,
+    time_unit: TimeUnits | str | None = None,
+    axes_names: Collection[str] | None = None,
+    name: str | None = None,
+    chunks: Collection[int] | None = None,
+    dtype: str = "uint16",
+    overwrite: bool = False,
+    version: str = "0.4",
+) -> ZarrGroupHandler:
+    """Create an empty OME-Zarr image with the given shape and metadata.
+
+    Args:
+        store (StoreOrGroup): The Zarr store or group to create the image in.
+        shape (Collection[int]): The shape of the image.
+        pixelsize (float): The pixel size in x and y dimensions.
+        z_spacing (float, optional): The spacing between z slices. Defaults to 1.0.
+        time_spacing (float, optional): The spacing between time points.
+            Defaults to 1.0.
+        levels (int | list[str], optional): The number of levels in the pyramid or a
+            list of level names. Defaults to 5.
+        yx_scaling_factor (float, optional): The down-scaling factor in x and y
+            dimensions. Defaults to 2.0.
+        z_scaling_factor (float, optional): The down-scaling factor in z dimension.
+            Defaults to 1.0.
+        space_unit (SpaceUnits | str | None, optional): The unit of space. Defaults to
+            None.
+        time_unit (TimeUnits | str | None, optional): The unit of time. Defaults to
+            None.
+        axes_names (Collection[str] | None, optional): The names of the axes.
+            If None the canonical names are used. Defaults to None.
+        name (str | None, optional): The name of the image. Defaults to None.
+        chunks (Collection[int] | None, optional): The chunk shape. If None the shape
+            is used. Defaults to None.
+        dtype (str, optional): The data type of the image. Defaults to "uint16".
+        overwrite (bool, optional): Whether to overwrite an existing image.
+            Defaults to True.
+        version (str, optional): The version of the OME-Zarr specification.
+            Defaults to "0.4".
+
+    """
+    if axes_names is None:
+        axes_names = canonical_axes_order()[-len(shape) :]
+
+    if len(axes_names) != len(shape):
+        raise NgioValueError(
+            f"Number of axes names {axes_names} does not match the number of "
+            f"dimensions {shape}."
+        )
+
+    meta, scaling_factors = _init_generic_meta(
+        meta_type=NgioImageMeta,
+        pixelsize=pixelsize,
+        z_spacing=z_spacing,
+        time_spacing=time_spacing,
+        levels=levels,
+        yx_scaling_factor=yx_scaling_factor,
+        z_scaling_factor=z_scaling_factor,
+        space_unit=space_unit,
+        time_unit=time_unit,
+        axes_names=axes_names,
+        name=name,
+        version=version,
+    )
+    mode = "w" if overwrite else "w-"
+    group_handler = ZarrGroupHandler(store=store, mode=mode, cache=False)
+    image_handler = get_image_meta_handler(version=version, group_handler=group_handler)
+    image_handler.write_meta(meta)
+
+    init_empty_pyramid(
+        store=store,
+        paths=meta.paths,
+        scaling_factors=scaling_factors,
+        ref_shape=shape,
+        chunks=chunks,
+        dtype=dtype,
+        mode="a",
+    )
+
+    group_handler._mode = "r+"
+    return group_handler