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

ngio/images/_abstract_image.py

@@ -1,32 +1,19 @@
 """Generic class to handle Image-like data in a OME-NGFF file."""
 
-import warnings
-from abc import ABC, abstractmethod
-from collections.abc import Mapping, Sequence
-from typing import Any, Literal
+from collections.abc import Sequence
+from typing import Generic, Literal, TypeVar
 
 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,
@@ -44,32 +31,15 @@ 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,
-    NgioValueError,
-    StoreOrGroup,
-    ZarrGroupHandler,
-)
-from ngio.utils._zarr_utils import find_dimension_separator
+from ngio.utils import NgioFileExistsError, ZarrGroupHandler
+
+_image_handler = TypeVar("_image_handler", ImageMetaHandler, LabelMetaHandler)
 
 
-class AbstractImage(ABC):
+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.
@@ -79,7 +49,7 @@ class AbstractImage(ABC):
         self,
         group_handler: ZarrGroupHandler,
         path: str,
-        meta_handler: ImageMetaHandler | LabelMetaHandler,
+        meta_handler: _image_handler,
     ) -> None:
         """Initialize the Image at a single level.
 
@@ -108,21 +78,14 @@ class AbstractImage(ABC):
         return self._path
 
     @property
-    @abstractmethod
-    def meta_handler(self) -> ImageMetaHandler | LabelMetaHandler:
+    def meta_handler(self) -> _image_handler:
         """Return the metadata."""
-        pass
-
-    @property
-    @abstractmethod
-    def meta(self) -> NgioImageMeta | NgioLabelMeta:
-        """Return the metadata."""
-        pass
+        return self._meta_handler
 
     @property
     def dataset(self) -> Dataset:
         """Return the dataset of the image."""
-        return self.meta_handler.get_meta().get_dataset(path=self.path)
+        return self.meta_handler.meta.get_dataset(path=self.path)
 
     @property
     def dimensions(self) -> Dimensions:
@@ -143,11 +106,6 @@ class AbstractImage(ABC):
         """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."""
@@ -217,47 +175,6 @@ class AbstractImage(ABC):
         """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,
@@ -285,7 +202,7 @@ class AbstractImage(ABC):
 
     def _get_roi_as_numpy(
         self,
-        roi: Roi,
+        roi: Roi | RoiPixels,
         axes_order: Sequence[str] | None = None,
         transforms: Sequence[TransformProtocol] | None = None,
         **slicing_kwargs: SlicingInputType,
@@ -335,7 +252,7 @@ class AbstractImage(ABC):
 
     def _get_roi_as_dask(
         self,
-        roi: Roi,
+        roi: Roi | RoiPixels,
         axes_order: Sequence[str] | None = None,
         transforms: Sequence[TransformProtocol] | None = None,
         **slicing_kwargs: SlicingInputType,
@@ -392,7 +309,7 @@ class AbstractImage(ABC):
 
     def _get_roi(
         self,
-        roi: Roi,
+        roi: Roi | RoiPixels,
         axes_order: Sequence[str] | None = None,
         transforms: Sequence[TransformProtocol] | None = None,
         mode: Literal["numpy", "dask"] = "numpy",
@@ -468,7 +385,7 @@ class AbstractImage(ABC):
 
     def _set_roi(
         self,
-        roi: Roi,
+        roi: Roi | RoiPixels,
         patch: np.ndarray | da.Array,
         axes_order: Sequence[str] | None = None,
         transforms: Sequence[TransformProtocol] | None = None,
@@ -527,14 +444,25 @@ class AbstractImage(ABC):
 
     def roi(self, name: str | None = "image") -> Roi:
         """Return the ROI covering the entire image."""
-        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)
+        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)
 
     def build_image_roi_table(self, name: str | None = "image") -> RoiTable:
         """Build the ROI table containing the ROI covering the entire image."""
@@ -648,7 +576,7 @@ def consolidate_image(
     mode: Literal["dask", "numpy", "coarsen"] = "dask",
 ) -> None:
     """Consolidate the image on disk."""
-    target_paths = image.meta_handler.get_meta().paths
+    target_paths = image._meta_handler.meta.paths
     targets = [
         image._group_handler.get_array(path)
         for path in target_paths
@@ -657,435 +585,3 @@ 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