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

ngio/experimental/iterators/_segmentation.py

@@ -0,0 +1,235 @@
+from collections.abc import Sequence
+
+import dask.array as da
+import numpy as np
+
+from ngio.common import Roi
+from ngio.experimental.iterators._abstract_iterator import AbstractIteratorBuilder
+from ngio.images import Image, Label
+from ngio.images._image import (
+    ChannelSlicingInputType,
+    add_channel_selection_to_slicing_dict,
+)
+from ngio.images._masked_image import MaskedImage
+from ngio.io_pipes import (
+    DaskGetterMasked,
+    DaskRoiGetter,
+    DaskRoiSetter,
+    DaskSetterMasked,
+    NumpyGetterMasked,
+    NumpyRoiGetter,
+    NumpyRoiSetter,
+    NumpySetterMasked,
+    TransformProtocol,
+)
+from ngio.io_pipes._io_pipes_types import DataGetterProtocol, DataSetterProtocol
+
+
+class SegmentationIterator(AbstractIteratorBuilder[np.ndarray, da.Array]):
+    """Base class for iterators over ROIs."""
+
+    def __init__(
+        self,
+        input_image: Image,
+        output_label: Label,
+        channel_selection: ChannelSlicingInputType = None,
+        axes_order: Sequence[str] | None = None,
+        input_transforms: Sequence[TransformProtocol] | None = None,
+        output_transforms: Sequence[TransformProtocol] | None = None,
+    ) -> None:
+        """Initialize the iterator with a ROI table and input/output images.
+
+        Args:
+            input_image (Image): The input image to be used as input for the
+                segmentation.
+            output_label (Label): The label image where the ROIs will be written.
+            channel_selection (ChannelSlicingInputType): Optional
+                selection of channels to use for the segmentation.
+            axes_order (Sequence[str] | None): Optional axes order for the
+                segmentation.
+            input_transforms (Sequence[TransformProtocol] | None): Optional
+                transforms to apply to the input image.
+            output_transforms (Sequence[TransformProtocol] | None): Optional
+                transforms to apply to the output label.
+        """
+        self._input = input_image
+        self._output = output_label
+        self._ref_image = input_image
+        self._rois = input_image.build_image_roi_table(name=None).rois()
+
+        # Set iteration parameters
+        self._input_slicing_kwargs = add_channel_selection_to_slicing_dict(
+            image=self._input, channel_selection=channel_selection, slicing_dict={}
+        )
+        self._channel_selection = channel_selection
+        self._axes_order = axes_order
+        self._input_transforms = input_transforms
+        self._output_transforms = output_transforms
+
+        self._input.require_dimensions_match(self._output, allow_singleton=False)
+
+    def get_init_kwargs(self) -> dict:
+        """Return the initialization arguments for the iterator."""
+        return {
+            "input_image": self._input,
+            "output_label": self._output,
+            "channel_selection": self._channel_selection,
+            "axes_order": self._axes_order,
+            "input_transforms": self._input_transforms,
+            "output_transforms": self._output_transforms,
+        }
+
+    def build_numpy_getter(self, roi: Roi) -> DataGetterProtocol[np.ndarray]:
+        return NumpyRoiGetter(
+            zarr_array=self._input.zarr_array,
+            dimensions=self._input.dimensions,
+            roi=roi,
+            axes_order=self._axes_order,
+            transforms=self._input_transforms,
+            slicing_dict=self._input_slicing_kwargs,
+        )
+
+    def build_numpy_setter(self, roi: Roi) -> DataSetterProtocol[np.ndarray]:
+        return NumpyRoiSetter(
+            zarr_array=self._output.zarr_array,
+            dimensions=self._output.dimensions,
+            roi=roi,
+            axes_order=self._axes_order,
+            transforms=self._output_transforms,
+            remove_channel_selection=True,
+        )
+
+    def build_dask_getter(self, roi: Roi) -> DataGetterProtocol[da.Array]:
+        return DaskRoiGetter(
+            zarr_array=self._input.zarr_array,
+            dimensions=self._input.dimensions,
+            roi=roi,
+            axes_order=self._axes_order,
+            transforms=self._input_transforms,
+            slicing_dict=self._input_slicing_kwargs,
+        )
+
+    def build_dask_setter(self, roi: Roi) -> DataSetterProtocol[da.Array]:
+        return DaskRoiSetter(
+            zarr_array=self._output.zarr_array,
+            dimensions=self._output.dimensions,
+            roi=roi,
+            axes_order=self._axes_order,
+            transforms=self._output_transforms,
+            remove_channel_selection=True,
+        )
+
+    def post_consolidate(self):
+        self._output.consolidate()
+
+
+class MaskedSegmentationIterator(SegmentationIterator):
+    """Base class for iterators over ROIs."""
+
+    def __init__(
+        self,
+        input_image: MaskedImage,
+        output_label: Label,
+        channel_selection: ChannelSlicingInputType = None,
+        axes_order: Sequence[str] | None = None,
+        input_transforms: Sequence[TransformProtocol] | None = None,
+        output_transforms: Sequence[TransformProtocol] | None = None,
+    ) -> None:
+        """Initialize the iterator with a ROI table and input/output images.
+
+        Args:
+            input_image (MaskedImage): The input image to be used as input for the
+                segmentation.
+            output_label (Label): The label image where the ROIs will be written.
+            channel_selection (ChannelSlicingInputType): Optional
+                selection of channels to use for the segmentation.
+            axes_order (Sequence[str] | None): Optional axes order for the
+                segmentation.
+            input_transforms (Sequence[TransformProtocol] | None): Optional
+                transforms to apply to the input image.
+            output_transforms (Sequence[TransformProtocol] | None): Optional
+                transforms to apply to the output label.
+        """
+        self._input = input_image
+        self._output = output_label
+
+        self._ref_image = input_image
+        self._set_rois(input_image._masking_roi_table.rois())
+
+        # Set iteration parameters
+        self._input_slicing_kwargs = add_channel_selection_to_slicing_dict(
+            image=self._input, channel_selection=channel_selection, slicing_dict={}
+        )
+        self._channel_selection = channel_selection
+        self._axes_order = axes_order
+        self._input_transforms = input_transforms
+        self._output_transforms = output_transforms
+
+        # Check compatibility between input and output images
+        # if not self._input.dimensions.is_compatible_with(self._output.dimensions):
+        #    raise NgioValidationError(
+        #        "Input image and output label have incompatible dimensions. "
+        #        f"Input: {self._input.dimensions}, Output: {self._output.dimensions}."
+        #    )
+
+    def get_init_kwargs(self) -> dict:
+        """Return the initialization arguments for the iterator."""
+        return {
+            "input_image": self._input,
+            "output_label": self._output,
+            "channel_selection": self._channel_selection,
+            "axes_order": self._axes_order,
+            "input_transforms": self._input_transforms,
+            "output_transforms": self._output_transforms,
+        }
+
+    def build_numpy_getter(self, roi: Roi):
+        return NumpyGetterMasked(
+            zarr_array=self._input.zarr_array,
+            dimensions=self._input.dimensions,
+            roi=roi,
+            label_zarr_array=self._input._label.zarr_array,
+            label_dimensions=self._input._label.dimensions,
+            axes_order=self._axes_order,
+            transforms=self._input_transforms,
+            slicing_dict=self._input_slicing_kwargs,
+        )
+
+    def build_numpy_setter(self, roi: Roi):
+        return NumpySetterMasked(
+            roi=roi,
+            zarr_array=self._output.zarr_array,
+            dimensions=self._output.dimensions,
+            label_zarr_array=self._input._label.zarr_array,
+            label_dimensions=self._input._label.dimensions,
+            axes_order=self._axes_order,
+            transforms=self._output_transforms,
+            remove_channel_selection=True,
+        )
+
+    def build_dask_getter(self, roi: Roi):
+        return DaskGetterMasked(
+            roi=roi,
+            zarr_array=self._input.zarr_array,
+            dimensions=self._input.dimensions,
+            label_zarr_array=self._input._label.zarr_array,
+            label_dimensions=self._input._label.dimensions,
+            axes_order=self._axes_order,
+            transforms=self._input_transforms,
+            slicing_dict=self._input_slicing_kwargs,
+        )
+
+    def build_dask_setter(self, roi: Roi):
+        return DaskSetterMasked(
+            roi=roi,
+            zarr_array=self._output.zarr_array,
+            dimensions=self._output.dimensions,
+            label_zarr_array=self._input._label.zarr_array,
+            label_dimensions=self._input._label.dimensions,
+            axes_order=self._axes_order,
+            transforms=self._output_transforms,
+            remove_channel_selection=True,
+        )
+
+    def post_consolidate(self):
+        self._output.consolidate()

ngio/hcs/_plate.py

@@ -1,11 +1,12 @@
-"""A module for handling the Plate Collection in an OME-Zarr file."""
+"""A module for handling the Plate Sequence in an OME-Zarr file."""
 
 import asyncio
 import warnings
-from collections.abc import Collection
+from collections.abc import Sequence
 from typing import Literal
 
-from ngio.common import (
+from ngio.images import (
+    OmeZarrContainer,
     concatenate_image_tables,
     concatenate_image_tables_as,
     concatenate_image_tables_as_async,
@@ -13,7 +14,6 @@ from ngio.common import (
     list_image_tables,
     list_image_tables_async,
 )
-from ngio.images import OmeZarrContainer
 from ngio.ome_zarr_meta import (
     ImageInWellPath,
     NgffVersions,
@@ -67,7 +67,7 @@ class MockLock:
 
 
 class OmeZarrWell:
-    """A class to handle the Well Collection in an OME-Zarr file."""
+    """A class to handle the Well Sequence in an OME-Zarr file."""
 
     def __init__(self, group_handler: ZarrGroupHandler) -> None:
         """Initialize the LabelGroupHandler.
@@ -190,7 +190,7 @@ class OmeZarrWell:
             acquisition_id (int | None): The acquisition id to filter the images.
             strict (bool): Whether to check if the acquisition id is already exists
                 in the well. Defaults to True. If False this might lead to
-                acquision in a well that does not exist at the plate level.
+                acquisition in a well that does not exist at the plate level.
         """
         return self._add_image(
             image_path=image_path,
@@ -200,11 +200,11 @@ class OmeZarrWell:
         )
 
 
-def _build_extras(paths: Collection[str]) -> list[dict[str, str]]:
+def _build_extras(paths: Sequence[str]) -> list[dict[str, str]]:
     """Build the extras for the images.
 
     Args:
-        paths (Collection[str]): The paths of the images.
+        paths (Sequence[str]): The paths of the images.
 
     Returns:
         list[dict[str, str]]: The extras for the images.
@@ -223,7 +223,7 @@ def _build_extras(paths: Collection[str]) -> list[dict[str, str]]:
 
 
 class OmeZarrPlate:
-    """A class to handle the Plate Collection in an OME-Zarr file."""
+    """A class to handle the Plate Sequence in an OME-Zarr file."""
 
     def __init__(
         self,
@@ -382,7 +382,8 @@ class OmeZarrPlate:
         """
         wells = self._group_handler.get_from_cache("wells")
         if wells is not None:
-            return wells  # type: ignore[return-value]
+            assert isinstance(wells, dict)
+            return wells
 
         def process_well(well_path):
             group_handler = self._group_handler.derive_handler(well_path)
@@ -410,7 +411,8 @@ class OmeZarrPlate:
         """
         wells = self._group_handler.get_from_cache("wells")
         if wells is not None:
-            return wells  # type: ignore[return-value]
+            assert isinstance(wells, dict)
+            return wells
 
         def process_well(well_path):
             group_handler = self._group_handler.derive_handler(well_path)
@@ -442,7 +444,8 @@ class OmeZarrPlate:
         """
         images = self._group_handler.get_from_cache("images")
         if images is not None:
-            return images  # type: ignore[return-value]
+            assert isinstance(images, dict)
+            return images
 
         paths = await self.images_paths_async(acquisition=acquisition)
 
@@ -473,7 +476,8 @@ class OmeZarrPlate:
         """
         images = self._group_handler.get_from_cache("images")
         if images is not None:
-            return images  # type: ignore[return-value]
+            assert isinstance(images, dict)
+            return images
         paths = self.images_paths(acquisition=acquisition)
 
         def process_image(image_path):
@@ -969,9 +973,9 @@ class OmeZarrPlate:
                 If 'common', return only common tables between all images.
                 If 'all', return all tables. Defaults to 'common'.
         """
-        images = self.get_images(acquisition=acquisition)
+        images = tuple(self.get_images(acquisition=acquisition).values())
         return list_image_tables(
-            images=images.values(),
+            images=images,
             filter_types=filter_types,
             mode=mode,
         )
@@ -993,15 +997,16 @@ class OmeZarrPlate:
                 If 'all', return all tables. Defaults to 'common'.
         """
         images = await self.get_images_async(acquisition=acquisition)
+        images = tuple(images.values())
         return await list_image_tables_async(
-            images=images.values(),
+            images=images,
             filter_types=filter_types,
             mode=mode,
         )
 
     def concatenate_image_tables(
         self,
-        table_name: str,
+        name: str,
         acquisition: int | None = None,
         strict: bool = True,
         index_key: str | None = None,
@@ -1010,7 +1015,7 @@ class OmeZarrPlate:
         """Concatenate tables from all images in the plate.
 
         Args:
-            table_name: The name of the table to concatenate.
+            name: The name of the table to concatenate.
             index_key: The key to use for the index of the concatenated table.
             acquisition: The acquisition id to filter the images.
             strict: If True, raise an error if the table is not found in the image.
@@ -1021,11 +1026,11 @@ class OmeZarrPlate:
                 if 'lazy', the table will be loaded as a lazy frame.
         """
         images = self.get_images(acquisition=acquisition)
-        extras = _build_extras(images.keys())
+        extras = _build_extras(tuple(images.keys()))
         return concatenate_image_tables(
-            images=images.values(),
+            images=tuple(images.values()),
             extras=extras,
-            table_name=table_name,
+            name=name,
             index_key=index_key,
             strict=strict,
             mode=mode,
@@ -1033,7 +1038,7 @@ class OmeZarrPlate:
 
     def concatenate_image_tables_as(
         self,
-        table_name: str,
+        name: str,
         table_cls: type[TableType],
         acquisition: int | None = None,
         index_key: str | None = None,
@@ -1043,7 +1048,7 @@ class OmeZarrPlate:
         """Concatenate tables from all images in the plate as a specific type.
 
         Args:
-            table_name: The name of the table to concatenate.
+            name: The name of the table to concatenate.
             table_cls: The type of the table to concatenate.
             index_key: The key to use for the index of the concatenated table.
             acquisition: The acquisition id to filter the images.
@@ -1055,11 +1060,11 @@ class OmeZarrPlate:
                 if 'lazy', the table will be loaded as a lazy frame.
         """
         images = self.get_images(acquisition=acquisition)
-        extras = _build_extras(images.keys())
+        extras = _build_extras(tuple(images.keys()))
         return concatenate_image_tables_as(
-            images=images.values(),
+            images=tuple(images.values()),
             extras=extras,
-            table_name=table_name,
+            name=name,
             table_cls=table_cls,
             index_key=index_key,
             strict=strict,
@@ -1068,7 +1073,7 @@ class OmeZarrPlate:
 
     async def concatenate_image_tables_async(
         self,
-        table_name: str,
+        name: str,
         acquisition: int | None = None,
         index_key: str | None = None,
         strict: bool = True,
@@ -1077,7 +1082,7 @@ class OmeZarrPlate:
         """Concatenate tables from all images in the plate asynchronously.
 
         Args:
-            table_name: The name of the table to concatenate.
+            name: The name of the table to concatenate.
             index_key: The key to use for the index of the concatenated table.
             acquisition: The acquisition id to filter the images.
             index_key: If a string is provided, a new index column will be created
@@ -1088,11 +1093,11 @@ class OmeZarrPlate:
                 if 'lazy', the table will be loaded as a lazy frame.
         """
         images = await self.get_images_async(acquisition=acquisition)
-        extras = _build_extras(images.keys())
+        extras = _build_extras(tuple(images.keys()))
         return await concatenate_image_tables_async(
-            images=images.values(),
+            images=tuple(images.values()),
             extras=extras,
-            table_name=table_name,
+            name=name,
             index_key=index_key,
             strict=strict,
             mode=mode,
@@ -1100,7 +1105,7 @@ class OmeZarrPlate:
 
     async def concatenate_image_tables_as_async(
         self,
-        table_name: str,
+        name: str,
         table_cls: type[TableType],
         acquisition: int | None = None,
         index_key: str | None = None,
@@ -1110,7 +1115,7 @@ class OmeZarrPlate:
         """Concatenate tables from all images in the plate as a specific type.
 
         Args:
-            table_name: The name of the table to concatenate.
+            name: The name of the table to concatenate.
             table_cls: The type of the table to concatenate.
             index_key: The key to use for the index of the concatenated table.
             acquisition: The acquisition id to filter the images.
@@ -1122,11 +1127,11 @@ class OmeZarrPlate:
                 if 'lazy', the table will be loaded as a lazy frame.
         """
         images = await self.get_images_async(acquisition=acquisition)
-        extras = _build_extras(images.keys())
+        extras = _build_extras(tuple(images.keys()))
         return await concatenate_image_tables_as_async(
-            images=images.values(),
+            images=tuple(images.values()),
             extras=extras,
-            table_name=table_name,
+            name=name,
             table_cls=table_cls,
             index_key=index_key,
             strict=strict,

ngio/images/__init__.py

@@ -1,23 +1,44 @@
 """OME-Zarr object models."""
 
-from ngio.images._image import Image, ImagesContainer
+from ngio.images._create_synt_container import create_synthetic_ome_zarr
+from ngio.images._image import ChannelSelectionModel, Image, ImagesContainer
 from ngio.images._label import Label, LabelsContainer
 from ngio.images._ome_zarr_container import (
     OmeZarrContainer,
     create_empty_ome_zarr,
     create_ome_zarr_from_array,
     open_image,
+    open_label,
     open_ome_zarr_container,
 )
+from ngio.images._table_ops import (
+    concatenate_image_tables,
+    concatenate_image_tables_as,
+    concatenate_image_tables_as_async,
+    concatenate_image_tables_async,
+    conctatenate_tables,
+    list_image_tables,
+    list_image_tables_async,
+)
 
 __all__ = [
+    "ChannelSelectionModel",
     "Image",
     "ImagesContainer",
     "Label",
     "LabelsContainer",
     "OmeZarrContainer",
+    "concatenate_image_tables",
+    "concatenate_image_tables_as",
+    "concatenate_image_tables_as_async",
+    "concatenate_image_tables_async",
+    "conctatenate_tables",
     "create_empty_ome_zarr",
     "create_ome_zarr_from_array",
+    "create_synthetic_ome_zarr",
+    "list_image_tables",
+    "list_image_tables_async",
     "open_image",
+    "open_label",
     "open_ome_zarr_container",
 ]