ngio 0.1.6__py3-none-any.whl → 0.2.0__py3-none-any.whl

ngio/images/label.py

@@ -0,0 +1,285 @@
+"""A module for handling label images in OME-NGFF files."""
+
+from collections.abc import Collection
+from typing import Literal
+
+from ngio.common import compute_masking_roi
+from ngio.images.abstract_image import AbstractImage, consolidate_image
+from ngio.images.create import _create_empty_label
+from ngio.images.image import Image
+from ngio.ome_zarr_meta import (
+    LabelMetaHandler,
+    NgioLabelMeta,
+    PixelSize,
+    find_label_meta_handler,
+)
+from ngio.tables import MaskingRoiTable
+from ngio.utils import (
+    NgioValidationError,
+    NgioValueError,
+    StoreOrGroup,
+    ZarrGroupHandler,
+)
+
+
+class Label(AbstractImage[LabelMetaHandler]):
+    """Placeholder class for a label."""
+
+    def __init__(
+        self,
+        group_handler: ZarrGroupHandler,
+        path: str,
+        meta_handler: LabelMetaHandler | None,
+    ) -> 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.
+
+        """
+        if meta_handler is None:
+            meta_handler = find_label_meta_handler(group_handler)
+        super().__init__(
+            group_handler=group_handler, path=path, meta_handler=meta_handler
+        )
+
+    def __repr__(self) -> str:
+        """Return the string representation of the label."""
+        return f"Label(path={self.path}, {self.dimensions})"
+
+    @property
+    def meta(self) -> NgioLabelMeta:
+        """Return the metadata."""
+        return self._meta_handler.meta
+
+    def build_masking_roi_table(self) -> MaskingRoiTable:
+        """Compute the masking ROI table."""
+        return build_masking_roi_table(self)
+
+    def consolidate(
+        self,
+        mode: Literal["dask", "numpy", "coarsen"] = "dask",
+    ) -> None:
+        """Consolidate the label on disk."""
+        consolidate_image(self, mode=mode, order=0)
+
+
+class LabelsContainer:
+    """A class to handle the /labels group in an OME-NGFF file."""
+
+    def __init__(self, group_handler: ZarrGroupHandler) -> None:
+        """Initialize the LabelGroupHandler."""
+        self._group_handler = group_handler
+
+        # Validate the group
+        # Either contains a labels attribute or is empty
+        attrs = self._group_handler.load_attrs()
+        if len(attrs) == 0:
+            # It's an empty group
+            pass
+        elif "labels" in attrs and isinstance(attrs["labels"], list):
+            # It's a valid group
+            pass
+        else:
+            raise NgioValidationError(
+                f"Invalid /labels group. "
+                f"Expected a single labels attribute with a list of label names. "
+                f"Found: {attrs}"
+            )
+
+    def list(self) -> list[str]:
+        """Create the /labels group if it doesn't exist."""
+        attrs = self._group_handler.load_attrs()
+        return attrs.get("labels", [])
+
+    def get(
+        self,
+        name: str,
+        path: str | None = None,
+        pixel_size: PixelSize | None = None,
+        strict: bool = False,
+    ) -> Label:
+        """Get a label from the group.
+
+        Args:
+            name (str): The name of the label.
+            path (str | None): The path to the image in the ome_zarr file.
+            pixel_size (PixelSize | None): The pixel size of the image.
+            strict (bool): Only used if the pixel size is provided. If True, the
+                pixel size must match the image pixel size exactly. If False, the
+                closest pixel size level will be returned.
+
+        """
+        group_handler = self._group_handler.derive_handler(name)
+        label_meta_handler = find_label_meta_handler(group_handler)
+        path = label_meta_handler.meta.get_dataset(
+            path=path, pixel_size=pixel_size, strict=strict
+        ).path
+        return Label(group_handler, path, label_meta_handler)
+
+    def derive(
+        self,
+        name: str,
+        ref_image: Image,
+        shape: Collection[int] | None = None,
+        pixel_size: PixelSize | None = None,
+        axes_names: Collection[str] | None = None,
+        chunks: Collection[int] | None = None,
+        dtype: str | None = None,
+        overwrite: bool = False,
+    ) -> "Label":
+        """Create an empty OME-Zarr label from a reference image.
+
+        And add the label to the /labels group.
+
+        Args:
+            store (StoreOrGroup): The Zarr store or group to create the image in.
+            ref_image (Image): The reference image.
+            name (str): The name of the new image.
+            shape (Collection[int] | None): The shape 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.
+                For labels, the channel axis is not allowed.
+            chunks (Collection[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.
+
+        Returns:
+            Label: The new label.
+
+        """
+        existing_labels = self.list()
+        if name in existing_labels and not overwrite:
+            raise NgioValueError(
+                f"Table '{name}' already exists in the group. "
+                "Use overwrite=True to replace it."
+            )
+
+        label_group = self._group_handler.get_group(name, create_mode=True)
+
+        _derive_label(
+            store=label_group,
+            ref_image=ref_image,
+            name=name,
+            shape=shape,
+            pixel_size=pixel_size,
+            axes_names=axes_names,
+            chunks=chunks,
+            dtype=dtype,
+            overwrite=overwrite,
+        )
+
+        if name not in existing_labels:
+            existing_labels.append(name)
+            self._group_handler.write_attrs({"labels": existing_labels})
+
+        return self.get(name)
+
+
+def _derive_label(
+    store: StoreOrGroup,
+    ref_image: Image,
+    name: str,
+    shape: Collection[int] | None = None,
+    pixel_size: PixelSize | None = None,
+    axes_names: Collection[str] | None = None,
+    chunks: Collection[int] | None = None,
+    dtype: str | None = None,
+    overwrite: bool = False,
+) -> None:
+    """Create an empty OME-Zarr label from a reference image.
+
+    Args:
+        store (StoreOrGroup): The Zarr store or group to create the image in.
+        ref_image (Image): The reference image.
+        name (str): The name of the new image.
+        shape (Collection[int] | None): The shape 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.
+            For labels, the channel axis is not allowed.
+        chunks (Collection[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.
+
+    Returns:
+        None
+
+    """
+    ref_meta = ref_image.meta
+
+    if shape is None:
+        shape = ref_image.shape
+
+    if pixel_size is None:
+        pixel_size = ref_image.pixel_size
+
+    if axes_names is None:
+        axes_names = ref_meta.axes_mapper.on_disk_axes_names
+        c_axis = ref_meta.axes_mapper.get_index("c")
+    else:
+        if "c" in axes_names:
+            raise NgioValidationError(
+                "Labels cannot have a channel axis. "
+                "Please remove the channel axis from the axes names."
+            )
+        c_axis = None
+
+    if len(axes_names) != len(shape):
+        raise NgioValidationError(
+            "The axes names of the new image does not match the reference image."
+            f"Got {axes_names} for shape {shape}."
+        )
+
+    if chunks is None:
+        chunks = ref_image.chunks
+
+    if len(chunks) != len(shape):
+        raise NgioValidationError(
+            "The chunks of the new image does not match the reference image."
+            f"Got {chunks} for shape {shape}."
+        )
+
+    if dtype is None:
+        dtype = ref_image.dtype
+
+    if c_axis is not None:
+        # remove channel if present
+        shape = list(shape)
+        shape = shape[:c_axis] + shape[c_axis + 1 :]
+        chunks = list(chunks)
+        chunks = chunks[:c_axis] + chunks[c_axis + 1 :]
+        axes_names = list(axes_names)
+        axes_names = axes_names[:c_axis] + axes_names[c_axis + 1 :]
+
+    _ = _create_empty_label(
+        store=store,
+        shape=shape,
+        pixelsize=ref_image.pixel_size.x,
+        z_spacing=ref_image.pixel_size.z,
+        time_spacing=ref_image.pixel_size.t,
+        levels=ref_meta.levels,
+        yx_scaling_factor=ref_meta.yx_scaling(),
+        z_scaling_factor=ref_meta.z_scaling(),
+        time_unit=ref_image.pixel_size.time_unit,
+        space_unit=ref_image.pixel_size.space_unit,
+        axes_names=axes_names,
+        chunks=chunks,
+        dtype=dtype,
+        overwrite=overwrite,
+        version=ref_meta.version,
+        name=name,
+    )
+    return None
+
+
+def build_masking_roi_table(label: Label) -> MaskingRoiTable:
+    """Compute the masking ROI table for a label."""
+    if label.dimensions.is_time_series:
+        raise NgioValueError("Time series labels are not supported.")
+
+    array = label.get_array(axes_order=["z", "y", "x"], mode="dask")
+
+    rois = compute_masking_roi(array, label.pixel_size)
+    return MaskingRoiTable(rois, reference_label=label.meta.name)

ngio/images/masked_image.py

@@ -0,0 +1,273 @@
+"""A module for handling label images in OME-NGFF files."""
+
+from collections.abc import Collection, Iterable
+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
+from ngio.images.label import Label
+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(
+        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:
+        """Return the array for a given ROI."""
+        roi = self._masking_roi_table.get(label)
+        roi = roi.zoom(zoom_factor)
+        return super().get_roi(
+            roi=roi, axes_order=axes_order, mode=mode, **slice_kwargs
+        )
+
+    def set_roi(
+        self,
+        label: int,
+        patch: ArrayLike,
+        zoom_factor: float = 1.0,
+        axes_order: Collection[str] | None = None,
+        **slice_kwargs: slice | int | Iterable[int],
+    ) -> None:
+        """Set the array for a given ROI."""
+        roi = self._masking_roi_table.get(label)
+        roi = roi.zoom(zoom_factor)
+        return super().set_roi(
+            roi=roi, patch=patch, axes_order=axes_order, **slice_kwargs
+        )
+
+    def get_roi_masked(
+        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=axes_order,
+            mode=mode,
+            zoom_factor=zoom_factor,
+            **slice_kwargs,
+        )
+
+    def set_roi_masked(
+        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=axes_order,
+            zoom_factor=zoom_factor,
+            **slice_kwargs,
+        )
+
+
+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(
+        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:
+        """Return the array for a given ROI."""
+        roi = self._masking_roi_table.get(label)
+        roi = roi.zoom(zoom_factor)
+        return super().get_roi(
+            roi=roi, axes_order=axes_order, mode=mode, **slice_kwargs
+        )
+
+    def set_roi(
+        self,
+        label: int,
+        patch: ArrayLike,
+        zoom_factor: float = 1.0,
+        axes_order: Collection[str] | None = None,
+        **slice_kwargs: slice | int | Iterable[int],
+    ) -> None:
+        """Set the array for a given ROI."""
+        roi = self._masking_roi_table.get(label)
+        roi = roi.zoom(zoom_factor)
+        return super().set_roi(
+            roi=roi, patch=patch, axes_order=axes_order, **slice_kwargs
+        )
+
+    def get_roi_masked(
+        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=axes_order,
+            mode=mode,
+            zoom_factor=zoom_factor,
+            **slice_kwargs,
+        )
+
+    def set_roi_masked(
+        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=axes_order,
+            zoom_factor=zoom_factor,
+            **slice_kwargs,
+        )
+
+
+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,
+    )