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

ngio/images/_image.py

@@ -1,7 +1,7 @@
 """Generic class to handle Image-like data in a OME-NGFF file."""
 
-from collections.abc import Sequence
-from typing import Literal
+from collections.abc import Mapping, Sequence
+from typing import Any, Literal
 
 import dask.array as da
 import numpy as np
@@ -12,10 +12,9 @@ from ngio.common import (
     Dimensions,
     InterpolationOrder,
     Roi,
-    RoiPixels,
 )
-from ngio.images._abstract_image import AbstractImage
-from ngio.images._create import create_empty_image_container
+from ngio.common._pyramid import ChunksLike, ShardsLike
+from ngio.images._abstract_image import AbstractImage, abstract_derive
 from ngio.io_pipes import (
     SlicingInputType,
     TransformProtocol,
@@ -24,7 +23,6 @@ from ngio.ome_zarr_meta import (
     ImageMetaHandler,
     NgioImageMeta,
     PixelSize,
-    find_image_meta_handler,
 )
 from ngio.ome_zarr_meta.ngio_specs import (
     Channel,
@@ -37,11 +35,10 @@ from ngio.ome_zarr_meta.ngio_specs import (
     TimeUnits,
 )
 from ngio.utils import (
-    NgioValidationError,
+    NgioValueError,
     StoreOrGroup,
     ZarrGroupHandler,
 )
-from ngio.utils._zarr_utils import find_dimension_separator
 
 
 class ChannelSelectionModel(BaseModel):
@@ -90,7 +87,7 @@ def _check_channel_meta(meta: NgioImageMeta, dimension: Dimensions) -> ChannelsM
         return ChannelsMeta.default_init(labels=c_dim)
 
     if len(meta.channels_meta.channels) != c_dim:
-        raise NgioValidationError(
+        raise NgioValueError(
             "The number of channels does not match the image. "
             f"Expected {len(meta.channels_meta.channels)} channels, got {c_dim}."
         )
@@ -98,7 +95,7 @@ def _check_channel_meta(meta: NgioImageMeta, dimension: Dimensions) -> ChannelsM
     return meta.channels_meta
 
 
-class Image(AbstractImage[ImageMetaHandler]):
+class Image(AbstractImage):
     """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.
@@ -119,15 +116,23 @@ class Image(AbstractImage[ImageMetaHandler]):
 
         """
         if meta_handler is None:
-            meta_handler = find_image_meta_handler(group_handler)
+            meta_handler = ImageMetaHandler(group_handler)
         super().__init__(
             group_handler=group_handler, path=path, meta_handler=meta_handler
         )
 
+    @property
+    def meta_handler(self) -> ImageMetaHandler:
+        """Return the metadata handler."""
+        assert isinstance(self._meta_handler, ImageMetaHandler)
+        return self._meta_handler
+
     @property
     def meta(self) -> NgioImageMeta:
         """Return the metadata."""
-        return self._meta_handler.meta
+        meta = self.meta_handler.get_meta()
+        assert isinstance(meta, NgioImageMeta)
+        return meta
 
     @property
     def channels_meta(self) -> ChannelsMeta:
@@ -187,7 +192,7 @@ class Image(AbstractImage[ImageMetaHandler]):
 
     def get_roi_as_numpy(
         self,
-        roi: Roi | RoiPixels,
+        roi: Roi,
         channel_selection: ChannelSlicingInputType = None,
         axes_order: Sequence[str] | None = None,
         transforms: Sequence[TransformProtocol] | None = None,
@@ -241,7 +246,7 @@ class Image(AbstractImage[ImageMetaHandler]):
 
     def get_roi_as_dask(
         self,
-        roi: Roi | RoiPixels,
+        roi: Roi,
         channel_selection: ChannelSlicingInputType = None,
         axes_order: Sequence[str] | None = None,
         transforms: Sequence[TransformProtocol] | None = None,
@@ -298,7 +303,7 @@ class Image(AbstractImage[ImageMetaHandler]):
 
     def get_roi(
         self,
-        roi: Roi | RoiPixels,
+        roi: Roi,
         channel_selection: ChannelSlicingInputType = None,
         axes_order: Sequence[str] | None = None,
         transforms: Sequence[TransformProtocol] | None = None,
@@ -358,7 +363,7 @@ class Image(AbstractImage[ImageMetaHandler]):
 
     def set_roi(
         self,
-        roi: Roi | RoiPixels,
+        roi: Roi,
         patch: np.ndarray | da.Array,
         channel_selection: ChannelSlicingInputType = None,
         axes_order: Sequence[str] | None = None,
@@ -401,22 +406,22 @@ class ImagesContainer:
     def __init__(self, group_handler: ZarrGroupHandler) -> None:
         """Initialize the LabelGroupHandler."""
         self._group_handler = group_handler
-        self._meta_handler = find_image_meta_handler(group_handler)
+        self._meta_handler = ImageMetaHandler(group_handler)
 
     @property
     def meta(self) -> NgioImageMeta:
         """Return the metadata."""
-        return self._meta_handler.meta
+        return self._meta_handler.get_meta()
 
     @property
     def levels(self) -> int:
         """Return the number of levels in the image."""
-        return self._meta_handler.meta.levels
+        return self._meta_handler.get_meta().levels
 
     @property
     def levels_paths(self) -> list[str]:
         """Return the paths of the levels in the image."""
-        return self._meta_handler.meta.paths
+        return self._meta_handler.get_meta().paths
 
     @property
     def num_channels(self) -> int:
@@ -456,6 +461,15 @@ class ImagesContainer:
             channel_label=channel_label, wavelength_id=wavelength_id
         )
 
+    def _set_channel_meta(
+        self,
+        channels_meta: ChannelsMeta,
+    ) -> None:
+        """Set the channels metadata."""
+        meta = self.meta
+        meta.set_channels_meta(channels_meta)
+        self._meta_handler.update_meta(meta)
+
     def set_channel_meta(
         self,
         labels: Sequence[str | None] | int | None = None,
@@ -492,16 +506,12 @@ class ImagesContainer:
         ref_image = self.get(path=low_res_dataset.path)
 
         if start is not None and end is None:
-            raise NgioValidationError(
-                "If start is provided, end must be provided as well."
-            )
+            raise NgioValueError("If start is provided, end must be provided as well.")
         if end is not None and start is None:
-            raise NgioValidationError(
-                "If end is provided, start must be provided as well."
-            )
+            raise NgioValueError("If end is provided, start must be provided as well.")
 
         if start is not None and percentiles is not None:
-            raise NgioValidationError(
+            raise NgioValueError(
                 "If start and end are provided, percentiles must be None."
             )
 
@@ -513,11 +523,11 @@ class ImagesContainer:
             )
         elif start is not None and end is not None:
             if len(start) != len(end):
-                raise NgioValidationError(
+                raise NgioValueError(
                     "The start and end lists must have the same length."
                 )
             if len(start) != self.num_channels:
-                raise NgioValidationError(
+                raise NgioValueError(
                     "The start and end lists must have the same length as "
                     "the number of channels."
                 )
@@ -541,10 +551,7 @@ class ImagesContainer:
             data_type=ref_image.dtype,
             **omero_kwargs,
         )
-
-        meta = self.meta
-        meta.set_channels_meta(channel_meta)
-        self._meta_handler.write_meta(meta)
+        self._set_channel_meta(channel_meta)
 
     def set_channel_percentiles(
         self,
@@ -553,7 +560,7 @@ class ImagesContainer:
     ) -> None:
         """Update the percentiles of the channels."""
         if self.meta._channels_meta is None:
-            raise NgioValidationError("The channels meta is not initialized.")
+            raise NgioValueError("The channels meta is not initialized.")
 
         low_res_dataset = self.meta.get_lowest_resolution_dataset()
         ref_image = self.get(path=low_res_dataset.path)
@@ -578,7 +585,7 @@ class ImagesContainer:
 
         meta = self.meta
         meta.set_channels_meta(new_meta)
-        self._meta_handler.write_meta(meta)
+        self._meta_handler.update_meta(meta)
 
     def set_axes_unit(
         self,
@@ -593,62 +600,90 @@ class ImagesContainer:
         """
         meta = self.meta
         meta = meta.to_units(space_unit=space_unit, time_unit=time_unit)
-        self._meta_handler.write_meta(meta)
+        self._meta_handler.update_meta(meta)
 
     def derive(
         self,
         store: StoreOrGroup,
         ref_path: str | None = None,
+        # Metadata parameters
         shape: Sequence[int] | None = None,
-        labels: Sequence[str] | None = None,
-        pixel_size: PixelSize | None = None,
-        axes_names: Sequence[str] | None = None,
+        pixelsize: float | tuple[float, float] | None = None,
+        z_spacing: float | None = None,
+        time_spacing: float | None = None,
         name: str | None = None,
-        chunks: Sequence[int] | None = None,
-        dtype: str | None = None,
-        dimension_separator: Literal[".", "/"] | None = None,
-        compressors: CompressorLike | None = None,
+        channels_meta: Sequence[str | Channel] | None = None,
         ngff_version: NgffVersions | None = None,
+        # Zarr Array parameters
+        chunks: ChunksLike = "auto",
+        shards: ShardsLike | None = None,
+        dtype: str = "uint16",
+        dimension_separator: Literal[".", "/"] = "/",
+        compressors: CompressorLike = "auto",
+        extra_array_kwargs: Mapping[str, Any] | None = None,
         overwrite: bool = False,
+        # Deprecated arguments
+        labels: Sequence[str] | None = None,
+        pixel_size: PixelSize | None = None,
     ) -> "ImagesContainer":
         """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:
+            image_container (ImagesContainer): The image container to derive the new
+                image.
             store (StoreOrGroup): The Zarr store or group to create the image in.
-            ref_path (str | None): The path to the reference image in
-                the image container.
+            ref_path (str | None): The path to the reference image in the image
+                container.
             shape (Sequence[int] | None): The shape of the new image.
-            labels (Sequence[str] | None): The labels of the new image.
-            pixel_size (PixelSize | None): The pixel size of the new image.
-            axes_names (Sequence[str] | None): The axes names 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.
-            chunks (Sequence[int] | Literal["auto"]): The chunk shape of the new image.
-            dimension_separator (DIMENSION_SEPARATOR | None): The separator to use for
-                dimensions. If None it will use the same as the reference image.
-            compressors: The compressor to use. If None it will use
-                the same as the reference image.
+            channels_meta (Sequence[str | Channel] | None): The channels metadata
+                of the new image.
+            ngff_version (NgffVersions | None): The NGFF version to use.
+            chunks (Sequence[int] | 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.
-            ngff_version (NgffVersions): The NGFF version to use.
+            dimension_separator (DIMENSION_SEPARATOR | 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.
             overwrite (bool): Whether to overwrite an existing image.
+            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 image
+            ImagesContainer: The new derived image.
+
         """
         return derive_image_container(
             image_container=self,
             store=store,
             ref_path=ref_path,
             shape=shape,
-            labels=labels,
-            pixel_size=pixel_size,
-            axes_names=axes_names,
+            pixelsize=pixelsize,
+            z_spacing=z_spacing,
+            time_spacing=time_spacing,
             name=name,
+            channels_meta=channels_meta,
+            ngff_version=ngff_version,
             chunks=chunks,
+            shards=shards,
             dtype=dtype,
             dimension_separator=dimension_separator,
             compressors=compressors,
-            ngff_version=ngff_version,
+            extra_array_kwargs=extra_array_kwargs,
             overwrite=overwrite,
+            labels=labels,
+            pixel_size=pixel_size,
         )
 
     def get(
@@ -667,7 +702,7 @@ class ImagesContainer:
                 closest pixel size level will be returned.
 
         """
-        dataset = self._meta_handler.meta.get_dataset(
+        dataset = self._meta_handler.get_meta().get_dataset(
             path=path, pixel_size=pixel_size, strict=strict
         )
         return Image(
@@ -720,147 +755,98 @@ def compute_image_percentile(
 
 
 def derive_image_container(
+    *,
     image_container: ImagesContainer,
     store: StoreOrGroup,
     ref_path: str | None = None,
+    # Metadata parameters
     shape: Sequence[int] | None = None,
-    labels: Sequence[str] | None = None,
-    pixel_size: PixelSize | None = None,
-    axes_names: Sequence[str] | None = None,
+    pixelsize: float | tuple[float, float] | None = None,
+    z_spacing: float | None = None,
+    time_spacing: float | None = None,
     name: str | None = None,
-    chunks: Sequence[int] | None = None,
+    channels_policy: Literal["same", "squeeze", "singleton"] | int = "same",
+    channels_meta: Sequence[str | Channel] | None = None,
+    ngff_version: NgffVersions | None = None,
+    # Zarr Array parameters
+    chunks: ChunksLike = "auto",
+    shards: ShardsLike | None = None,
     dtype: str | None = None,
     dimension_separator: Literal[".", "/"] | None = None,
     compressors: CompressorLike | None = None,
-    ngff_version: NgffVersions | None = None,
+    extra_array_kwargs: Mapping[str, Any] | None = None,
     overwrite: bool = False,
+    # Deprecated arguments
+    labels: Sequence[str] | None = None,
+    pixel_size: PixelSize | None = None,
 ) -> ImagesContainer:
-    """Create an empty OME-Zarr image from an existing image.
+    """Derive a new OME-Zarr image container from an existing image.
+
+    If a kwarg is not provided, the value from the reference image will be used.
 
     Args:
-        image_container (ImagesContainer): The image container to derive the new image.
+        image_container (ImagesContainer): The image container to derive the new image
+            from.
         store (StoreOrGroup): The Zarr store or group to create the image in.
         ref_path (str | None): The path to the reference image in the image container.
         shape (Sequence[int] | None): The shape of the new image.
-        labels (Sequence[str] | None): The labels of the new image.
-        pixel_size (PixelSize | None): The pixel size of the new image.
-        axes_names (Sequence[str] | None): The axes names 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.
-        chunks (Sequence[int] | None): The chunk shape of the new image.
-        dimension_separator (DIMENSION_SEPARATOR | None): The separator to use for
-            dimensions. If None it will use the same as the reference image.
-        compressors (CompressorLike | None): The compressors to use. If None it will use
-            the same as the reference image.
-        ngff_version (NgffVersions): The NGFF version to use.
+        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): The chunk shape of the new image. Defaults to "auto".
+        shards (ShardsLike | None): The shard shape of the new image.
         dtype (str | None): The data type of the new image.
-        overwrite (bool): Whether to overwrite an existing 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.
+        overwrite (bool): Whether to overwrite an existing image. Defaults to False.
+        labels (Sequence[str] | None): Deprecated. This argument is deprecated,
+            please use channels_meta instead.
+        pixel_size (PixelSize | None): Deprecated. The pixel size of the new image.
+            This argument is deprecated, please use pixelsize, z_spacing,
+            and time_spacing instead.
 
     Returns:
-        ImagesContainer: The new image
+        ImagesContainer: The new derived image container.
 
     """
-    if ref_path is None:
-        ref_image = image_container.get()
-    else:
-        ref_image = image_container.get(path=ref_path)
-
-    ref_meta = ref_image.meta
-
-    if shape is None:
-        shape = ref_image.shape
-
-    if pixel_size is None:
-        pixel_size = ref_image.pixel_size
-
-    if axes_names is None:
-        axes_names = ref_meta.axes_handler.axes_names
-
-    if len(axes_names) != len(shape):
-        raise NgioValidationError(
-            "The axes names of the new image does not match the reference image."
-            f"Got {axes_names} for shape {shape}."
-        )
-
-    if chunks is None:
-        chunks = ref_image.chunks
-
-    if len(chunks) != len(shape):
-        raise NgioValidationError(
-            "The chunks of the new image does not match the reference image."
-            f"Got {chunks} for shape {shape}."
-        )
-
-    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 ngff_version is None:
-        ngff_version = ref_meta.version
-
-    handler = create_empty_image_container(
+    ref_image = image_container.get(path=ref_path)
+    group_handler = abstract_derive(
+        ref_image=ref_image,
+        meta_type=NgioImageMeta,
         store=store,
         shape=shape,
-        pixelsize=pixel_size.x,
-        z_spacing=pixel_size.z,
-        time_spacing=pixel_size.t,
-        levels=ref_meta.paths,
-        yx_scaling_factor=ref_meta.yx_scaling(),
-        z_scaling_factor=ref_meta.z_scaling(),
-        time_unit=pixel_size.time_unit,
-        space_unit=pixel_size.space_unit,
-        axes_names=axes_names,
+        pixelsize=pixelsize,
+        z_spacing=z_spacing,
+        time_spacing=time_spacing,
         name=name,
+        channels_meta=channels_meta,
+        channels_policy=channels_policy,
+        ngff_version=ngff_version,
         chunks=chunks,
+        shards=shards,
         dtype=dtype,
         dimension_separator=dimension_separator,
         compressors=compressors,
+        extra_array_kwargs=extra_array_kwargs,
         overwrite=overwrite,
-        ngff_version=ngff_version,
-    )
-    image_container = ImagesContainer(handler)
-
-    if ref_image.num_channels == image_container.num_channels:
-        _labels = ref_image.channel_labels
-        wavelength_id = ref_image.wavelength_ids
-
-        channel_meta = ref_image.channels_meta
-        colors = [c.channel_visualisation.color for c in channel_meta.channels]
-        active = [c.channel_visualisation.active for c in channel_meta.channels]
-        start = [c.channel_visualisation.start for c in channel_meta.channels]
-        end = [c.channel_visualisation.end for c in channel_meta.channels]
-    else:
-        _labels = None
-        wavelength_id = None
-        colors = None
-        active = None
-        start = None
-        end = None
-
-    if labels is not None:
-        if len(labels) != image_container.num_channels:
-            raise NgioValidationError(
-                "The number of labels does not match the number of channels."
-            )
-        _labels = labels
-
-    image_container.set_channel_meta(
-        labels=_labels,
-        wavelength_id=wavelength_id,
-        percentiles=None,
-        colors=colors,
-        active=active,
-        start=start,
-        end=end,
+        labels=labels,
+        pixel_size=pixel_size,
     )
-    return image_container
+    return ImagesContainer(group_handler=group_handler)
 
 
 def _parse_str_or_model(
@@ -869,9 +855,9 @@ def _parse_str_or_model(
     """Parse a string or ChannelSelectionModel to an integer channel index."""
     if isinstance(channel_selection, int):
         if channel_selection < 0:
-            raise NgioValidationError("Channel index must be a non-negative integer.")
+            raise NgioValueError("Channel index must be a non-negative integer.")
         if channel_selection >= image.num_channels:
-            raise NgioValidationError(
+            raise NgioValueError(
                 "Channel index must be less than the number "
                 f"of channels ({image.num_channels})."
             )
@@ -889,7 +875,7 @@ def _parse_str_or_model(
             )
         elif channel_selection.mode == "index":
             return int(channel_selection.identifier)
-    raise NgioValidationError(
+    raise NgioValueError(
         "Invalid channel selection type. "
         f"{channel_selection} is of type {type(channel_selection)} ",
         "supported types are str, ChannelSelectionModel, and int.",
@@ -908,7 +894,7 @@ def _parse_channel_selection(
     elif isinstance(channel_selection, Sequence):
         _sequence = [_parse_str_or_model(image, cs) for cs in channel_selection]
         return {"c": _sequence}
-    raise NgioValidationError(
+    raise NgioValueError(
         f"Invalid channel selection type {type(channel_selection)}. "
         "Supported types are int, str, ChannelSelectionModel, and Sequence."
     )
@@ -922,7 +908,7 @@ def add_channel_selection_to_slicing_dict(
     """Add channel selection information to the slicing dictionary."""
     channel_info = _parse_channel_selection(image, channel_selection)
     if "c" in slicing_dict and channel_info:
-        raise NgioValidationError(
+        raise NgioValueError(
             "Both channel_selection and 'c' in slicing_kwargs are provided. "
             "Which channel selection should be used is ambiguous. "
             "Please provide only one."