ngio 0.2.0a2__py3-none-any.whl → 0.5.0b4__py3-none-any.whl

ngio/images/_masked_image.py

@@ -0,0 +1,531 @@
+"""A module for handling label images in OME-NGFF files."""
+
+from collections.abc import Sequence
+from typing import Literal
+
+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 (
+    ZarrGroupHandler,
+)
+
+
+class MaskedImage(Image):
+    """Placeholder class for a label."""
+
+    def __init__(
+        self,
+        group_handler: ZarrGroupHandler,
+        path: str,
+        meta_handler: ImageMetaHandler | None,
+        label: Label,
+        masking_roi_table: MaskingRoiTable,
+    ) -> None:
+        """Initialize the Image at a single level.
+
+        Args:
+            group_handler: The Zarr group handler.
+            path: The path to the image in the ome_zarr file.
+            meta_handler: The image metadata handler.
+            label: The label image.
+            masking_roi_table: The masking ROI table.
+
+        """
+        super().__init__(
+            group_handler=group_handler, path=path, meta_handler=meta_handler
+        )
+        self._label = label
+        self._masking_roi_table = masking_roi_table
+
+    def __repr__(self) -> str:
+        """Return a string representation of the object."""
+        label_name = self._label.meta.name
+        if label_name is None:
+            label_name = self._masking_roi_table.reference_label
+        return f"MaskedImage(path={self.path}, {self.dimensions}, {label_name})"
+
+    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: 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(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,
+            channel_selection=channel_selection,
+            axes_order=axes_order,
+            transforms=transforms,
+            mode=mode,
+            **slicing_kwargs,
+        )
+
+    def set_roi(  # type: ignore (this ignore the method override issue)
+        self,
+        label: int,
+        patch: np.ndarray | da.Array,
+        zoom_factor: float = 1.0,
+        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(label)
+        roi = roi.zoom(zoom_factor)
+        return super().set_roi(
+            roi=roi,
+            patch=patch,
+            channel_selection=channel_selection,
+            axes_order=axes_order,
+            transforms=transforms,
+            **slicing_kwargs,
+        )
+
+    def get_roi_masked_as_numpy(
+        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,
+    ) -> 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,
+            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: np.ndarray | da.Array,
+        channel_selection: ChannelSlicingInputType | None = None,
+        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."""
+        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."""
+
+    def __init__(
+        self,
+        group_handler: ZarrGroupHandler,
+        path: str,
+        meta_handler: LabelMetaHandler | None,
+        label: Label,
+        masking_roi_table: MaskingRoiTable,
+    ) -> None:
+        """Initialize the Image at a single level.
+
+        Args:
+            group_handler: The Zarr group handler.
+            path: The path to the image in the ome_zarr file.
+            meta_handler: The image metadata handler.
+            label: The label image.
+            masking_roi_table: The masking ROI table.
+
+        """
+        super().__init__(
+            group_handler=group_handler, path=path, meta_handler=meta_handler
+        )
+        self._label = label
+        self._masking_roi_table = masking_roi_table
+
+    def __repr__(self) -> str:
+        """Return a string representation of the object."""
+        label_name = self._label.meta.name
+        if label_name is None:
+            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: 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(label)
+        roi = roi.zoom(zoom_factor)
+        return super().get_roi(
+            roi=roi,
+            axes_order=axes_order,
+            mode=mode,
+            transforms=transforms,
+            **slicing_kwargs,
+        )
+
+    def set_roi(
+        self,
+        label: int,
+        patch: np.ndarray | da.Array,
+        zoom_factor: float = 1.0,
+        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(label)
+        roi = roi.zoom(zoom_factor)
+        return super().set_roi(
+            roi=roi,
+            patch=patch,
+            axes_order=axes_order,
+            transforms=transforms,
+            **slicing_kwargs,
+        )
+
+    def get_roi_masked_as_numpy(
+        self,
+        label: int,
+        zoom_factor: float = 1.0,
+        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,
+            transforms=transforms,
+            slicing_dict=slicing_kwargs,
+            allow_rescaling=allow_rescaling,
+        )
+        return masked_getter()
+
+    def get_roi_masked_as_dask(
+        self,
+        label: int,
+        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."""
+        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,
+        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 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."
+            )