wip new lazy to_multiscale() for labels

Author: LucaMarconato

src/spatialdata/_io/io_raster.py

@@ -13,7 +13,7 @@
 from ome_zarr.writer import write_labels as write_labels_ngff
 from ome_zarr.writer import write_multiscale as write_multiscale_ngff
 from ome_zarr.writer import write_multiscale_labels as write_multiscale_labels_ngff
-from xarray import DataArray, Dataset, DataTree
+from xarray import DataArray, DataTree
 
 from spatialdata._io._utils import (
     _get_transformations_from_ngff_dict,
@@ -27,6 +27,7 @@
 from spatialdata._utils import get_pyramid_levels
 from spatialdata.models._utils import get_channel_names
 from spatialdata.models.models import ATTRS_KEY
+from spatialdata.models.pyramids_utils import dask_arrays_to_datatree
 from spatialdata.transformations._utils import (
     _get_transformations,
     _get_transformations_xarray,
@@ -91,20 +92,8 @@ def _read_multiscale(
             channels = [d["label"] for d in omero_metadata["channels"]]
     axes = [i["name"] for i in node.metadata["axes"]]
     if len(datasets) > 1:
-        multiscale_image = {}
-        for i, d in enumerate(datasets):
-            data = node.load(Multiscales).array(resolution=d)
-            multiscale_image[f"scale{i}"] = Dataset(
-                {
-                    "image": DataArray(
-                        data,
-                        name="image",
-                        dims=axes,
-                        coords={"c": channels} if channels is not None else {},
-                    )
-                }
-            )
-        msi = DataTree.from_dict(multiscale_image)
+        arrays = [node.load(Multiscales).array(resolution=d) for d in datasets]
+        msi = dask_arrays_to_datatree(arrays, dims=axes, channels=channels)
         _set_transformations(msi, transformations)
         return compute_coordinates(msi)
 

src/spatialdata/datasets.py

@@ -20,7 +20,15 @@
 from spatialdata._core.query.relational_query import get_element_instances
 from spatialdata._core.spatialdata import SpatialData
 from spatialdata._types import ArrayLike
-from spatialdata.models import Image2DModel, Labels2DModel, PointsModel, ShapesModel, TableModel
+from spatialdata.models import (
+    Image2DModel,
+    Image3DModel,
+    Labels2DModel,
+    Labels3DModel,
+    PointsModel,
+    ShapesModel,
+    TableModel,
+)
 from spatialdata.transformations import Identity
 
 __all__ = ["blobs", "raccoon"]
@@ -172,37 +180,47 @@ def _image_blobs(
         n_channels: int = 3,
         c_coords: str | list[str] | None = None,
         multiscale: bool = False,
+        ndim: int = 2,
     ) -> DataArray | DataTree:
         masks = []
         for i in range(n_channels):
-            mask = self._generate_blobs(length=length, seed=i)
+            mask = self._generate_blobs(length=length, seed=i, ndim=ndim)
             mask = (mask - mask.min()) / np.ptp(mask)
             masks.append(mask)
 
         x = np.stack(masks, axis=0)
-        dims = ["c", "y", "x"]
+        if ndim == 2:
+            dims = ["c", "y", "x"]
+            model = Image2DModel
+        else:
+            dims = ["c", "z", "y", "x"]
+            model = Image3DModel
         if not multiscale:
-            return Image2DModel.parse(x, transformations=transformations, dims=dims, c_coords=c_coords)
-        return Image2DModel.parse(
-            x, transformations=transformations, dims=dims, c_coords=c_coords, scale_factors=[2, 2]
-        )
+            return model.parse(x, transformations=transformations, dims=dims, c_coords=c_coords)
+        return model.parse(x, transformations=transformations, dims=dims, c_coords=c_coords, scale_factors=[2, 2])
 
     def _labels_blobs(
-        self, transformations: dict[str, Any] | None = None, length: int = 512, multiscale: bool = False
+        self,
+        transformations: dict[str, Any] | None = None,
+        length: int = 512,
+        multiscale: bool = False,
+        ndim: int = 2,
     ) -> DataArray | DataTree:
-        """Create a 2D labels."""
+        """Create labels in 2D or 3D."""
         from scipy.ndimage import watershed_ift
 
         # from skimage
-        mask = self._generate_blobs(length=length)
+        mask = self._generate_blobs(length=length, ndim=ndim)
         threshold = np.percentile(mask, 100 * (1 - 0.3))
         inputs = np.logical_not(mask < threshold).astype(np.uint8)
         # use watershed from scipy
-        xm, ym = np.ogrid[0:length:10, 0:length:10]
+        grid = np.ogrid[tuple(slice(0, length, 10) for _ in range(ndim))]
         markers = np.zeros_like(inputs).astype(np.int16)
-        markers[xm, ym] = np.arange(xm.size * ym.size).reshape((xm.size, ym.size))
+        grid_shape = tuple(g.size for g in grid)
+        markers[tuple(grid)] = np.arange(np.prod(grid_shape)).reshape(grid_shape)
         out = watershed_ift(inputs, markers)
-        out[xm, ym] = out[xm - 1, ym - 1]  # remove the isolate seeds
+        shifted = tuple(g - 1 for g in grid)
+        out[tuple(grid)] = out[tuple(shifted)]  # remove the isolated seeds
         # reindex by frequency
         val, counts = np.unique(out, return_counts=True)
         sorted_idx = np.argsort(counts)
@@ -211,20 +229,25 @@ def _labels_blobs(
                 out[out == val[idx]] = 0
             else:
                 out[out == val[idx]] = i
-        dims = ["y", "x"]
+        if ndim == 2:
+            dims = ["y", "x"]
+            model = Labels2DModel
+        else:
+            dims = ["z", "y", "x"]
+            model = Labels3DModel
         if not multiscale:
-            return Labels2DModel.parse(out, transformations=transformations, dims=dims)
-        return Labels2DModel.parse(out, transformations=transformations, dims=dims, scale_factors=[2, 2])
+            return model.parse(out, transformations=transformations, dims=dims)
+        return model.parse(out, transformations=transformations, dims=dims, scale_factors=[2, 2])
 
-    def _generate_blobs(self, length: int = 512, seed: int | None = None) -> ArrayLike:
+    def _generate_blobs(self, length: int = 512, seed: int | None = None, ndim: int = 2) -> ArrayLike:
         from scipy.ndimage import gaussian_filter
 
         rng = default_rng(42) if seed is None else default_rng(seed)
         # from skimage
-        shape = tuple([length] * 2)
+        shape = (length,) * ndim
         mask = np.zeros(shape)
-        n_pts = max(int(1.0 / 0.1) ** 2, 1)
-        points = (length * rng.random((2, n_pts))).astype(int)
+        n_pts = max(int(1.0 / 0.1) ** ndim, 1)
+        points = (length * rng.random((ndim, n_pts))).astype(int)
         mask[tuple(indices for indices in points)] = 1
         mask = gaussian_filter(mask, sigma=0.25 * length * 0.1)
         assert isinstance(mask, np.ndarray)

src/spatialdata/models/models.py

@@ -14,7 +14,7 @@
 from dask.array.core import from_array
 from dask.dataframe import DataFrame as DaskDataFrame
 from geopandas import GeoDataFrame, GeoSeries
-from multiscale_spatial_image import to_multiscale
+from multiscale_spatial_image import to_multiscale as to_multiscale_msi
 from multiscale_spatial_image.to_multiscale.to_multiscale import Methods
 from pandas import CategoricalDtype
 from shapely._geometry import GeometryType
@@ -38,17 +38,15 @@
     _validate_mapping_to_coordinate_system_type,
     convert_region_column_to_categorical,
 )
+from spatialdata.models.pyramids_utils import Chunks_t, ScaleFactors_t
+from spatialdata.models.pyramids_utils import to_multiscale as to_multiscale_ozp  # ozp -> ome-zarr-py
 from spatialdata.transformations._utils import (
     _get_transformations,
     _set_transformations,
     compute_coordinates,
 )
 from spatialdata.transformations.transformations import Identity
 
-# Types
-Chunks_t: TypeAlias = int | tuple[int, ...] | tuple[tuple[int, ...], ...] | Mapping[Any, None | int | tuple[int, ...]]
-ScaleFactors_t = Sequence[dict[str, int] | int]
-
 ATTRS_KEY = "spatialdata_attrs"
 
 
@@ -225,12 +223,19 @@ def parse(
                 chunks = {dim: chunks[index] for index, dim in enumerate(data.dims)}
             if isinstance(chunks, float):
                 chunks = {dim: chunks for index, dim in data.dims}
-            data = to_multiscale(
-                data,
-                scale_factors=scale_factors,
-                method=method,
-                chunks=chunks,
-            )
+            if method is not None:
+                data = to_multiscale_msi(
+                    data,
+                    scale_factors=scale_factors,
+                    method=method,
+                    chunks=chunks,
+                )
+            else:
+                data = to_multiscale_ozp(
+                    data,
+                    scale_factors=scale_factors,
+                    chunks=chunks,
+                )
             _parse_transformations(data, parsed_transform)
         else:
             # Chunk single scale images
@@ -375,9 +380,6 @@ def parse(  # noqa: D102
     ) -> DataArray | DataTree:
         if kwargs.get("c_coords") is not None:
             raise ValueError("`c_coords` is not supported for labels")
-        if kwargs.get("scale_factors") is not None and kwargs.get("method") is None:
-            # Override default scaling method to preserve labels
-            kwargs["method"] = Methods.DASK_IMAGE_NEAREST
         return super().parse(*args, **kwargs)
 
 
@@ -388,9 +390,6 @@ class Labels3DModel(RasterSchema):
     def parse(self, *args: Any, **kwargs: Any) -> DataArray | DataTree:  # noqa: D102
         if kwargs.get("c_coords") is not None:
             raise ValueError("`c_coords` is not supported for labels")
-        if kwargs.get("scale_factors") is not None and kwargs.get("method") is None:
-            # Override default scaling method to preserve labels
-            kwargs["method"] = Methods.DASK_IMAGE_NEAREST
         return super().parse(*args, **kwargs)
 
 

src/spatialdata/models/pyramids_utils.py

@@ -0,0 +1,84 @@
+from collections.abc import Mapping, Sequence
+from typing import Any, TypeAlias
+
+import dask.array as da
+from ome_zarr.dask_utils import resize
+from xarray import DataArray, Dataset, DataTree
+
+Chunks_t: TypeAlias = int | tuple[int, ...] | tuple[tuple[int, ...], ...] | Mapping[Any, None | int | tuple[int, ...]]
+ScaleFactors_t = Sequence[dict[str, int] | int]
+
+
+def dask_arrays_to_datatree(
+    arrays: Sequence[da.Array],
+    dims: Sequence[str],
+    channels: list[Any] | None = None,
+) -> DataTree:
+    """Build a multiscale DataTree from a sequence of dask arrays.
+
+    Parameters
+    ----------
+    arrays
+        Sequence of dask arrays, one per scale level (scale0, scale1, ...).
+    dims
+        Dimension names for the arrays (e.g. ``("c", "y", "x")``).
+    channels
+        Optional channel coordinate values. If provided, a ``"c"`` coordinate
+        is added to each scale level.
+
+    Returns
+    -------
+    DataTree with one child per scale level.
+    """
+    coords = {"c": channels} if channels is not None else {}
+    d = {}
+    for i, arr in enumerate(arrays):
+        d[f"scale{i}"] = Dataset(
+            {
+                "image": DataArray(
+                    arr,
+                    name="image",
+                    dims=list(dims),
+                    coords=coords,
+                )
+            }
+        )
+    return DataTree.from_dict(d)
+
+
+def to_multiscale(
+    image: DataArray,
+    scale_factors: ScaleFactors_t,
+    chunks: Chunks_t | None = None,
+) -> DataTree:
+    dims = [str(dim) for dim in image.dims]
+    spatial_dims = [d for d in dims if d != "c"]
+    order = 1 if "c" in dims else 0
+    pyramid = [image.data]
+    for sf in scale_factors:
+        prev = pyramid[-1]
+        # Compute per-axis scale factors: int applies to spatial axes only, dict to specific ones.
+        sf_by_axis = dict.fromkeys(dims, 1)
+        if isinstance(sf, int):
+            sf_by_axis.update(dict.fromkeys(spatial_dims, sf))
+        else:
+            sf_by_axis.update(sf)
+        # Clamp: skip axes where the scale factor exceeds the axis size.
+        for ax, factor in sf_by_axis.items():
+            ax_size = prev.shape[dims.index(ax)]
+            if factor > ax_size:
+                sf_by_axis[ax] = 1
+        output_shape = tuple(prev.shape[dims.index(ax)] // f for ax, f in sf_by_axis.items())
+        resized = resize(
+            image=prev.astype(float),
+            output_shape=output_shape,
+            order=order,
+            mode="reflect",
+            anti_aliasing=False,
+        )
+        pyramid.append(resized.astype(prev.dtype))
+    if chunks is not None:
+        if isinstance(chunks, Mapping):
+            chunks = {dims.index(k) if isinstance(k, str) else k: v for k, v in chunks.items()}
+        pyramid = [arr.rechunk(chunks) for arr in pyramid]
+    return dask_arrays_to_datatree(pyramid, dims=dims)

tests/models/test_pyramids_utils.py

@@ -0,0 +1,52 @@
+import dask
+import numpy as np
+import pytest
+from multiscale_spatial_image.to_multiscale.to_multiscale import Methods
+
+from spatialdata.datasets import BlobsDataset
+from spatialdata.models import Image2DModel, Image3DModel, Labels2DModel, Labels3DModel
+
+CHUNK_SIZE = 32
+
+
+@pytest.mark.parametrize(
+    ("model", "length", "ndim", "n_channels", "scale_factors", "method"),
+    [
+        (Image2DModel, 128, 2, 3, (2, 2), Methods.XARRAY_COARSEN),
+        (Image3DModel, 32, 3, 3, (2, 2), Methods.XARRAY_COARSEN),
+        (Labels2DModel, 128, 2, 0, (2, 2), Methods.DASK_IMAGE_NEAREST),
+        (Labels3DModel, 32, 3, 0, (2, 2), Methods.DASK_IMAGE_NEAREST),
+    ],
+)
+def test_to_multiscale_via_ome_zarr_scaler(model, length, ndim, n_channels, scale_factors, method):
+    blob_gen = BlobsDataset()
+
+    if n_channels > 0:
+        # Image: stack multiple blob channels
+        masks = []
+        for i in range(n_channels):
+            mask = blob_gen._generate_blobs(length=length, seed=i, ndim=ndim)
+            mask = (mask - mask.min()) / np.ptp(mask)
+            masks.append(mask)
+        array = np.stack(masks, axis=0)
+    else:
+        # Labels: threshold blob pattern to get integer labels
+        mask = blob_gen._generate_blobs(length=length, ndim=ndim)
+        threshold = np.percentile(mask, 70)
+        array = (mask >= threshold).astype(np.int64)
+
+    dims = model.dims
+    dask_data = dask.array.from_array(array).rechunk(CHUNK_SIZE)
+
+    # multiscale-spatial-image path (explicit method)
+    result_msi = model.parse(dask_data, dims=dims, scale_factors=scale_factors, chunks=CHUNK_SIZE, method=method)
+
+    # ome-zarr-py scaler path (method=None triggers the ome-zarr-py scaler)
+    result_ozp = model.parse(dask_data, dims=dims, scale_factors=scale_factors, chunks=CHUNK_SIZE)
+
+    # Compare data values at each scale level
+    for scale_name in result_msi.children:
+        msi_arr = result_msi[scale_name].ds["image"]
+        ozp_arr = result_ozp[scale_name].ds["image"]
+        assert msi_arr.sizes == ozp_arr.sizes
+        np.testing.assert_allclose(msi_arr.values, ozp_arr.values)