ngio 0.3.5__py3-none-any.whl → 0.4.0a2__py3-none-any.whl

ngio/images/_image.py

@@ -1,12 +1,21 @@
 """Generic class to handle Image-like data in a OME-NGFF file."""
 
-from collections.abc import Collection
+from collections.abc import Sequence
 from typing import Literal
 
 import dask.array as da
-
-from ngio.common import Dimensions
-from ngio.images._abstract_image import AbstractImage, consolidate_image
+import numpy as np
+from pydantic import BaseModel, model_validator
+
+from ngio.common import (
+    ArrayLike,
+    Dimensions,
+    Roi,
+    RoiPixels,
+    SlicingInputType,
+    TransformProtocol,
+)
+from ngio.images._abstract_image import AbstractImage
 from ngio.images._create import create_empty_image_container
 from ngio.ome_zarr_meta import (
     ImageMetaHandler,
@@ -30,18 +39,55 @@ from ngio.utils import (
 )
 
 
+class ChannelSelectionModel(BaseModel):
+    """Model for channel selection.
+
+    This model is used to select a channel by label, wavelength ID, or index.
+
+    Args:
+        identifier (str): Unique identifier for the channel.
+            This can be a channel label, wavelength ID, or index.
+        mode (Literal["label", "wavelength_id", "index"]): Specifies how to
+            interpret the identifier. Can be "label", "wavelength_id", or
+            "index" (must be an integer).
+
+    """
+
+    mode: Literal["label", "wavelength_id", "index"] = "label"
+    identifier: str
+
+    @model_validator(mode="after")
+    def check_channel_selection(self):
+        if self.mode == "index":
+            try:
+                int(self.identifier)
+            except ValueError as e:
+                raise ValueError(
+                    "Identifier must be an integer when mode is 'index'"
+                ) from e
+        return self
+
+
+ChannelSlicingInputType = (
+    None
+    | int
+    | str
+    | ChannelSelectionModel
+    | Sequence[str | ChannelSelectionModel | int]
+)
+
+
 def _check_channel_meta(meta: NgioImageMeta, dimension: Dimensions) -> ChannelsMeta:
     """Check the channel metadata."""
-    c_dim = dimension.get("c", strict=False)
-    c_dim = 1 if c_dim is None else c_dim
+    c_dim = dimension.get("c", default=1)
 
     if meta.channels_meta is None:
         return ChannelsMeta.default_init(labels=c_dim)
 
-    if len(meta.channels) != c_dim:
+    if len(meta.channels_meta.channels) != c_dim:
         raise NgioValidationError(
             "The number of channels does not match the image. "
-            f"Expected {len(meta.channels)} channels, got {c_dim}."
+            f"Expected {len(meta.channels_meta.channels)} channels, got {c_dim}."
         )
 
     return meta.channels_meta
@@ -72,33 +118,268 @@ class Image(AbstractImage[ImageMetaHandler]):
         super().__init__(
             group_handler=group_handler, path=path, meta_handler=meta_handler
         )
-        self._channels_meta = _check_channel_meta(self.meta, self.dimensions)
 
     @property
     def meta(self) -> NgioImageMeta:
         """Return the metadata."""
         return self._meta_handler.meta
 
+    @property
+    def channels_meta(self) -> ChannelsMeta:
+        """Return the channels metadata."""
+        return _check_channel_meta(self.meta, self.dimensions)
+
     @property
     def channel_labels(self) -> list[str]:
         """Return the channels of the image."""
-        channel_labels = []
-        for c in self._channels_meta.channels:
-            channel_labels.append(c.label)
-        return channel_labels
+        return self.channels_meta.channel_labels
 
     @property
     def wavelength_ids(self) -> list[str | None]:
         """Return the list of wavelength of the image."""
-        wavelength_ids = []
-        for c in self._channels_meta.channels:
-            wavelength_ids.append(c.wavelength_id)
-        return wavelength_ids
+        return self.channels_meta.channel_wavelength_ids
 
     @property
     def num_channels(self) -> int:
         """Return the number of channels."""
-        return len(self._channels_meta.channels)
+        return len(self.channel_labels)
+
+    def get_channel_idx(
+        self, channel_label: str | None = None, wavelength_id: str | None = None
+    ) -> int:
+        """Get the index of a channel by its label or wavelength ID."""
+        return self.channels_meta.get_channel_idx(
+            channel_label=channel_label, wavelength_id=wavelength_id
+        )
+
+    def get_as_numpy(
+        self,
+        channel_selection: ChannelSlicingInputType = None,
+        axes_order: Sequence[str] | None = None,
+        transforms: Sequence[TransformProtocol] | None = None,
+        **slicing_kwargs: slice | int | Sequence[int] | None,
+    ) -> np.ndarray:
+        """Get the image as a numpy array.
+
+        Args:
+            channel_selection: Select a specific channel by label.
+                If None, all channels are returned.
+                Alternatively, you can slice arbitrary channels
+                using the slice_kwargs (c=[0, 2]).
+            axes_order: The order of the axes to return the array.
+            transforms: The transforms to apply to the array.
+            **slicing_kwargs: The slices to get the array.
+
+        Returns:
+            The array of the region of interest.
+        """
+        _slicing_kwargs = add_channel_selection_to_slicing_dict(
+            image=self, channel_selection=channel_selection, slicing_dict=slicing_kwargs
+        )
+        return self._get_as_numpy(
+            axes_order=axes_order, transforms=transforms, **_slicing_kwargs
+        )
+
+    def get_roi_as_numpy(
+        self,
+        roi: Roi | RoiPixels,
+        channel_selection: ChannelSlicingInputType = None,
+        axes_order: Sequence[str] | None = None,
+        transforms: Sequence[TransformProtocol] | None = None,
+        **slicing_kwargs: SlicingInputType,
+    ) -> np.ndarray:
+        """Get the image as a numpy array for a region of interest.
+
+        Args:
+            roi: The region of interest to get the array.
+            channel_selection: Select a what subset of channels to return.
+                If None, all channels are returned.
+            axes_order: The order of the axes to return the array.
+            transforms: The transforms to apply to the array.
+            **slicing_kwargs: The slices to get the array.
+
+        Returns:
+            The array of the region of interest.
+        """
+        _slicing_kwargs = add_channel_selection_to_slicing_dict(
+            image=self, channel_selection=channel_selection, slicing_dict=slicing_kwargs
+        )
+        return self._get_roi_as_numpy(
+            roi=roi, axes_order=axes_order, transforms=transforms, **_slicing_kwargs
+        )
+
+    def get_as_dask(
+        self,
+        channel_selection: ChannelSlicingInputType = None,
+        axes_order: Sequence[str] | None = None,
+        transforms: Sequence[TransformProtocol] | None = None,
+        **slicing_kwargs: SlicingInputType,
+    ) -> da.Array:
+        """Get the image as a dask array.
+
+        Args:
+            channel_selection: Select a what subset of channels to return.
+                If None, all channels are returned.
+            axes_order: The order of the axes to return the array.
+            transforms: The transforms to apply to the array.
+            **slicing_kwargs: The slices to get the array.
+
+        Returns:
+            The dask array of the region of interest.
+        """
+        _slicing_kwargs = add_channel_selection_to_slicing_dict(
+            image=self, channel_selection=channel_selection, slicing_dict=slicing_kwargs
+        )
+        return self._get_as_dask(
+            axes_order=axes_order, transforms=transforms, **_slicing_kwargs
+        )
+
+    def get_roi_as_dask(
+        self,
+        roi: Roi | RoiPixels,
+        channel_selection: ChannelSlicingInputType = None,
+        axes_order: Sequence[str] | None = None,
+        transforms: Sequence[TransformProtocol] | None = None,
+        **slicing_kwargs: SlicingInputType,
+    ) -> da.Array:
+        """Get the image as a dask array for a region of interest.
+
+        Args:
+            roi: The region of interest to get the array.
+            channel_selection: Select a what subset of channels to return.
+                If None, all channels are returned.
+            axes_order: The order of the axes to return the array.
+            transforms: The transforms to apply to the array.
+            **slicing_kwargs: The slices to get the array.
+
+        Returns:
+            The dask array of the region of interest.
+        """
+        _slicing_kwargs = add_channel_selection_to_slicing_dict(
+            image=self, channel_selection=channel_selection, slicing_dict=slicing_kwargs
+        )
+        return self._get_roi_as_dask(
+            roi=roi, axes_order=axes_order, transforms=transforms, **_slicing_kwargs
+        )
+
+    def get_array(
+        self,
+        channel_selection: ChannelSlicingInputType = None,
+        axes_order: Sequence[str] | None = None,
+        transforms: Sequence[TransformProtocol] | None = None,
+        mode: Literal["numpy", "dask"] = "numpy",
+        **slicing_kwargs: SlicingInputType,
+    ) -> ArrayLike:
+        """Get the image as a zarr array.
+
+        Args:
+            channel_selection: Select a what subset of channels to return.
+                If None, all channels are returned.
+            axes_order: The order of the axes to return the array.
+            transforms: The transforms to apply to the array.
+            mode: The object type to return.
+                Can be "dask", "numpy".
+            **slicing_kwargs: The slices to get the array.
+
+        Returns:
+            The zarr array of the region of interest.
+        """
+        _slicing_kwargs = add_channel_selection_to_slicing_dict(
+            image=self, channel_selection=channel_selection, slicing_dict=slicing_kwargs
+        )
+        return self._get_array(
+            axes_order=axes_order, mode=mode, transforms=transforms, **_slicing_kwargs
+        )
+
+    def get_roi(
+        self,
+        roi: Roi | RoiPixels,
+        channel_selection: ChannelSlicingInputType = None,
+        axes_order: Sequence[str] | None = None,
+        transforms: Sequence[TransformProtocol] | None = None,
+        mode: Literal["numpy", "dask"] = "numpy",
+        **slicing_kwargs: SlicingInputType,
+    ) -> ArrayLike:
+        """Get the image as a zarr array for a region of interest.
+
+        Args:
+            roi: The region of interest to get the array.
+            channel_selection: Select a what subset of channels to return.
+                If None, all channels are returned.
+            axes_order: The order of the axes to return the array.
+            transforms: The transforms to apply to the array.
+            mode: The object type to return.
+                Can be "dask", "numpy".
+            **slicing_kwargs: The slices to get the array.
+
+        Returns:
+            The zarr array of the region of interest.
+        """
+        _slicing_kwargs = add_channel_selection_to_slicing_dict(
+            image=self, channel_selection=channel_selection, slicing_dict=slicing_kwargs
+        )
+        return self._get_roi(
+            roi=roi,
+            axes_order=axes_order,
+            mode=mode,
+            transforms=transforms,
+            **_slicing_kwargs,
+        )
+
+    def set_array(
+        self,
+        patch: ArrayLike,
+        channel_selection: ChannelSlicingInputType = None,
+        axes_order: Sequence[str] | None = None,
+        transforms: Sequence[TransformProtocol] | None = None,
+        **slicing_kwargs: SlicingInputType,
+    ) -> None:
+        """Set the image array.
+
+        Args:
+            patch: The array to set.
+            channel_selection: Select a what subset of channels to return.
+                If None, all channels are set.
+            axes_order: The order of the axes to set the array.
+            transforms: The transforms to apply to the array.
+            **slicing_kwargs: The slices to set the array.
+        """
+        _slicing_kwargs = add_channel_selection_to_slicing_dict(
+            image=self, channel_selection=channel_selection, slicing_dict=slicing_kwargs
+        )
+        self._set_array(
+            patch=patch, axes_order=axes_order, transforms=transforms, **_slicing_kwargs
+        )
+
+    def set_roi(
+        self,
+        roi: Roi | RoiPixels,
+        patch: ArrayLike,
+        channel_selection: ChannelSlicingInputType = None,
+        axes_order: Sequence[str] | None = None,
+        transforms: Sequence[TransformProtocol] | None = None,
+        **slicing_kwargs: SlicingInputType,
+    ) -> None:
+        """Set the image array for a region of interest.
+
+        Args:
+            roi: The region of interest to set the array.
+            patch: The array to set.
+            channel_selection: Select a what subset of channels to return.
+            axes_order: The order of the axes to set the array.
+            transforms: The transforms to apply to the array.
+            **slicing_kwargs: The slices to set the array.
+        """
+        _slicing_kwargs = add_channel_selection_to_slicing_dict(
+            image=self, channel_selection=channel_selection, slicing_dict=slicing_kwargs
+        )
+        self._set_roi(
+            roi=roi,
+            patch=patch,
+            axes_order=axes_order,
+            transforms=transforms,
+            **_slicing_kwargs,
+        )
 
     def consolidate(
         self,
@@ -106,7 +387,7 @@ class Image(AbstractImage[ImageMetaHandler]):
         mode: Literal["dask", "numpy", "coarsen"] = "dask",
     ) -> None:
         """Consolidate the label on disk."""
-        consolidate_image(self, order=order, mode=mode)
+        self._consolidate(order=order, mode=mode)
 
 
 class ImagesContainer:
@@ -150,35 +431,55 @@ class ImagesContainer:
         image = self.get()
         return image.wavelength_ids
 
+    def get_channel_idx(
+        self, channel_label: str | None = None, wavelength_id: str | None = None
+    ) -> int:
+        """Get the index of a channel by label or wavelength ID.
+
+        Args:
+            channel_label (str | None): The label of the channel.
+                If None a wavelength ID must be provided.
+            wavelength_id (str | None): The wavelength ID of the channel.
+                If None a channel label must be provided.
+
+        Returns:
+            int: The index of the channel.
+
+        """
+        image = self.get()
+        return image.get_channel_idx(
+            channel_label=channel_label, wavelength_id=wavelength_id
+        )
+
     def set_channel_meta(
         self,
-        labels: Collection[str | None] | int | None = None,
-        wavelength_id: Collection[str | None] | None = None,
-        start: Collection[float | None] | None = None,
-        end: Collection[float | None] | None = None,
+        labels: Sequence[str | None] | int | None = None,
+        wavelength_id: Sequence[str | None] | None = None,
+        start: Sequence[float | None] | None = None,
+        end: Sequence[float | None] | None = None,
         percentiles: tuple[float, float] | None = None,
-        colors: Collection[str | None] | None = None,
-        active: Collection[bool | None] | None = None,
+        colors: Sequence[str | None] | None = None,
+        active: Sequence[bool | None] | None = None,
         **omero_kwargs: dict,
     ) -> None:
         """Create a ChannelsMeta object with the default unit.
 
         Args:
-            labels(Collection[str | None] | int): The list of channels names
+            labels(Sequence[str | None] | int): The list of channels names
                 in the image. If an integer is provided, the channels will
                 be named "channel_i".
-            wavelength_id(Collection[str | None]): The wavelength ID of the channel.
+            wavelength_id(Sequence[str | None]): The wavelength ID of the channel.
                 If None, the wavelength ID will be the same as the channel name.
-            start(Collection[float | None]): The start value for each channel.
+            start(Sequence[float | None]): The start value for each channel.
                 If None, the start value will be computed from the image.
-            end(Collection[float | None]): The end value for each channel.
+            end(Sequence[float | None]): The end value for each channel.
                 If None, the end value will be computed from the image.
             percentiles(tuple[float, float] | None): The start and end
                 percentiles for each channel. If None, the percentiles will
                 not be computed.
-            colors(Collection[str | None]): The list of colors for the
+            colors(Sequence[str | None]): The list of colors for the
                 channels. If None, the colors will be random.
-            active (Collection[bool | None]): Whether the channel should
+            active (Sequence[bool | None]): Whether the channel should
                 be shown by default.
             omero_kwargs(dict): Extra fields to store in the omero attributes.
         """
@@ -293,12 +594,12 @@ class ImagesContainer:
         self,
         store: StoreOrGroup,
         ref_path: str | None = None,
-        shape: Collection[int] | None = None,
-        labels: Collection[str] | None = None,
+        shape: Sequence[int] | None = None,
+        labels: Sequence[str] | None = None,
         pixel_size: PixelSize | None = None,
-        axes_names: Collection[str] | None = None,
+        axes_names: Sequence[str] | None = None,
         name: str | None = None,
-        chunks: Collection[int] | None = None,
+        chunks: Sequence[int] | None = None,
         dtype: str | None = None,
         overwrite: bool = False,
     ) -> "ImagesContainer":
@@ -308,12 +609,12 @@ class ImagesContainer:
             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 (Collection[int] | None): The shape of the new image.
-            labels (Collection[str] | None): The labels of the new image.
+            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 (Collection[str] | None): The axes names of the new image.
+            axes_names (Sequence[str] | None): The axes names of the new image.
             name (str | None): The name of the new image.
-            chunks (Collection[int] | None): The chunk shape of the new image.
+            chunks (Sequence[int] | None): The chunk shape of the new image.
             dtype (str | None): The data type of the new image.
             overwrite (bool): Whether to overwrite an existing image.
 
@@ -378,11 +679,10 @@ def compute_image_percentile(
     starts, ends = [], []
     for c in range(image.num_channels):
         if image.num_channels == 1:
-            data = image.get_array(mode="dask")
+            data = image.get_as_dask()
         else:
-            data = image.get_array(c=c, mode="dask")
+            data = image.get_as_dask(c=c)
 
-        assert isinstance(data, da.Array), "Data must be a Dask array."
         data = da.ravel(data)
         # remove all the zeros
         mask = data > 1e-16
@@ -396,7 +696,7 @@ def compute_image_percentile(
         # compute the percentiles
         _s_perc, _e_perc = da.percentile(
             data, [start_percentile, end_percentile], method="nearest"
-        ).compute()  # type: ignore
+        ).compute()  # type: ignore (return type is a tuple of floats)
 
         starts.append(float(_s_perc))
         ends.append(float(_e_perc))
@@ -407,12 +707,12 @@ def derive_image_container(
     image_container: ImagesContainer,
     store: StoreOrGroup,
     ref_path: str | None = None,
-    shape: Collection[int] | None = None,
-    labels: Collection[str] | None = None,
+    shape: Sequence[int] | None = None,
+    labels: Sequence[str] | None = None,
     pixel_size: PixelSize | None = None,
-    axes_names: Collection[str] | None = None,
+    axes_names: Sequence[str] | None = None,
     name: str | None = None,
-    chunks: Collection[int] | None = None,
+    chunks: Sequence[int] | None = None,
     dtype: str | None = None,
     overwrite: bool = False,
 ) -> ImagesContainer:
@@ -422,12 +722,12 @@ def derive_image_container(
         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.
-        shape (Collection[int] | None): The shape of the new image.
-        labels (Collection[str] | None): The labels of the new image.
+        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 (Collection[str] | None): The axes names of the new image.
+        axes_names (Sequence[str] | None): The axes names of the new image.
         name (str | None): The name of the new image.
-        chunks (Collection[int] | None): The chunk shape of the new image.
+        chunks (Sequence[int] | None): The chunk shape of the new image.
         dtype (str | None): The data type of the new image.
         overwrite (bool): Whether to overwrite an existing image.
 
@@ -449,7 +749,7 @@ def derive_image_container(
         pixel_size = ref_image.pixel_size
 
     if axes_names is None:
-        axes_names = ref_meta.axes_mapper.on_disk_axes_names
+        axes_names = ref_meta.axes_mapper.axes_names
 
     if len(axes_names) != len(shape):
         raise NgioValidationError(
@@ -495,16 +795,11 @@ def derive_image_container(
         _labels = ref_image.channel_labels
         wavelength_id = ref_image.wavelength_ids
 
-        colors = [
-            c.channel_visualisation.color for c in ref_image._channels_meta.channels
-        ]
-        active = [
-            c.channel_visualisation.active for c in ref_image._channels_meta.channels
-        ]
-        start = [
-            c.channel_visualisation.start for c in ref_image._channels_meta.channels
-        ]
-        end = [c.channel_visualisation.end for c in ref_image._channels_meta.channels]
+        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
@@ -530,3 +825,71 @@ def derive_image_container(
         end=end,
     )
     return image_container
+
+
+def _parse_str_or_model(
+    image: Image, channel_selection: int | str | ChannelSelectionModel
+) -> int:
+    """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.")
+        if channel_selection >= image.num_channels:
+            raise NgioValidationError(
+                "Channel index must be less than the number "
+                f"of channels ({image.num_channels})."
+            )
+        return channel_selection
+    elif isinstance(channel_selection, str):
+        return image.get_channel_idx(channel_label=channel_selection)
+    elif isinstance(channel_selection, ChannelSelectionModel):
+        if channel_selection.mode == "label":
+            return image.get_channel_idx(
+                channel_label=str(channel_selection.identifier)
+            )
+        elif channel_selection.mode == "wavelength_id":
+            return image.get_channel_idx(
+                channel_label=str(channel_selection.identifier)
+            )
+        elif channel_selection.mode == "index":
+            return int(channel_selection.identifier)
+    raise NgioValidationError(
+        "Invalid channel selection type. "
+        f"{channel_selection} is of type {type(channel_selection)} ",
+        "supported types are str, ChannelSelectionModel, and int.",
+    )
+
+
+def _parse_channel_selection(
+    image: Image, channel_selection: ChannelSlicingInputType
+) -> dict[str, SlicingInputType]:
+    """Parse the channel selection input into a list of channel indices."""
+    if channel_selection is None:
+        return {}
+    if isinstance(channel_selection, int | str | ChannelSelectionModel):
+        channel_index = _parse_str_or_model(image, channel_selection)
+        return {"c": channel_index}
+    elif isinstance(channel_selection, Sequence):
+        _sequence = [_parse_str_or_model(image, cs) for cs in channel_selection]
+        return {"c": _sequence}
+    raise NgioValidationError(
+        f"Invalid channel selection type {type(channel_selection)}. "
+        "Supported types are int, str, ChannelSelectionModel, and Sequence."
+    )
+
+
+def add_channel_selection_to_slicing_dict(
+    image: Image,
+    channel_selection: ChannelSlicingInputType,
+    slicing_dict: dict[str, SlicingInputType],
+) -> dict[str, SlicingInputType]:
+    """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(
+            "Both channel_selection and 'c' in slicing_kwargs are provided. "
+            "Which channel selection should be used is ambiguous. "
+            "Please provide only one."
+        )
+    slicing_dict = slicing_dict | channel_info
+    return slicing_dict