ngio 0.4.7__py3-none-any.whl → 0.5.0__py3-none-any.whl

ngio/images/_abstract_image.py

@@ -1,19 +1,32 @@
 """Generic class to handle Image-like data in a OME-NGFF file."""
 
-from collections.abc import Sequence
-from typing import Generic, Literal, TypeVar
+import warnings
+from abc import ABC, abstractmethod
+from collections.abc import Mapping, Sequence
+from typing import Any, Literal
 
 import dask.array as da
 import numpy as np
 import zarr
+from zarr.core.array import CompressorLike
 
 from ngio.common import (
     Dimensions,
     InterpolationOrder,
     Roi,
-    RoiPixels,
     consolidate_pyramid,
 )
+from ngio.common._pyramid import (
+    ChunksLike,
+    ShardsLike,
+    compute_scales_from_shapes,
+    compute_shapes_from_scaling_factors,
+)
+from ngio.images._create_utils import (
+    _image_or_label_meta,
+    compute_base_scale,
+    init_image_like_from_shapes,
+)
 from ngio.io_pipes import (
     DaskGetter,
     DaskRoiGetter,
@@ -31,15 +44,32 @@ from ngio.ome_zarr_meta import (
     Dataset,
     ImageMetaHandler,
     LabelMetaHandler,
+    NgioImageMeta,
     PixelSize,
 )
+from ngio.ome_zarr_meta.ngio_specs import (
+    Channel,
+    NgffVersions,
+    NgioLabelMeta,
+)
+from ngio.ome_zarr_meta.ngio_specs._axes import (
+    AxesSetup,
+    DefaultSpaceUnit,
+    DefaultTimeUnit,
+    SpaceUnits,
+    TimeUnits,
+)
 from ngio.tables import RoiTable
-from ngio.utils import NgioFileExistsError, ZarrGroupHandler
-
-_image_handler = TypeVar("_image_handler", ImageMetaHandler, LabelMetaHandler)
+from ngio.utils import (
+    NgioFileExistsError,
+    NgioValueError,
+    StoreOrGroup,
+    ZarrGroupHandler,
+)
+from ngio.utils._zarr_utils import find_dimension_separator
 
 
-class AbstractImage(Generic[_image_handler]):
+class AbstractImage(ABC):
     """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.
@@ -49,7 +79,7 @@ class AbstractImage(Generic[_image_handler]):
         self,
         group_handler: ZarrGroupHandler,
         path: str,
-        meta_handler: _image_handler,
+        meta_handler: ImageMetaHandler | LabelMetaHandler,
     ) -> None:
         """Initialize the Image at a single level.
 
@@ -78,14 +108,21 @@ class AbstractImage(Generic[_image_handler]):
         return self._path
 
     @property
-    def meta_handler(self) -> _image_handler:
+    @abstractmethod
+    def meta_handler(self) -> ImageMetaHandler | LabelMetaHandler:
         """Return the metadata."""
-        return self._meta_handler
+        pass
+
+    @property
+    @abstractmethod
+    def meta(self) -> NgioImageMeta | NgioLabelMeta:
+        """Return the metadata."""
+        pass
 
     @property
     def dataset(self) -> Dataset:
         """Return the dataset of the image."""
-        return self.meta_handler.meta.get_dataset(path=self.path)
+        return self.meta_handler.get_meta().get_dataset(path=self.path)
 
     @property
     def dimensions(self) -> Dimensions:
@@ -106,6 +143,11 @@ class AbstractImage(Generic[_image_handler]):
         """Return the axes handler of the image."""
         return self.dataset.axes_handler
 
+    @property
+    def axes_setup(self) -> AxesSetup:
+        """Return the axes setup of the image."""
+        return self.axes_handler.axes_setup
+
     @property
     def axes(self) -> tuple[str, ...]:
         """Return the axes of the image."""
@@ -175,6 +217,47 @@ class AbstractImage(Generic[_image_handler]):
         """Return True if the image has the given axis."""
         return self.axes_handler.has_axis(axis)
 
+    def set_axes_unit(
+        self,
+        space_unit: SpaceUnits = DefaultSpaceUnit,
+        time_unit: TimeUnits = DefaultTimeUnit,
+    ) -> None:
+        """Set the axes unit of the image.
+
+        Args:
+            space_unit (SpaceUnits): The space unit of the image.
+            time_unit (TimeUnits): The time unit of the image.
+        """
+        meta = self._meta_handler.get_meta()
+        meta = meta.to_units(space_unit=space_unit, time_unit=time_unit)
+        self._meta_handler.update_meta(meta)  # type: ignore
+
+    def set_axes_names(self, axes_names: Sequence[str]) -> None:
+        """Set the axes names of the label.
+
+        Args:
+            axes_names (Sequence[str]): The axes names to set.
+        """
+        meta = self._meta_handler.get_meta()
+        meta = meta.rename_axes(axes_names=axes_names)
+        self._meta_handler._axes_setup = meta.axes_handler.axes_setup
+        self._meta_handler.update_meta(meta)  # type: ignore
+
+    def set_name(
+        self,
+        name: str,
+    ) -> None:
+        """Set the name of the image in the metadata.
+
+        This does not change the group name or any paths.
+
+        Args:
+            name (str): The name of the image.
+        """
+        meta = self._meta_handler.get_meta()
+        meta = meta.rename_image(name=name)
+        self._meta_handler.update_meta(meta)  # type: ignore
+
     def _get_as_numpy(
         self,
         axes_order: Sequence[str] | None = None,
@@ -202,7 +285,7 @@ class AbstractImage(Generic[_image_handler]):
 
     def _get_roi_as_numpy(
         self,
-        roi: Roi | RoiPixels,
+        roi: Roi,
         axes_order: Sequence[str] | None = None,
         transforms: Sequence[TransformProtocol] | None = None,
         **slicing_kwargs: SlicingInputType,
@@ -252,7 +335,7 @@ class AbstractImage(Generic[_image_handler]):
 
     def _get_roi_as_dask(
         self,
-        roi: Roi | RoiPixels,
+        roi: Roi,
         axes_order: Sequence[str] | None = None,
         transforms: Sequence[TransformProtocol] | None = None,
         **slicing_kwargs: SlicingInputType,
@@ -309,7 +392,7 @@ class AbstractImage(Generic[_image_handler]):
 
     def _get_roi(
         self,
-        roi: Roi | RoiPixels,
+        roi: Roi,
         axes_order: Sequence[str] | None = None,
         transforms: Sequence[TransformProtocol] | None = None,
         mode: Literal["numpy", "dask"] = "numpy",
@@ -385,7 +468,7 @@ class AbstractImage(Generic[_image_handler]):
 
     def _set_roi(
         self,
-        roi: Roi | RoiPixels,
+        roi: Roi,
         patch: np.ndarray | da.Array,
         axes_order: Sequence[str] | None = None,
         transforms: Sequence[TransformProtocol] | None = None,
@@ -444,25 +527,14 @@ class AbstractImage(Generic[_image_handler]):
 
     def roi(self, name: str | None = "image") -> Roi:
         """Return the ROI covering the entire image."""
-        dim_x = self.dimensions.get("x")
-        dim_y = self.dimensions.get("y")
-        assert dim_x is not None and dim_y is not None
-        dim_z = self.dimensions.get("z")
-        z = None if dim_z is None else 0
-        dim_t = self.dimensions.get("t")
-        t = None if dim_t is None else 0
-        roi_px = RoiPixels(
-            name=name,
-            x=0,
-            y=0,
-            z=z,
-            t=t,
-            x_length=dim_x,
-            y_length=dim_y,
-            z_length=dim_z,
-            t_length=dim_t,
-        )
-        return roi_px.to_roi(pixel_size=self.pixel_size)
+        slices = {}
+        for ax_name in ["t", "z", "y", "x"]:
+            axis_size = self.dimensions.get(ax_name, default=None)
+            if axis_size is None:
+                continue
+            slices[ax_name] = slice(0, axis_size)
+        roi_px = Roi.from_values(name=name, slices=slices, space="pixel")
+        return roi_px.to_world(pixel_size=self.pixel_size)
 
     def build_image_roi_table(self, name: str | None = "image") -> RoiTable:
         """Build the ROI table containing the ROI covering the entire image."""
@@ -576,7 +648,7 @@ def consolidate_image(
     mode: Literal["dask", "numpy", "coarsen"] = "dask",
 ) -> None:
     """Consolidate the image on disk."""
-    target_paths = image._meta_handler.meta.paths
+    target_paths = image.meta_handler.get_meta().paths
     targets = [
         image._group_handler.get_array(path)
         for path in target_paths
@@ -585,3 +657,435 @@ def consolidate_image(
     consolidate_pyramid(
         source=image.zarr_array, targets=targets, order=order, mode=mode
     )
+
+
+def _shapes_from_ref_image(
+    ref_image: AbstractImage,
+) -> tuple[list[tuple[int, ...]], list[tuple[float, ...]]]:
+    """Rebuild base shape based on a new shape."""
+    meta = ref_image.meta
+    paths = meta.paths
+    index_path = paths.index(ref_image.path)
+    sub_paths = paths[index_path:]
+    group_handler = ref_image._group_handler
+    shapes, scales = [], []
+    for path in sub_paths:
+        zarr_array = group_handler.get_array(path)
+        shapes.append(zarr_array.shape)
+        scales.append(meta.get_dataset(path=path).scale)
+    if len(shapes) == len(paths):
+        return shapes, scales
+    missing_levels = len(paths) - len(shapes)
+    extended_shapes = compute_shapes_from_scaling_factors(
+        base_shape=shapes[-1],
+        scaling_factors=ref_image.meta.scaling_factor(),
+        num_levels=missing_levels + 1,
+    )
+    shapes.extend(extended_shapes[1:])
+    extended_scales = compute_scales_from_shapes(
+        shapes=extended_shapes,
+        base_scale=scales[-1],
+    )
+    scales.extend(extended_scales[1:])
+    return shapes, scales
+
+
+def _shapes_from_new_shape(
+    ref_image: AbstractImage,
+    shape: Sequence[int],
+) -> tuple[list[tuple[int, ...]], list[tuple[float, ...]]]:
+    """Rebuild pyramid shapes based on a new base shape."""
+    if len(shape) != len(ref_image.shape):
+        raise NgioValueError(
+            "The shape of the new image does not match the reference image."
+            f"Got shape {shape} for reference shape {ref_image.shape}."
+        )
+    base_shape = tuple(shape)
+    scaling_factors = ref_image.meta.scaling_factor()
+    num_levels = len(ref_image.meta.paths)
+    shapes = compute_shapes_from_scaling_factors(
+        base_shape=base_shape,
+        scaling_factors=scaling_factors,
+        num_levels=num_levels,
+    )
+    scales = compute_scales_from_shapes(
+        shapes=shapes,
+        base_scale=ref_image.dataset.scale,
+    )
+    return shapes, scales
+
+
+def _compute_pyramid_shapes(
+    ref_image: AbstractImage,
+    shape: Sequence[int] | None,
+) -> tuple[list[tuple[int, ...]], list[tuple[float, ...]]]:
+    """Rebuild pyramid shapes based on a new base shape."""
+    if shape is None:
+        return _shapes_from_ref_image(ref_image=ref_image)
+    return _shapes_from_new_shape(ref_image=ref_image, shape=shape)
+
+
+def _check_len_compatibility(
+    ref_shape: tuple[int, ...],
+    chunks: ChunksLike,
+    shards: ShardsLike | None,
+    translation: Sequence[float] | None = None,
+) -> None:
+    """Check if the chunks and shards are compatible with the reference shape.
+
+    Args:
+        ref_shape: The reference shape.
+        chunks: The chunks to check.
+        shards: The shards to check.
+        translation: The translation to check.
+    """
+    if chunks != "auto":
+        if len(chunks) != len(ref_shape):
+            raise NgioValueError(
+                "The length of the chunks must be the same as the number of dimensions."
+            )
+    if shards is not None and shards != "auto":
+        if len(shards) != len(ref_shape):
+            raise NgioValueError(
+                "The length of the shards must be the same as the number of dimensions."
+            )
+    if translation is not None:
+        if len(translation) != len(ref_shape):
+            raise NgioValueError(
+                "The length of the translation must be the same as the number of "
+                "dimensions."
+            )
+
+
+def _apply_channel_policy(
+    ref_image: AbstractImage,
+    channels_policy: Literal["squeeze", "same", "singleton"] | int,
+    shapes: list[tuple[int, ...]],
+    axes: tuple[str, ...],
+    chunks: ChunksLike,
+    shards: ShardsLike | None,
+    translation: Sequence[float],
+    scales: list[tuple[float, ...]] | tuple[float, ...],
+) -> tuple[
+    list[tuple[int, ...]],
+    tuple[str, ...],
+    ChunksLike,
+    ShardsLike | None,
+    tuple[float, ...],
+    list[tuple[float, ...]] | tuple[float, ...],
+]:
+    """Apply the channel policy to the shapes and axes.
+
+    Args:
+        ref_image: The reference image.
+        channels_policy: The channels policy to apply.
+        shapes: The shapes of the pyramid levels.
+        axes: The axes of the image.
+        chunks: The chunks of the image.
+        shards: The shards of the image.
+        translation: The translation of the image.
+        scales: The scales of the image.
+
+    Returns:
+        The new shapes and axes after applying the channel policy.
+    """
+    translation = tuple(translation)
+    if channels_policy == "same":
+        return shapes, axes, chunks, shards, translation, scales
+
+    if channels_policy == "singleton":
+        # Treat 'singleton' as setting channel size to 1
+        channels_policy = 1
+
+    channel_index = ref_image.axes_handler.get_index("c")
+    if channel_index is None:
+        if channels_policy == "squeeze":
+            return shapes, axes, chunks, shards, translation, scales
+        raise NgioValueError(
+            f"Cannot apply channel policy {channels_policy=} to an image "
+            "without channels axis."
+        )
+    if channels_policy == "squeeze":
+        new_shapes = []
+        for shape in shapes:
+            new_shape = shape[:channel_index] + shape[channel_index + 1 :]
+            new_shapes.append(new_shape)
+
+        if isinstance(scales, tuple):
+            new_scales = scales[:channel_index] + scales[channel_index + 1 :]
+        else:
+            new_scales = []
+            for scale in scales:
+                new_scale = scale[:channel_index] + scale[channel_index + 1 :]
+                new_scales.append(new_scale)
+
+        new_axes = axes[:channel_index] + axes[channel_index + 1 :]
+        if chunks == "auto":
+            new_chunks: ChunksLike = "auto"
+        else:
+            new_chunks = chunks[:channel_index] + chunks[channel_index + 1 :]
+        if shards == "auto" or shards is None:
+            new_shards: ShardsLike | None = shards
+        else:
+            new_shards = shards[:channel_index] + shards[channel_index + 1 :]
+
+        translation = translation[:channel_index] + translation[channel_index + 1 :]
+        return new_shapes, new_axes, new_chunks, new_shards, translation, new_scales
+    elif isinstance(channels_policy, int):
+        new_shapes = []
+        for shape in shapes:
+            new_shape = (
+                *shape[:channel_index],
+                channels_policy,
+                *shape[channel_index + 1 :],
+            )
+            new_shapes.append(new_shape)
+        return new_shapes, axes, chunks, shards, translation, scales
+    else:
+        raise NgioValueError(
+            f"Invalid channels policy: {channels_policy}. "
+            "Must be 'squeeze', 'same', or an integer."
+        )
+
+
+def _check_channels_meta_compatibility(
+    meta_type: type[_image_or_label_meta],
+    ref_image: AbstractImage,
+    channels_meta: Sequence[str | Channel] | None,
+) -> Sequence[str | Channel] | None:
+    """Check if the channels metadata is compatible with the reference image.
+
+    Args:
+        meta_type: The metadata type.
+        ref_image: The reference image.
+        channels_meta: The channels metadata to check.
+
+    Returns:
+        The channels metadata if compatible, None otherwise.
+    """
+    if issubclass(meta_type, NgioLabelMeta):
+        if channels_meta is not None:
+            raise NgioValueError("Cannot set channels_meta for a label image.")
+        return None
+    if channels_meta is not None:
+        return channels_meta
+    assert isinstance(ref_image.meta, NgioImageMeta)
+    ref_meta = ref_image.meta
+    index_c = ref_meta.axes_handler.get_index("c")
+    if index_c is None:
+        return None
+
+    # If the channels number does not match, return None
+    # Else return the channels metadata from the reference image
+    ref_shape = ref_image.shape
+    ref_num_channels = ref_shape[index_c] if index_c is not None else 1
+    channels_ = ref_meta.channels_meta.channels if ref_meta.channels_meta else []
+    # Reset to None if number of channels do not match
+    channels_meta_ = channels_ if ref_num_channels == len(channels_) else None
+    return channels_meta_
+
+
+def adapt_scales(
+    scales: list[tuple[float, ...]],
+    pixelsize: float | tuple[float, float] | None,
+    z_spacing: float | None,
+    time_spacing: float | None,
+    ref_image: AbstractImage,
+) -> list[tuple[float, ...]] | tuple[float, ...]:
+    if pixelsize is None and z_spacing is None and time_spacing is None:
+        return scales
+    pixel_size = ref_image.pixel_size
+    if pixelsize is None:
+        pixelsize = (pixel_size.y, pixel_size.x)
+    if z_spacing is None:
+        z_spacing = pixel_size.z
+    else:
+        z_spacing = z_spacing
+    if time_spacing is None:
+        time_spacing = pixel_size.t
+    else:
+        time_spacing = time_spacing
+    base_scale = compute_base_scale(
+        pixelsize=pixelsize,
+        z_spacing=z_spacing,
+        time_spacing=time_spacing,
+        axes_handler=ref_image.axes_handler,
+    )
+    return base_scale
+
+
+def abstract_derive(
+    *,
+    ref_image: AbstractImage,
+    meta_type: type[_image_or_label_meta],
+    store: StoreOrGroup,
+    overwrite: bool = False,
+    # Metadata parameters
+    shape: Sequence[int] | None = None,
+    pixelsize: float | tuple[float, float] | None = None,
+    z_spacing: float | None = None,
+    time_spacing: float | None = None,
+    name: str | None = None,
+    translation: Sequence[float] | None = None,
+    channels_policy: Literal["squeeze", "same", "singleton"] | int = "same",
+    channels_meta: Sequence[str | Channel] | None = None,
+    ngff_version: NgffVersions | None = None,
+    # Zarr Array parameters
+    chunks: ChunksLike | None = None,
+    shards: ShardsLike | None = None,
+    dtype: str | None = None,
+    dimension_separator: Literal[".", "/"] | None = None,
+    compressors: CompressorLike | None = None,
+    extra_array_kwargs: Mapping[str, Any] | None = None,
+    # Deprecated arguments
+    labels: Sequence[str] | None = None,
+    pixel_size: PixelSize | None = None,
+) -> tuple[ZarrGroupHandler, AxesSetup]:
+    """Create an empty OME-Zarr image from an existing image.
+
+    If a kwarg is not provided, the value from the reference image will be used.
+
+    Args:
+        ref_image (AbstractImage): The reference image to derive from.
+        meta_type (type[_image_or_label_meta]): The metadata type to use.
+        store (StoreOrGroup): The Zarr store or group to create the image in.
+        overwrite (bool): Whether to overwrite an existing image.
+        shape (Sequence[int] | None): The shape of the new image.
+        pixelsize (float | tuple[float, float] | None): The pixel size of the new image.
+        z_spacing (float | None): The z spacing of the new image.
+        time_spacing (float | None): The time spacing of the new image.
+        name (str | None): The name of the new image.
+        translation (Sequence[float] | None): The translation for each axis
+            at the highest resolution level. Defaults to None.
+        channels_policy (Literal["squeeze", "same", "singleton"] | int):
+            Possible policies:
+            - If "squeeze", the channels axis will be removed (no matter its size).
+            - If "same", the channels axis will be kept as is (if it exists).
+            - If "singleton", the channels axis will be set to size 1.
+            - If an integer is provided, the channels axis will be changed to have that
+                size.
+        channels_meta (Sequence[str | Channel] | None): The channels metadata
+            of the new image.
+        ngff_version (NgffVersions | None): The NGFF version to use.
+        chunks (ChunksLike | None): The chunk shape of the new image.
+        shards (ShardsLike | None): The shard shape of the new image.
+        dtype (str | None): The data type of the new image.
+        dimension_separator (Literal[".", "/"] | None): The separator to use for
+            dimensions.
+        compressors (CompressorLike | None): The compressors to use.
+        extra_array_kwargs (Mapping[str, Any] | None): Extra arguments to pass to
+            the zarr array creation.
+        labels (Sequence[str] | None): The labels of the new image.
+            This argument is DEPRECATED please use channels_meta instead.
+        pixel_size (PixelSize | None): The pixel size of the new image.
+            This argument is DEPRECATED please use pixelsize, z_spacing,
+            and time_spacing instead.
+
+    Returns:
+        ImagesContainer: The new derived image.
+
+    """
+    # TODO: remove in ngio 0.6
+    if labels is not None:
+        warnings.warn(
+            "The 'labels' argument is deprecated and will be removed in "
+            "ngio=0.6. Please use 'channels_meta' instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        channels_meta = list(labels)
+    if pixel_size is not None:
+        warnings.warn(
+            "The 'pixel_size' argument is deprecated and will be removed in "
+            "ngio=0.6. Please use 'pixelsize', 'z_spacing', and 'time_spacing'"
+            "instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        pixelsize = (pixel_size.y, pixel_size.x)
+    # End of deprecated arguments handling
+    ref_meta = ref_image.meta
+
+    shapes, scales = _compute_pyramid_shapes(
+        shape=shape,
+        ref_image=ref_image,
+    )
+    ref_shape = next(iter(shapes))
+
+    scales = adapt_scales(
+        scales=scales,
+        pixelsize=pixelsize,
+        z_spacing=z_spacing,
+        time_spacing=time_spacing,
+        ref_image=ref_image,
+    )
+
+    if name is None:
+        name = ref_meta.name
+
+    if dtype is None:
+        dtype = ref_image.dtype
+
+    if dimension_separator is None:
+        dimension_separator = find_dimension_separator(ref_image.zarr_array)
+
+    if compressors is None:
+        compressors = ref_image.zarr_array.compressors  # type: ignore
+
+    if translation is None:
+        translation = ref_image.dataset.translation
+
+    if chunks is None:
+        chunks = ref_image.zarr_array.chunks
+    if shards is None:
+        shards = ref_image.zarr_array.shards
+
+    _check_len_compatibility(
+        ref_shape=ref_shape,
+        chunks=chunks,
+        shards=shards,
+        translation=translation,
+    )
+
+    if ngff_version is None:
+        ngff_version = ref_meta.version
+
+    shapes, axes, chunks, shards, translation, scales = _apply_channel_policy(
+        ref_image=ref_image,
+        channels_policy=channels_policy,
+        shapes=shapes,
+        axes=ref_image.axes,
+        chunks=chunks,
+        shards=shards,
+        translation=translation,
+        scales=scales,
+    )
+    channels_meta_ = _check_channels_meta_compatibility(
+        meta_type=meta_type,
+        ref_image=ref_image,
+        channels_meta=channels_meta,
+    )
+
+    handler, axes_setup = init_image_like_from_shapes(
+        store=store,
+        meta_type=meta_type,
+        shapes=shapes,
+        base_scale=scales,
+        levels=ref_meta.paths,
+        translation=translation,
+        time_unit=ref_image.time_unit,
+        space_unit=ref_image.space_unit,
+        axes_names=axes,
+        name=name,
+        axes_setup=ref_image.axes_setup,
+        channels_meta=channels_meta_,
+        chunks=chunks,
+        shards=shards,
+        dtype=dtype,
+        dimension_separator=dimension_separator,
+        compressors=compressors,
+        overwrite=overwrite,
+        ngff_version=ngff_version,
+        extra_array_kwargs=extra_array_kwargs,
+    )
+    return handler, axes_setup