ngio 0.3.4__py3-none-any.whl → 0.4.0a1__py3-none-any.whl

ngio/images/_abstract_image.py

@@ -1,8 +1,10 @@
 """Generic class to handle Image-like data in a OME-NGFF file."""
 
-from collections.abc import Collection, Iterable
+from collections.abc import Sequence
 from typing import Generic, Literal, TypeVar
 
+import dask.array as da
+import numpy as np
 import zarr
 
 from ngio.common import (
@@ -10,10 +12,17 @@ from ngio.common import (
     Dimensions,
     Roi,
     RoiPixels,
+    SlicingInputType,
+    TransformProtocol,
+    build_dask_getter,
+    build_dask_setter,
+    build_numpy_getter,
+    build_numpy_setter,
+    build_roi_dask_getter,
+    build_roi_dask_setter,
+    build_roi_numpy_getter,
+    build_roi_numpy_setter,
     consolidate_pyramid,
-    get_pipe,
-    roi_to_slice_kwargs,
-    set_pipe,
 )
 from ngio.ome_zarr_meta import (
     AxesMapper,
@@ -165,80 +174,223 @@ class AbstractImage(Generic[_image_handler]):
         self.axes_mapper.get_index("x")
         return self.dimensions.has_axis(axis)
 
-    def get_array(
+    def _get_as_numpy(
         self,
-        axes_order: Collection[str] | None = None,
-        mode: Literal["numpy", "dask", "delayed"] = "numpy",
-        **slice_kwargs: slice | int | Iterable[int],
-    ) -> ArrayLike:
-        """Get a slice of the image.
+        axes_order: Sequence[str] | None = None,
+        transforms: Sequence[TransformProtocol] | None = None,
+        **slicing_kwargs: SlicingInputType,
+    ) -> np.ndarray:
+        """Get the image as a numpy array.
 
         Args:
             axes_order: The order of the axes to return the array.
-            mode: The mode to return the array.
-            **slice_kwargs: The slices to get the array.
+            transforms: The transforms to apply to the array.
+            **slicing_kwargs: The slices to get the array.
 
         Returns:
             The array of the region of interest.
         """
-        return get_pipe(
-            array=self.zarr_array,
+        numpy_getter = build_numpy_getter(
+            zarr_array=self.zarr_array,
             dimensions=self.dimensions,
             axes_order=axes_order,
-            mode=mode,
-            **slice_kwargs,
+            transforms=transforms,
+            slicing_dict=slicing_kwargs,
         )
+        return numpy_getter()
+
+    def _get_roi_as_numpy(
+        self,
+        roi: Roi | RoiPixels,
+        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.
+            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.
+        """
+        numpy_roi_getter = build_roi_numpy_getter(
+            zarr_array=self.zarr_array,
+            dimensions=self.dimensions,
+            roi=roi,
+            pixel_size=self.pixel_size,
+            axes_order=axes_order,
+            transforms=transforms,
+            slicing_dict=slicing_kwargs,
+        )
+        return numpy_roi_getter()
+
+    def _get_as_dask(
+        self,
+        axes_order: Sequence[str] | None = None,
+        transforms: Sequence[TransformProtocol] | None = None,
+        **slicing_kwargs: SlicingInputType,
+    ) -> da.Array:
+        """Get the image as a dask array.
+
+        Args:
+            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.
+        """
+        dask_getter = build_dask_getter(
+            zarr_array=self.zarr_array,
+            dimensions=self.dimensions,
+            axes_order=axes_order,
+            transforms=transforms,
+            slicing_dict=slicing_kwargs,
+        )
+        return dask_getter()
+
+    def _get_roi_as_dask(
+        self,
+        roi: Roi | RoiPixels,
+        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.
+            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.
+        """
+        roi_dask_getter = build_roi_dask_getter(
+            zarr_array=self.zarr_array,
+            dimensions=self.dimensions,
+            roi=roi,
+            pixel_size=self.pixel_size,
+            axes_order=axes_order,
+            transforms=transforms,
+            slicing_dict=slicing_kwargs,
+        )
+        return roi_dask_getter()
+
+    def _get_array(
+        self,
+        axes_order: Sequence[str] | None = None,
+        transforms: Sequence[TransformProtocol] | None = None,
+        mode: Literal["numpy", "dask"] = "numpy",
+        **slicing_kwargs: SlicingInputType,
+    ) -> ArrayLike:
+        """Get a slice of the image.
+
+        Args:
+            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.
 
-    def get_roi(
+        Returns:
+            The array of the region of interest.
+        """
+        if mode == "numpy":
+            return self._get_as_numpy(
+                axes_order=axes_order, transforms=transforms, **slicing_kwargs
+            )
+        elif mode == "dask":
+            return self._get_as_dask(
+                axes_order=axes_order, transforms=transforms, **slicing_kwargs
+            )
+        else:
+            raise ValueError(
+                f"Unknown mode: {mode}. Choose from 'numpy', 'dask', or 'delayed'."
+            )
+
+    def _get_roi(
         self,
-        roi: Roi,
-        axes_order: Collection[str] | None = None,
-        mode: Literal["numpy", "dask", "delayed"] = "numpy",
-        **slice_kwargs: slice | int | Iterable[int],
+        roi: Roi | RoiPixels,
+        axes_order: Sequence[str] | None = None,
+        transforms: Sequence[TransformProtocol] | None = None,
+        mode: Literal["numpy", "dask"] = "numpy",
+        **slice_kwargs: SlicingInputType,
     ) -> ArrayLike:
         """Get a slice of the image.
 
         Args:
             roi: The region of interest to get the array.
             axes_order: The order of the axes to return the array.
+            transforms: The transforms to apply to the array.
             mode: The mode to return the array.
+                Can be "dask", "numpy".
             **slice_kwargs: The slices to get the array.
 
         Returns:
             The array of the region of interest.
         """
-        return get_roi_pipe(
-            image=self, roi=roi, axes_order=axes_order, mode=mode, **slice_kwargs
-        )
-
-    def set_array(
+        if mode == "numpy":
+            return self._get_roi_as_numpy(
+                roi=roi, axes_order=axes_order, transforms=transforms, **slice_kwargs
+            )
+        elif mode == "dask":
+            return self._get_roi_as_dask(
+                roi=roi, axes_order=axes_order, transforms=transforms, **slice_kwargs
+            )
+        else:
+            raise ValueError(
+                f"Unknown mode: {mode}. Choose from 'numpy', 'dask', or 'delayed'."
+            )
+
+    def _set_array(
         self,
         patch: ArrayLike,
-        axes_order: Collection[str] | None = None,
-        **slice_kwargs: slice | int | Iterable[int],
+        axes_order: Sequence[str] | None = None,
+        transforms: Sequence[TransformProtocol] | None = None,
+        **slicing_kwargs: SlicingInputType,
     ) -> None:
         """Set a slice of the image.
 
         Args:
             patch: The patch to set.
             axes_order: The order of the axes to set the patch.
-            **slice_kwargs: The slices to set the patch.
+            transforms: The transforms to apply to the patch.
+            **slicing_kwargs: The slices to set the patch.
 
         """
-        set_pipe(
-            array=self.zarr_array,
-            patch=patch,
-            dimensions=self.dimensions,
-            axes_order=axes_order,
-            **slice_kwargs,
-        )
-
-    def set_roi(
+        if isinstance(patch, np.ndarray):
+            numpy_setter = build_numpy_setter(
+                zarr_array=self.zarr_array,
+                dimensions=self.dimensions,
+                axes_order=axes_order,
+                transforms=transforms,
+                slicing_dict=slicing_kwargs,
+            )
+            numpy_setter(patch)
+
+        elif isinstance(patch, da.Array):
+            dask_setter = build_dask_setter(
+                zarr_array=self.zarr_array,
+                dimensions=self.dimensions,
+                axes_order=axes_order,
+                transforms=transforms,
+                slicing_dict=slicing_kwargs,
+            )
+            dask_setter(patch)
+        else:
+            raise TypeError(
+                f"Unsupported patch type: {type(patch)}. "
+                "Supported types are: "
+                "numpy.ndarray, dask.array.Array."
+            )
+
+    def _set_roi(
         self,
-        roi: Roi,
+        roi: Roi | RoiPixels,
         patch: ArrayLike,
-        axes_order: Collection[str] | None = None,
-        **slice_kwargs: slice | int | Iterable[int],
+        axes_order: Sequence[str] | None = None,
+        transforms: Sequence[TransformProtocol] | None = None,
+        **slicing_kwargs: SlicingInputType,
     ) -> None:
         """Set a slice of the image.
 
@@ -246,12 +398,52 @@ class AbstractImage(Generic[_image_handler]):
             roi: The region of interest to set the patch.
             patch: The patch to set.
             axes_order: The order of the axes to set the patch.
-            **slice_kwargs: The slices to set the patch.
+            transforms: The transforms to apply to the patch.
+            **slicing_kwargs: The slices to set the patch.
 
         """
-        return set_roi_pipe(
-            image=self, roi=roi, patch=patch, axes_order=axes_order, **slice_kwargs
-        )
+        if isinstance(patch, np.ndarray):
+            roi_numpy_setter = build_roi_numpy_setter(
+                zarr_array=self.zarr_array,
+                dimensions=self.dimensions,
+                roi=roi,
+                pixel_size=self.pixel_size,
+                axes_order=axes_order,
+                transforms=transforms,
+                slicing_dict=slicing_kwargs,
+            )
+            roi_numpy_setter(patch)
+
+        elif isinstance(patch, da.Array):
+            roi_dask_setter = build_roi_dask_setter(
+                zarr_array=self.zarr_array,
+                dimensions=self.dimensions,
+                roi=roi,
+                pixel_size=self.pixel_size,
+                axes_order=axes_order,
+                transforms=transforms,
+                slicing_dict=slicing_kwargs,
+            )
+            roi_dask_setter(patch)
+        else:
+            raise TypeError(
+                f"Unsupported patch type: {type(patch)}. "
+                "Supported types are: "
+                "numpy.ndarray, dask.array.Array."
+            )
+
+    def _consolidate(
+        self,
+        order: Literal[0, 1, 2] = 1,
+        mode: Literal["dask", "numpy", "coarsen"] = "dask",
+    ) -> None:
+        """Consolidate the image on disk.
+
+        Args:
+            order: The order of the consolidation.
+            mode: The mode of the consolidation.
+        """
+        consolidate_image(image=self, order=order, mode=mode)
 
     def build_image_roi_table(self, name: str = "image") -> RoiTable:
         """Build the ROI table for an image."""
@@ -275,86 +467,24 @@ def consolidate_image(
     )
 
 
-def get_roi_pipe(
-    image: AbstractImage,
-    roi: Roi,
-    axes_order: Collection[str] | None = None,
-    mode: Literal["numpy", "dask", "delayed"] = "numpy",
-    **slice_kwargs: slice | int | Iterable[int],
-) -> ArrayLike:
-    """Get a slice of the image.
-
-    Args:
-        image: The image to get the ROI.
-        roi: The region of interest to get the array.
-        axes_order: The order of the axes to return the array.
-        mode: The mode to return the array.
-        **slice_kwargs: The slices to get the array.
-
-    Returns:
-        The array of the region of interest.
-    """
-    slice_kwargs = roi_to_slice_kwargs(
-        roi=roi,
-        pixel_size=image.pixel_size,
-        dimensions=image.dimensions,
-        **slice_kwargs,
-    )
-    return get_pipe(
-        array=image.zarr_array,
-        dimensions=image.dimensions,
-        axes_order=axes_order,
-        mode=mode,
-        **slice_kwargs,
-    )
-
-
-def set_roi_pipe(
-    image: AbstractImage,
-    roi: Roi,
-    patch: ArrayLike,
-    axes_order: Collection[str] | None = None,
-    **slice_kwargs: slice | int | Iterable[int],
-) -> None:
-    """Set a slice of the image.
-
-    Args:
-        image: The image to set the ROI.
-        roi: The region of interest to set the patch.
-        patch: The patch to set.
-        axes_order: The order of the axes to set the patch.
-        **slice_kwargs: The slices to set the patch.
-
-    """
-    slice_kwargs = roi_to_slice_kwargs(
-        roi=roi,
-        pixel_size=image.pixel_size,
-        dimensions=image.dimensions,
-        **slice_kwargs,
-    )
-    set_pipe(
-        array=image.zarr_array,
-        patch=patch,
-        dimensions=image.dimensions,
-        axes_order=axes_order,
-        **slice_kwargs,
-    )
-
-
 def build_image_roi_table(image: AbstractImage, name: str = "image") -> RoiTable:
     """Build the ROI table for an image."""
-    dim_z, dim_y, dim_x = (
-        image.dimensions.get("z", strict=False),
-        image.dimensions.get("y"),
-        image.dimensions.get("x"),
-    )
+    dim_x = image.dimensions.get("x")
+    dim_y = image.dimensions.get("y")
+    assert dim_x is not None and dim_y is not None
+    dim_z = image.dimensions.get("z")
+    z = None if dim_z is None else 0
+    dim_t = image.dimensions.get("t")
+    t = None if dim_t is None else 0
     image_roi = RoiPixels(
         name=name,
         x=0,
         y=0,
-        z=0,
+        z=z,
+        t=t,
         x_length=dim_x,
         y_length=dim_y,
         z_length=dim_z,
+        t_length=dim_t,
     )
     return RoiTable(rois=[image_roi.to_roi(pixel_size=image.pixel_size)])

ngio/images/_create.py

@@ -1,6 +1,6 @@
 """Utility functions for working with OME-Zarr images."""
 
-from collections.abc import Collection
+from collections.abc import Sequence
 from typing import TypeVar
 
 from ngio.common._pyramid import init_empty_pyramid
@@ -29,7 +29,7 @@ _image_or_label_meta = TypeVar("_image_or_label_meta", NgioImageMeta, NgioLabelM
 def _init_generic_meta(
     meta_type: type[_image_or_label_meta],
     pixelsize: float,
-    axes_names: Collection[str],
+    axes_names: Sequence[str],
     z_spacing: float = 1.0,
     time_spacing: float = 1.0,
     levels: int | list[str] = 5,
@@ -80,7 +80,7 @@ def _init_generic_meta(
 
 def create_empty_label_container(
     store: StoreOrGroup,
-    shape: Collection[int],
+    shape: Sequence[int],
     pixelsize: float,
     z_spacing: float = 1.0,
     time_spacing: float = 1.0,
@@ -89,10 +89,10 @@ def create_empty_label_container(
     z_scaling_factor: float = 1.0,
     space_unit: SpaceUnits | str | None = DefaultSpaceUnit,
     time_unit: TimeUnits | str | None = DefaultTimeUnit,
-    axes_names: Collection[str] | None = None,
+    axes_names: Sequence[str] | None = None,
     name: str | None = None,
-    chunks: Collection[int] | None = None,
-    dtype: str = "uint16",
+    chunks: Sequence[int] | None = None,
+    dtype: str = "uint32",
     overwrite: bool = False,
     version: NgffVersions = DefaultNgffVersion,
 ) -> ZarrGroupHandler:
@@ -100,7 +100,7 @@ def create_empty_label_container(
 
     Args:
         store (StoreOrGroup): The Zarr store or group to create the image in.
-        shape (Collection[int]): The shape of the image.
+        shape (Sequence[int]): The shape of the image.
         pixelsize (float): The pixel size in x and y dimensions.
         z_spacing (float, optional): The spacing between z slices. Defaults to 1.0.
         time_spacing (float, optional): The spacing between time points.
@@ -115,10 +115,10 @@ def create_empty_label_container(
             DefaultSpaceUnit.
         time_unit (TimeUnits, optional): The unit of time. Defaults to
             DefaultTimeUnit.
-        axes_names (Collection[str] | None, optional): The names of the axes.
+        axes_names (Sequence[str] | None, optional): The names of the axes.
             If None the canonical names are used. Defaults to None.
         name (str | None, optional): The name of the image. Defaults to None.
-        chunks (Collection[int] | None, optional): The chunk shape. If None the shape
+        chunks (Sequence[int] | None, optional): The chunk shape. If None the shape
             is used. Defaults to None.
         dtype (str, optional): The data type of the image. Defaults to "uint16".
         overwrite (bool, optional): Whether to overwrite an existing image.
@@ -171,7 +171,7 @@ def create_empty_label_container(
 
 def create_empty_image_container(
     store: StoreOrGroup,
-    shape: Collection[int],
+    shape: Sequence[int],
     pixelsize: float,
     z_spacing: float = 1.0,
     time_spacing: float = 1.0,
@@ -180,9 +180,9 @@ def create_empty_image_container(
     z_scaling_factor: float = 1.0,
     space_unit: SpaceUnits | str | None = DefaultSpaceUnit,
     time_unit: TimeUnits | str | None = DefaultTimeUnit,
-    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 = "uint16",
     overwrite: bool = False,
     version: NgffVersions = DefaultNgffVersion,
@@ -191,7 +191,7 @@ def create_empty_image_container(
 
     Args:
         store (StoreOrGroup): The Zarr store or group to create the image in.
-        shape (Collection[int]): The shape of the image.
+        shape (Sequence[int]): The shape of the image.
         pixelsize (float): The pixel size in x and y dimensions.
         z_spacing (float, optional): The spacing between z slices. Defaults to 1.0.
         time_spacing (float, optional): The spacing between time points.
@@ -206,10 +206,10 @@ def create_empty_image_container(
             DefaultSpaceUnit.
         time_unit (TimeUnits, optional): The unit of time. Defaults to
             DefaultTimeUnit.
-        axes_names (Collection[str] | None, optional): The names of the axes.
+        axes_names (Sequence[str] | None, optional): The names of the axes.
             If None the canonical names are used. Defaults to None.
         name (str | None, optional): The name of the image. Defaults to None.
-        chunks (Collection[int] | None, optional): The chunk shape. If None the shape
+        chunks (Sequence[int] | None, optional): The chunk shape. If None the shape
             is used. Defaults to None.
         dtype (str, optional): The data type of the image. Defaults to "uint16".
         overwrite (bool, optional): Whether to overwrite an existing image.

ngio/images/_create_synt_container.py

@@ -0,0 +1,128 @@
+"""Abstract class for handling OME-NGFF images."""
+
+from collections.abc import Sequence
+
+import numpy as np
+import PIL.Image
+
+from ngio.common._synt_images_utils import fit_to_shape
+from ngio.images._ome_zarr_container import OmeZarrContainer, create_ome_zarr_from_array
+from ngio.ome_zarr_meta.ngio_specs import (
+    DefaultNgffVersion,
+    NgffVersions,
+)
+from ngio.resources import AVAILABLE_SAMPLES, SampleInfo, get_sample_info
+from ngio.tables import (
+    DefaultTableBackend,
+    TableBackend,
+)
+from ngio.utils import (
+    StoreOrGroup,
+)
+
+
+def create_synthetic_ome_zarr(
+    store: StoreOrGroup,
+    shape: Sequence[int],
+    reference_sample: AVAILABLE_SAMPLES | SampleInfo = "Cardiomyocyte",
+    levels: int | list[str] = 5,
+    xy_scaling_factor: float = 2,
+    z_scaling_factor: float = 1.0,
+    axes_names: Sequence[str] | None = None,
+    chunks: Sequence[int] | None = None,
+    channel_labels: list[str] | None = None,
+    channel_wavelengths: list[str] | None = None,
+    channel_colors: Sequence[str] | None = None,
+    channel_active: Sequence[bool] | None = None,
+    table_backend: TableBackend = DefaultTableBackend,
+    overwrite: bool = False,
+    version: NgffVersions = DefaultNgffVersion,
+) -> OmeZarrContainer:
+    """Create an empty OME-Zarr image with the given shape and metadata.
+
+    Args:
+        store (StoreOrGroup): The Zarr store or group to create the image in.
+        shape (Sequence[int]): The shape of the image.
+        reference_sample (AVAILABLE_SAMPLES | SampleInfo): The reference sample to use.
+        levels (int | list[str], optional): The number of levels in the pyramid or a
+            list of level names. Defaults to 5.
+        xy_scaling_factor (float, optional): The down-scaling factor in x and y
+            dimensions. Defaults to 2.0.
+        z_scaling_factor (float, optional): The down-scaling factor in z dimension.
+            Defaults to 1.0.
+        axes_names (Sequence[str] | None, optional): The names of the axes.
+            If None the canonical names are used. Defaults to None.
+        chunks (Sequence[int] | None, optional): The chunk shape. If None the shape
+            is used. Defaults to None.
+        channel_labels (list[str] | None, optional): The labels of the channels.
+            Defaults to None.
+        channel_wavelengths (list[str] | None, optional): The wavelengths of the
+            channels. Defaults to None.
+        channel_colors (Sequence[str] | None, optional): The colors of the channels.
+            Defaults to None.
+        channel_active (Sequence[bool] | None, optional): Whether the channels are
+            active. Defaults to None.
+        table_backend (TableBackend): Table backend to be used to store tables
+        overwrite (bool, optional): Whether to overwrite an existing image.
+            Defaults to True.
+        version (NgffVersion, optional): The version of the OME-Zarr specification.
+            Defaults to DefaultNgffVersion.
+    """
+    if isinstance(reference_sample, str):
+        sample_info = get_sample_info(reference_sample)
+    else:
+        sample_info = reference_sample
+
+    raw = np.asarray(PIL.Image.open(sample_info.img_path))
+    raw = fit_to_shape(arr=raw, out_shape=tuple(shape))
+    raw = raw / np.max(raw) * (2**16 - 1)
+    raw = raw.astype(np.uint16)
+    ome_zarr = create_ome_zarr_from_array(
+        store=store,
+        array=raw,
+        xy_pixelsize=sample_info.xy_pixelsize,
+        z_spacing=sample_info.z_spacing,
+        time_spacing=sample_info.time_spacing,
+        levels=levels,
+        xy_scaling_factor=xy_scaling_factor,
+        z_scaling_factor=z_scaling_factor,
+        space_unit=sample_info.space_unit,
+        time_unit=sample_info.time_unit,
+        axes_names=axes_names,
+        channel_labels=channel_labels,
+        channel_wavelengths=channel_wavelengths,
+        channel_colors=channel_colors,
+        channel_active=channel_active,
+        name=sample_info.name,
+        chunks=chunks,
+        overwrite=overwrite,
+        version=version,
+    )
+
+    image = ome_zarr.get_image()
+    well_table = image.build_image_roi_table()
+    ome_zarr.add_table("well_ROI_table", table=well_table, backend=table_backend)
+
+    for label_info in sample_info.labels:
+        ome_zarr.derive_label(name=label_info.name)
+        label = ome_zarr.get_label(name=label_info.name)
+
+        ref_label = np.asarray(PIL.Image.open(label_info.label_path))
+        ref_label = fit_to_shape(
+            arr=ref_label,
+            out_shape=label.shape,
+            ensure_unique_info=label_info.ensure_unique_labels,
+        )
+        ref_label = ref_label.astype(np.uint32)
+        label.set_array(ref_label)
+        label.consolidate()
+
+        if label_info.create_masking_table:
+            masking_table = label.build_masking_roi_table()
+            ome_zarr.add_table(
+                name=f"{label_info.name}_masking_table",
+                table=masking_table,
+                backend=table_backend,
+            )
+
+    return ome_zarr