Feature/dask poc

README.rst

@@ -18,6 +18,19 @@ It can be built locally with:
     pip install spinx
     sphinx-build -b html docs/source/ docs/build/html
 
+Tests
+-----
+
+Tests can be run locally via `tox` with:
+
+    $ pip install tox
+    $ tox
+
+To enable pre-commit code validation:
+
+    $ pip install pre-commit
+    $ pre-commit install
+
 Release process
 ---------------
 

ome_zarr/dask_utils.py

@@ -0,0 +1,75 @@
+from typing import Any, Tuple
+
+import numpy as np
+import skimage.transform
+from dask import array as da
+
+# This module contributed by Andreas Eisenbarth @aeisenbarth
+# See https://github.com/toloudis/ome-zarr-py/pull/1
+
+
+def resize(
+    image: da.Array, output_shape: Tuple[int, ...], *args: Any, **kwargs: Any
+) -> da.Array:
+    r"""
+    Wrapped copy of "skimage.transform.resize"
+    Resize image to match a certain size.
+    :type image: :class:`dask.array`
+    :param image: The dask array to resize
+    :type output_shape: tuple
+    :param output_shape: The shape of the resize array
+    :type \*args: list
+    :param \*args: Arguments of skimage.transform.resize
+    :type \*\*kwargs: dict
+    :param \*\*kwargs: Keyword arguments of skimage.transform.resize
+    :return: Resized image.
+    """
+    factors = np.array(output_shape) / np.array(image.shape).astype(float)
+    # Rechunk the input blocks so that the factors achieve an output
+    # blocks size of full numbers.
+    better_chunksize = tuple(
+        np.maximum(1, np.round(np.array(image.chunksize) * factors) / factors).astype(
+            int
+        )
+    )
+    image_prepared = image.rechunk(better_chunksize)
+    block_output_shape = tuple(
+        np.floor(np.array(better_chunksize) * factors).astype(int)
+    )
+
+    # Map overlap
+    def resize_block(image_block: da.Array, block_info: dict) -> da.Array:
+        return skimage.transform.resize(
+            image_block, block_output_shape, *args, **kwargs
+        ).astype(image_block.dtype)
+
+    output_slices = tuple(slice(0, d) for d in output_shape)
+    output = da.map_blocks(
+        resize_block, image_prepared, dtype=image.dtype, chunks=block_output_shape
+    )[output_slices]
+    return output.rechunk(image.chunksize).astype(image.dtype)
+
+
+def downscale_nearest(image: da.Array, factors: Tuple[int, ...]) -> da.Array:
+    """
+    Primitive downscaling by integer factors using stepped slicing.
+    :type image: :class:`dask.array`
+    :param image: The dask array to resize
+    :type factors: tuple
+    :param factors: Sequence of integers factors for each dimension.
+    :return: Resized image.
+    """
+    if not len(factors) == image.ndim:
+        raise ValueError(
+            f"Dimension mismatch: {image.ndim} image dimensions, "
+            f"{len(factors)} scale factors"
+        )
+    if not (
+        all(isinstance(f, int) and 0 < f <= d for f, d in zip(factors, image.shape))
+    ):
+        raise ValueError(
+            f"All scale factors must not be greater than the dimension length: "
+            f"({tuple(factors)}) <= ({tuple(image.shape)})"
+        )
+    slices = tuple(slice(None, None, factor) for factor in factors)
+    return image[slices]

ome_zarr/reader.py

@@ -500,7 +500,7 @@ def get_pyramid_lazy(self, node: Node) -> None:
         well_node = Node(well_zarr, node)
         well_spec: Optional[Well] = well_node.first(Well)
         if well_spec is None:
-            raise Exception("could not find first well")
+            raise Exception("Could not find first well")
         self.numpy_type = well_spec.numpy_type
 
         LOGGER.debug(f"img_pyramid_shapes: {well_spec.img_pyramid_shapes}")

ome_zarr/scale.py

@@ -7,8 +7,9 @@
 import os
 from collections.abc import MutableMapping
 from dataclasses import dataclass
-from typing import Callable, Iterator, List
+from typing import Any, Callable, Iterator, List, Tuple, Union
 
+import dask.array as da
 import numpy as np
 import zarr
 from scipy.ndimage import zoom
@@ -19,10 +20,14 @@
     resize,
 )
 
+from .dask_utils import resize as dask_resize
 from .io import parse_url
 
 LOGGER = logging.getLogger("ome_zarr.scale")
 
+ListOfArrayLike = Union[List[da.Array], List[np.ndarray]]
+ArrayLike = Union[da.Array, np.ndarray]
+
 
 @dataclass
 class Scaler:
@@ -125,6 +130,30 @@ def __create_group(
             series.append({"path": path})
         return grp
 
+    def resize_image(self, image: ArrayLike) -> ArrayLike:
+        """
+        Resize a numpy array OR a dask array to a smaller array (not pyramid)
+        """
+        if isinstance(image, da.Array):
+
+            def _resize(image: ArrayLike, out_shape: Tuple, **kwargs: Any) -> ArrayLike:
+                return dask_resize(image, out_shape, **kwargs)
+
+        else:
+            _resize = resize
+
+        # only down-sample in X and Y dimensions for now...
+        new_shape = list(image.shape)
+        new_shape[-1] = np.ceil(float(image.shape[-1]) / self.downscale)
+        new_shape[-2] = np.ceil(float(image.shape[-2]) / self.downscale)
+        out_shape = tuple(new_shape)
+
+        dtype = image.dtype
+        image = _resize(
+            image.astype(float), out_shape, order=1, mode="reflect", anti_aliasing=False
+        )
+        return image.astype(dtype)
+
     def nearest(self, base: np.ndarray) -> List[np.ndarray]:
         """
         Downsample using :func:`skimage.transform.resize`.
@@ -133,9 +162,21 @@ def nearest(self, base: np.ndarray) -> List[np.ndarray]:
         """
         return self._by_plane(base, self.__nearest)
 
-    def __nearest(self, plane: np.ndarray, sizeY: int, sizeX: int) -> np.ndarray:
+    def __nearest(
+        self, plane: Union[np.ndarray, da.Array], sizeY: int, sizeX: int
+    ) -> np.ndarray:
         """Apply the 2-dimensional transformation."""
-        return resize(
+        if isinstance(plane, da.Array):
+
+            def _resize(
+                image: ArrayLike, output_shape: Tuple, **kwargs: Any
+            ) -> ArrayLike:
+                return dask_resize(image, output_shape, **kwargs)
+
+        else:
+            _resize = resize
+
+        return _resize(
             plane,
             output_shape=(sizeY // self.downscale, sizeX // self.downscale),
             order=0,