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

ngio/images/_masked_image.py

@@ -1,11 +1,25 @@
 """A module for handling label images in OME-NGFF files."""
 
-from collections.abc import Collection, Iterable
+from collections.abc import Sequence
 from typing import Literal
 
-from ngio.common import ArrayLike, get_masked_pipe, roi_to_slice_kwargs, set_masked_pipe
-from ngio.images._image import Image
+import dask.array as da
+import numpy as np
+
+from ngio.images._image import (
+    ChannelSlicingInputType,
+    Image,
+    SlicingInputType,
+    add_channel_selection_to_slicing_dict,
+)
 from ngio.images._label import Label
+from ngio.io_pipes import (
+    DaskGetterMasked,
+    DaskSetterMasked,
+    NumpyGetterMasked,
+    NumpySetterMasked,
+    TransformProtocol,
+)
 from ngio.ome_zarr_meta import ImageMetaHandler, LabelMetaHandler
 from ngio.tables import MaskingRoiTable
 from ngio.utils import (
@@ -47,72 +61,236 @@ class MaskedImage(Image):
             label_name = self._masking_roi_table.reference_label
         return f"MaskedImage(path={self.path}, {self.dimensions}, {label_name})"
 
-    def get_roi(
+    def get_roi_as_numpy(  # type: ignore (this ignore the method override issue)
         self,
         label: int,
+        channel_selection: ChannelSlicingInputType | None = None,
         zoom_factor: float = 1.0,
-        axes_order: Collection[str] | None = None,
-        mode: Literal["numpy", "dask", "delayed"] = "numpy",
-        **slice_kwargs: slice | int | Iterable[int],
-    ) -> ArrayLike:
+        axes_order: Sequence[str] | None = None,
+        transforms: Sequence[TransformProtocol] | None = None,
+        **slicing_kwargs: slice | int | Sequence[int],
+    ) -> np.ndarray:
         """Return the array for a given ROI."""
-        roi = self._masking_roi_table.get(label)
+        roi = self._masking_roi_table.get_label(label)
+        roi = roi.zoom(zoom_factor)
+        return super().get_roi_as_numpy(
+            roi=roi,
+            channel_selection=channel_selection,
+            axes_order=axes_order,
+            transforms=transforms,
+            **slicing_kwargs,
+        )
+
+    def get_roi_as_dask(  # type: ignore (this ignore the method override issue)
+        self,
+        label: int,
+        channel_selection: ChannelSlicingInputType | None = None,
+        zoom_factor: float = 1.0,
+        axes_order: Sequence[str] | None = None,
+        transforms: Sequence[TransformProtocol] | None = None,
+        **slicing_kwargs: slice | int | Sequence[int],
+    ) -> da.Array:
+        """Return the array for a given ROI as a Dask array."""
+        roi = self._masking_roi_table.get_label(label)
+        roi = roi.zoom(zoom_factor)
+        return super().get_roi_as_dask(
+            roi=roi,
+            channel_selection=channel_selection,
+            axes_order=axes_order,
+            transforms=transforms,
+            **slicing_kwargs,
+        )
+
+    def get_roi(  # type: ignore (this ignore the method override issue)
+        self,
+        label: int,
+        zoom_factor: float = 1.0,
+        channel_selection: ChannelSlicingInputType | None = None,
+        axes_order: Sequence[str] | None = None,
+        transforms: Sequence[TransformProtocol] | None = None,
+        mode: Literal["numpy", "dask"] = "numpy",
+        **slicing_kwargs: slice | int | Sequence[int],
+    ) -> np.ndarray | da.Array:
+        """Return the array for a given ROI."""
+        roi = self._masking_roi_table.get_label(label)
         roi = roi.zoom(zoom_factor)
         return super().get_roi(
-            roi=roi, axes_order=axes_order, mode=mode, **slice_kwargs
+            roi=roi,
+            channel_selection=channel_selection,
+            axes_order=axes_order,
+            transforms=transforms,
+            mode=mode,
+            **slicing_kwargs,
         )
 
-    def set_roi(
+    def set_roi(  # type: ignore (this ignore the method override issue)
         self,
         label: int,
-        patch: ArrayLike,
+        patch: np.ndarray | da.Array,
         zoom_factor: float = 1.0,
-        axes_order: Collection[str] | None = None,
-        **slice_kwargs: slice | int | Iterable[int],
+        channel_selection: ChannelSlicingInputType | None = None,
+        axes_order: Sequence[str] | None = None,
+        transforms: Sequence[TransformProtocol] | None = None,
+        **slicing_kwargs: slice | int | Sequence[int],
     ) -> None:
         """Set the array for a given ROI."""
-        roi = self._masking_roi_table.get(label)
+        roi = self._masking_roi_table.get_label(label)
         roi = roi.zoom(zoom_factor)
         return super().set_roi(
-            roi=roi, patch=patch, axes_order=axes_order, **slice_kwargs
+            roi=roi,
+            patch=patch,
+            channel_selection=channel_selection,
+            axes_order=axes_order,
+            transforms=transforms,
+            **slicing_kwargs,
         )
 
-    def get_roi_masked(
+    def get_roi_masked_as_numpy(
         self,
         label: int,
-        axes_order: Collection[str] | None = None,
-        mode: Literal["numpy", "dask", "delayed"] = "numpy",
+        channel_selection: ChannelSlicingInputType | None = None,
         zoom_factor: float = 1.0,
-        **slice_kwargs: slice | int | Iterable[int],
-    ) -> ArrayLike:
-        """Return the masked array for a given label."""
-        return get_masked_roi_pipe(
-            image=self,
-            label=label,
+        axes_order: Sequence[str] | None = None,
+        transforms: Sequence[TransformProtocol] | None = None,
+        allow_rescaling: bool = True,
+        **slicing_kwargs: SlicingInputType,
+    ) -> np.ndarray:
+        """Return the masked array for a given label as a NumPy array."""
+        slicing_kwargs = add_channel_selection_to_slicing_dict(
+            image=self, channel_selection=channel_selection, slicing_dict=slicing_kwargs
+        )
+
+        roi = self._masking_roi_table.get_label(label)
+        roi = roi.zoom(zoom_factor)
+        masked_getter = NumpyGetterMasked(
+            roi=roi,
+            zarr_array=self.zarr_array,
+            label_zarr_array=self._label.zarr_array,
+            dimensions=self.dimensions,
+            label_dimensions=self._label.dimensions,
             axes_order=axes_order,
-            mode=mode,
-            zoom_factor=zoom_factor,
-            **slice_kwargs,
+            transforms=transforms,
+            slicing_dict=slicing_kwargs,
+            allow_rescaling=allow_rescaling,
+        )
+        return masked_getter()
+
+    def get_roi_masked_as_dask(
+        self,
+        label: int,
+        channel_selection: ChannelSlicingInputType | None = None,
+        zoom_factor: float = 1.0,
+        axes_order: Sequence[str] | None = None,
+        transforms: Sequence[TransformProtocol] | None = None,
+        allow_rescaling: bool = True,
+        **slicing_kwargs: SlicingInputType,
+    ) -> da.Array:
+        """Return the masked array for a given label as a Dask array."""
+        slicing_kwargs = add_channel_selection_to_slicing_dict(
+            image=self, channel_selection=channel_selection, slicing_dict=slicing_kwargs
+        )
+
+        roi = self._masking_roi_table.get_label(label)
+        roi = roi.zoom(zoom_factor)
+        masked_getter = DaskGetterMasked(
+            roi=roi,
+            zarr_array=self.zarr_array,
+            label_zarr_array=self._label.zarr_array,
+            dimensions=self.dimensions,
+            label_dimensions=self._label.dimensions,
+            axes_order=axes_order,
+            transforms=transforms,
+            slicing_dict=slicing_kwargs,
+            allow_rescaling=allow_rescaling,
         )
+        return masked_getter()
+
+    def get_roi_masked(
+        self,
+        label: int,
+        channel_selection: ChannelSlicingInputType | None = None,
+        zoom_factor: float = 1.0,
+        axes_order: Sequence[str] | None = None,
+        transforms: Sequence[TransformProtocol] | None = None,
+        mode: Literal["numpy", "dask"] = "numpy",
+        allow_rescaling: bool = True,
+        **slicing_kwargs: SlicingInputType,
+    ) -> np.ndarray | da.Array:
+        """Return the masked array for a given label."""
+        if mode == "numpy":
+            return self.get_roi_masked_as_numpy(
+                label=label,
+                channel_selection=channel_selection,
+                zoom_factor=zoom_factor,
+                axes_order=axes_order,
+                transforms=transforms,
+                allow_rescaling=allow_rescaling,
+                **slicing_kwargs,
+            )
+
+        elif mode == "dask":
+            return self.get_roi_masked_as_dask(
+                label=label,
+                channel_selection=channel_selection,
+                zoom_factor=zoom_factor,
+                axes_order=axes_order,
+                transforms=transforms,
+                allow_rescaling=allow_rescaling,
+                **slicing_kwargs,
+            )
+        else:
+            raise ValueError(f"Unknown mode: {mode}")
 
     def set_roi_masked(
         self,
         label: int,
-        patch: ArrayLike,
-        axes_order: Collection[str] | None = None,
+        patch: np.ndarray | da.Array,
+        channel_selection: ChannelSlicingInputType | None = None,
+        axes_order: Sequence[str] | None = None,
         zoom_factor: float = 1.0,
-        **slice_kwargs: slice | int | Iterable[int],
+        transforms: Sequence[TransformProtocol] | None = None,
+        allow_rescaling: bool = True,
+        **slicing_kwargs: SlicingInputType,
     ) -> None:
         """Set the masked array for a given label."""
-        return set_masked_roi_pipe(
-            image=self,
-            label=label,
-            patch=patch,
-            axes_order=axes_order,
-            zoom_factor=zoom_factor,
-            **slice_kwargs,
+        slicing_kwargs = add_channel_selection_to_slicing_dict(
+            image=self, channel_selection=channel_selection, slicing_dict=slicing_kwargs
         )
 
+        roi = self._masking_roi_table.get_label(label)
+        roi = roi.zoom(zoom_factor)
+        if isinstance(patch, da.Array):
+            path_setter = DaskSetterMasked(
+                roi=roi,
+                zarr_array=self.zarr_array,
+                label_zarr_array=self._label.zarr_array,
+                dimensions=self.dimensions,
+                label_dimensions=self._label.dimensions,
+                axes_order=axes_order,
+                transforms=transforms,
+                slicing_dict=slicing_kwargs,
+                allow_rescaling=allow_rescaling,
+            )
+            path_setter(patch)
+        elif isinstance(patch, np.ndarray):
+            path_setter = NumpySetterMasked(
+                roi=roi,
+                zarr_array=self.zarr_array,
+                label_zarr_array=self._label.zarr_array,
+                dimensions=self.dimensions,
+                label_dimensions=self._label.dimensions,
+                axes_order=axes_order,
+                transforms=transforms,
+                slicing_dict=slicing_kwargs,
+                allow_rescaling=allow_rescaling,
+            )
+            path_setter(patch)
+        else:
+            raise TypeError(
+                f"Unsupported patch type: {type(patch)}. "
+                "Expected numpy.ndarray or dask.array.Array."
+            )
+
 
 class MaskedLabel(Label):
     """Placeholder class for a label."""
@@ -148,126 +326,206 @@ class MaskedLabel(Label):
             label_name = self._masking_roi_table.reference_label
         return f"MaskedLabel(path={self.path}, {self.dimensions}, {label_name})"
 
+    def get_roi_as_numpy(
+        self,
+        label: int,
+        zoom_factor: float = 1.0,
+        axes_order: Sequence[str] | None = None,
+        transforms: Sequence[TransformProtocol] | None = None,
+        **slicing_kwargs: slice | int | Sequence[int],
+    ) -> np.ndarray:
+        """Return the ROI as a NumPy array."""
+        roi = self._masking_roi_table.get_label(label)
+        roi = roi.zoom(zoom_factor)
+        return super().get_roi_as_numpy(
+            roi=roi,
+            axes_order=axes_order,
+            transforms=transforms,
+            **slicing_kwargs,
+        )
+
+    def get_roi_as_dask(
+        self,
+        label: int,
+        zoom_factor: float = 1.0,
+        axes_order: Sequence[str] | None = None,
+        transforms: Sequence[TransformProtocol] | None = None,
+        **slicing_kwargs: slice | int | Sequence[int],
+    ) -> da.Array:
+        """Return the ROI as a Dask array."""
+        roi = self._masking_roi_table.get_label(label)
+        roi = roi.zoom(zoom_factor)
+        return super().get_roi_as_dask(
+            roi=roi,
+            axes_order=axes_order,
+            transforms=transforms,
+            **slicing_kwargs,
+        )
+
     def get_roi(
         self,
         label: int,
         zoom_factor: float = 1.0,
-        axes_order: Collection[str] | None = None,
-        mode: Literal["numpy", "dask", "delayed"] = "numpy",
-        **slice_kwargs: slice | int | Iterable[int],
-    ) -> ArrayLike:
+        axes_order: Sequence[str] | None = None,
+        mode: Literal["numpy", "dask"] = "numpy",
+        transforms: Sequence[TransformProtocol] | None = None,
+        **slicing_kwargs: slice | int | Sequence[int],
+    ) -> np.ndarray | da.Array:
         """Return the array for a given ROI."""
-        roi = self._masking_roi_table.get(label)
+        roi = self._masking_roi_table.get_label(label)
         roi = roi.zoom(zoom_factor)
         return super().get_roi(
-            roi=roi, axes_order=axes_order, mode=mode, **slice_kwargs
+            roi=roi,
+            axes_order=axes_order,
+            mode=mode,
+            transforms=transforms,
+            **slicing_kwargs,
         )
 
     def set_roi(
         self,
         label: int,
-        patch: ArrayLike,
+        patch: np.ndarray | da.Array,
         zoom_factor: float = 1.0,
-        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: slice | int | Sequence[int],
     ) -> None:
         """Set the array for a given ROI."""
-        roi = self._masking_roi_table.get(label)
+        roi = self._masking_roi_table.get_label(label)
         roi = roi.zoom(zoom_factor)
         return super().set_roi(
-            roi=roi, patch=patch, axes_order=axes_order, **slice_kwargs
+            roi=roi,
+            patch=patch,
+            axes_order=axes_order,
+            transforms=transforms,
+            **slicing_kwargs,
         )
 
-    def get_roi_masked(
+    def get_roi_masked_as_numpy(
         self,
         label: int,
-        axes_order: Collection[str] | None = None,
-        mode: Literal["numpy", "dask", "delayed"] = "numpy",
         zoom_factor: float = 1.0,
-        **slice_kwargs: slice | int | Iterable[int],
-    ) -> ArrayLike:
-        """Return the masked array for a given label."""
-        return get_masked_roi_pipe(
-            image=self,
-            label=label,
+        axes_order: Sequence[str] | None = None,
+        transforms: Sequence[TransformProtocol] | None = None,
+        allow_rescaling: bool = True,
+        **slicing_kwargs: SlicingInputType,
+    ) -> np.ndarray:
+        """Return the masked array for a given label as a NumPy array."""
+        roi = self._masking_roi_table.get_label(label)
+        roi = roi.zoom(zoom_factor)
+        masked_getter = NumpyGetterMasked(
+            roi=roi,
+            zarr_array=self.zarr_array,
+            label_zarr_array=self._label.zarr_array,
+            dimensions=self.dimensions,
+            label_dimensions=self._label.dimensions,
             axes_order=axes_order,
-            mode=mode,
-            zoom_factor=zoom_factor,
-            **slice_kwargs,
+            transforms=transforms,
+            slicing_dict=slicing_kwargs,
+            allow_rescaling=allow_rescaling,
         )
+        return masked_getter()
 
-    def set_roi_masked(
+    def get_roi_masked_as_dask(
         self,
         label: int,
-        patch: ArrayLike,
-        axes_order: Collection[str] | None = None,
         zoom_factor: float = 1.0,
-        **slice_kwargs: slice | int | Iterable[int],
-    ) -> None:
-        """Set the masked array for a given label."""
-        return set_masked_roi_pipe(
-            image=self,
-            label=label,
-            patch=patch,
+        axes_order: Sequence[str] | None = None,
+        transforms: Sequence[TransformProtocol] | None = None,
+        allow_rescaling: bool = True,
+        **slicing_kwargs: SlicingInputType,
+    ) -> da.Array:
+        """Return the masked array for a given label as a Dask array."""
+        roi = self._masking_roi_table.get_label(label)
+        roi = roi.zoom(zoom_factor)
+        masked_getter = DaskGetterMasked(
+            roi=roi,
+            zarr_array=self.zarr_array,
+            label_zarr_array=self._label.zarr_array,
+            dimensions=self.dimensions,
+            label_dimensions=self._label.dimensions,
             axes_order=axes_order,
-            zoom_factor=zoom_factor,
-            **slice_kwargs,
+            transforms=transforms,
+            slicing_dict=slicing_kwargs,
+            allow_rescaling=allow_rescaling,
         )
+        return masked_getter()
+
+    def get_roi_masked(
+        self,
+        label: int,
+        zoom_factor: float = 1.0,
+        axes_order: Sequence[str] | None = None,
+        mode: Literal["numpy", "dask"] = "numpy",
+        transforms: Sequence[TransformProtocol] | None = None,
+        allow_rescaling: bool = True,
+        **slicing_kwargs: SlicingInputType,
+    ) -> np.ndarray | da.Array:
+        """Return the masked array for a given label."""
+        if mode == "numpy":
+            return self.get_roi_masked_as_numpy(
+                label=label,
+                zoom_factor=zoom_factor,
+                axes_order=axes_order,
+                transforms=transforms,
+                allow_rescaling=allow_rescaling,
+                **slicing_kwargs,
+            )
 
+        elif mode == "dask":
+            return self.get_roi_masked_as_dask(
+                label=label,
+                zoom_factor=zoom_factor,
+                axes_order=axes_order,
+                transforms=transforms,
+                allow_rescaling=allow_rescaling,
+                **slicing_kwargs,
+            )
+        else:
+            raise ValueError(f"Unknown mode: {mode}")
 
-def get_masked_roi_pipe(
-    image: MaskedImage | MaskedLabel,
-    label: int,
-    axes_order: Collection[str] | None = None,
-    mode: Literal["numpy", "dask", "delayed"] = "numpy",
-    zoom_factor: float = 1.0,
-    **slice_kwargs: slice | int | Iterable[int],
-) -> ArrayLike:
-    """Return the masked array for a given label."""
-    roi = image._masking_roi_table.get(label)
-    roi = roi.zoom(zoom_factor)
-    slice_kwargs = roi_to_slice_kwargs(
-        roi=roi,
-        pixel_size=image.pixel_size,
-        dimensions=image.dimensions,
-        **slice_kwargs,
-    )
-    return get_masked_pipe(
-        array=image.zarr_array,
-        label_array=image._label.zarr_array,
-        label=label,
-        dimensions_array=image.dimensions,
-        dimensions_label=image._label.dimensions,
-        axes_order=axes_order,
-        mode=mode,
-        **slice_kwargs,
-    )
-
-
-def set_masked_roi_pipe(
-    image: MaskedImage | MaskedLabel,
-    label: int,
-    patch: ArrayLike,
-    axes_order: Collection[str] | None = None,
-    zoom_factor: float = 1.0,
-    **slice_kwargs: slice | int | Iterable[int],
-) -> None:
-    """Set the masked array for a given label."""
-    roi = image._masking_roi_table.get(label)
-    roi = roi.zoom(zoom_factor)
-    slice_kwargs = roi_to_slice_kwargs(
-        roi=roi,
-        pixel_size=image.pixel_size,
-        dimensions=image.dimensions,
-        **slice_kwargs,
-    )
-    return set_masked_pipe(
-        array=image.zarr_array,
-        label_array=image._label.zarr_array,
-        label=label,
-        patch=patch,
-        dimensions_array=image.dimensions,
-        dimensions_label=image._label.dimensions,
-        axes_order=axes_order,
-        **slice_kwargs,
-    )
+    def set_roi_masked(
+        self,
+        label: int,
+        patch: np.ndarray | da.Array,
+        axes_order: Sequence[str] | None = None,
+        zoom_factor: float = 1.0,
+        transforms: Sequence[TransformProtocol] | None = None,
+        allow_rescaling: bool = True,
+        **slicing_kwargs: SlicingInputType,
+    ) -> None:
+        """Set the masked array for a given label."""
+        roi = self._masking_roi_table.get_label(label)
+        roi = roi.zoom(zoom_factor)
+        if isinstance(patch, da.Array):
+            path_setter = DaskSetterMasked(
+                roi=roi,
+                zarr_array=self.zarr_array,
+                label_zarr_array=self._label.zarr_array,
+                dimensions=self.dimensions,
+                label_dimensions=self._label.dimensions,
+                axes_order=axes_order,
+                transforms=transforms,
+                slicing_dict=slicing_kwargs,
+                allow_rescaling=allow_rescaling,
+            )
+            path_setter(patch)
+        elif isinstance(patch, np.ndarray):
+            path_setter = NumpySetterMasked(
+                roi=roi,
+                zarr_array=self.zarr_array,
+                label_zarr_array=self._label.zarr_array,
+                dimensions=self.dimensions,
+                label_dimensions=self._label.dimensions,
+                axes_order=axes_order,
+                transforms=transforms,
+                slicing_dict=slicing_kwargs,
+                allow_rescaling=allow_rescaling,
+            )
+            path_setter(patch)
+        else:
+            raise TypeError(
+                f"Unsupported patch type: {type(patch)}. "
+                "Expected numpy.ndarray or dask.array.Array."
+            )