Merge pull request #527 from thewtex/server-side-chunking

thewtex · web-flow · commit 9aede3dbd969 · 2022-09-20T19:39:00.000-07:00
ENH: Server-side multi-scale chunking of images
diff --git a/itkwidgets/_type_aliases.py b/itkwidgets/_type_aliases.py
@@ -1,8 +1,8 @@
 import itkwasm
 import zarr
 import numpy as np
+import dask
 
-from .integrations.dask import HAVE_DASK
 from .integrations.itk import HAVE_ITK
 from .integrations.pytorch import HAVE_TORCH
 from .integrations.vtk import HAVE_VTK
@@ -23,10 +23,8 @@
     import vtk
     Image = Union[Image, vtk.vtkImageData]
     Point_Sets = Union[Point_Sets, vtk.vtkPolyData]
-if HAVE_DASK:
-    import dask
-    Image = Union[Image, dask.array.core.Array]
-    Point_Sets = Union[Point_Sets, dask.array.core.Array]
+Image = Union[Image, dask.array.core.Array]
+Point_Sets = Union[Point_Sets, dask.array.core.Array]
 if HAVE_TORCH:
     import torch
     Image = Union[Image, torch.Tensor]
diff --git a/itkwidgets/imjoy.py b/itkwidgets/imjoy.py
@@ -1,10 +1,11 @@
-from dataclasses import dataclass, asdict
+from dataclasses import asdict
 
-from typing import Dict
+from typing import Dict 
 
 import itkwasm
 import numcodecs
 from imjoy_rpc import api
+import zarr
 
 _numcodec_encoder = numcodecs.Blosc(cname='lz4', clevel=3)
 _numcodec_config = _numcodec_encoder.get_config()
@@ -24,6 +25,26 @@ def encode_itkwasm_image(image):
 
     return image_dict
 
+def encode_zarr_store(store):
+    def getItem(key):
+        return store[key]
+
+    def setItem(key, value):
+        store[key] = value
+
+    def containsItem(key):
+        return key in store
+
+    return {
+        "_rintf": True,
+        "_rtype": 'zarr-store',
+        "getItem": getItem,
+        "setItem": setItem,
+        "containsItem": containsItem,
+    }
+
 def register_itkwasm_imjoy_codecs():
 
     api.registerCodec({'name': 'itkwasm-image', 'type': itkwasm.Image, 'encoder': encode_itkwasm_image})
+    api.registerCodec({'name': 'zarr-store', 'type': zarr.storage.BaseStore, 'encoder': encode_zarr_store})
+ 
diff --git a/itkwidgets/integrations/__init__.py b/itkwidgets/integrations/__init__.py
@@ -1,60 +1,148 @@
+from collections.abc import MutableMapping
+from re import S
+
 import itkwasm
 import numpy as np
 import zarr
+from multiscale_spatial_image import MultiscaleSpatialImage, to_multiscale, itk_image_to_multiscale
+from spatial_image import to_spatial_image, is_spatial_image
 
-from .dask import HAVE_DASK, dask_array_to_ndarray
+import dask
+import xarray as xr
 from .itk import HAVE_ITK, itk_image_to_wasm_image, itk_group_spatial_object_to_wasm_point_set
 from .pytorch import HAVE_TORCH
-from .vtk import HAVE_VTK, vtk_image_to_ndarray, vtk_polydata_to_vtkjs
-from .xarray import HAVE_XARRAY, xarray_data_array_to_numpy, xarray_data_set_to_numpy
+from .vtk import HAVE_VTK, vtk_image_to_spatial_image, vtk_polydata_to_vtkjs
+from .xarray import xarray_data_array_to_numpy, xarray_data_set_to_numpy
 from ..render_types import RenderType
 
+def _spatial_image_scale_factors(spatial_image, min_length):
+    sizes = dict(spatial_image.sizes)
+    scale_factors = []
+    dims = spatial_image.dims
+    previous = { d: 1 for d in { 'x', 'y', 'z' }.intersection(dims) }
+    while (np.array(list(sizes.values())) > min_length).any():
+        max_size = np.array(list(sizes.values())).max()
+        to_skip = { d: sizes[d] <= max_size / 2 for d in previous.keys() }
+        scale_factor = {}
+        for dim in previous.keys():
+            if to_skip[dim]:
+                scale_factor[dim] = previous[dim]
+                continue
+            scale_factor[dim] = 2 * previous[dim]
+
+            sizes[dim] = int(sizes[dim] / 2)
+        previous = scale_factor
+        scale_factors.append(scale_factor)
+
+    return scale_factors
+
+def _make_multiscale_store(multiscale):
+    # Todo: for very large images serialize to disk cache
+    store = zarr.storage.MemoryStore(dimension_separator='/')
+    multiscale.to_zarr(store, compute=True)
+    return store
 
 def _get_viewer_image(image):
+    min_length = 64
+    if isinstance(image, MultiscaleSpatialImage):
+        return _make_multiscale_store(image)
+
+    # Todo: support for itkwasm.Image
     if HAVE_ITK:
         import itk
         if isinstance(image, itk.Image):
-            return itk_image_to_wasm_image(image)
+            dimension = image.GetImageDimension()
+            size = np.array(itk.size(image))
+            scale_factors = []
+            dims = ('x', 'y', 'z')
+            previous = {'x': 1, 'y': 1, 'z': 1}
+            while (size > min_length).any():
+                to_skip = size <= size.max() / 2
+                scale_factor = {}
+                for dim in range(dimension):
+                    if to_skip[dim]:
+                        scale_factor[dims[dim]] = previous[dims[dim]]
+                        continue
+                    scale_factor[dims[dim]] = 2 * previous[dims[dim]]
+                    size[dim] = int(size[dim] / 2)
+                previous = scale_factor
+                scale_factors.append(scale_factor)
+
+            multiscale = itk_image_to_multiscale(image, scale_factors=scale_factors)
+            return _make_multiscale_store(multiscale)
+
     if HAVE_VTK:
         import vtk
         if isinstance(image, vtk.vtkImageData):
-            return vtk_image_to_ndarray(image)
-    if HAVE_DASK:
-        import dask
-        if isinstance(image, dask.array.core.Array):
-            return dask_array_to_ndarray(image)
+            spatial_image = vtk_image_to_spatial_image(image)
+            scale_factors = _spatial_image_scale_factors(spatial_image, min_length)
+            multiscale = to_multiscale(spatial_image, scale_factors)
+            return _make_multiscale_store(multiscale)
+
+    if isinstance(image, dask.array.core.Array):
+        spatial_image = to_spatial_image(image)
+        scale_factors = _spatial_image_scale_factors(spatial_image, min_length)
+        multiscale = to_multiscale(spatial_image, scale_factors)
+        return _make_multiscale_store(multiscale)
+
+    if isinstance(image, zarr.Array):
+        spatial_image = to_spatial_image(image)
+        scale_factors = _spatial_image_scale_factors(spatial_image, min_length)
+        multiscale = to_multiscale(spatial_image, scale_factors)
+        return _make_multiscale_store(multiscale)
+
+    # NGFF Zarr
+    if isinstance(image, zarr.Group) and 'multiscales' in image.attrs:
+        return _make_multiscale_store(image.store)
+
     if HAVE_TORCH:
         import torch
         if isinstance(image, torch.Tensor):
-            return image.numpy()
-    if HAVE_XARRAY:
-        import xarray
-        if isinstance(image, xarray.DataArray):
-            return xarray_data_array_to_numpy(image)
-        if isinstance(image, xarray.Dataset):
-            return xarray_data_set_to_numpy(image)
-    return image
+            spatial_image = to_spatial_image(image.numpy())
+            scale_factors = _spatial_image_scale_factors(spatial_image, min_length)
+            multiscale = to_multiscale(spatial_image, scale_factors)
+            return _make_multiscale_store(multiscale)
+
+    # Todo: preserve dask Array, if present, check if dims are NGFF -> use dims, coords
+    # Check if coords are uniform, if not, resample
+    if isinstance(image, xr.DataArray):
+        if is_spatial_image(image):
+            scale_factors = _spatial_image_scale_factors(image, min_length)
+            multiscale = to_multiscale(image, scale_factors)
+            return _make_multiscale_store(multiscale)
+
+        return xarray_data_array_to_numpy(image)
+    if isinstance(image, xr.Dataset):
+        da = image[next(iter(image.variables.keys()))]
+        if is_spatial_image(da):
+            scale_factors = _spatial_image_scale_factors(da, min_length)
+            multiscale = to_multiscale(da, scale_factors)
+            return _make_multiscale_store(multiscale)
+        return xarray_data_set_to_numpy(image)
+
+    if isinstance(image, np.ndarray):
+        spatial_image = to_spatial_image(image)
+        scale_factors = _spatial_image_scale_factors(spatial_image, min_length)
+        multiscale = to_multiscale(spatial_image, scale_factors)
+        return _make_multiscale_store(multiscale)
+    raise RuntimeError("Could not process the viewer image")
 
 
 def _get_viewer_point_sets(point_sets):
     if HAVE_VTK:
         import vtk
         if isinstance(point_sets, vtk.vtkPolyData):
             return vtk_polydata_to_vtkjs(point_sets)
-    if HAVE_DASK:
-        import dask
-        if isinstance(point_sets, dask.array.core.Array):
-            return dask_array_to_ndarray(point_sets)
+    if isinstance(point_sets, dask.array.core.Array):
+        return np.asarray(point_sets)
     if HAVE_TORCH:
         import torch
         if isinstance(point_sets, torch.Tensor):
             return point_sets.numpy()
-    if HAVE_XARRAY:
-        import xarray
-        if isinstance(point_sets, xarray.DataArray):
-            return xarray_data_array_to_numpy(point_sets)
-        if isinstance(point_sets, xarray.Dataset):
-            return xarray_data_set_to_numpy(point_sets)
+    if isinstance(point_sets, xr.DataArray):
+        return xarray_data_array_to_numpy(point_sets)
+    if isinstance(point_sets, xr.Dataset):
+        return xarray_data_set_to_numpy(point_sets)
     return point_sets
 
 
@@ -63,6 +151,13 @@ def _detect_render_type(data, input_type) -> RenderType:
         return RenderType.IMAGE
     elif isinstance(data, itkwasm.PointSet):
         return RenderType.POINT_SET
+    elif isinstance(data, MultiscaleSpatialImage):
+        return RenderType.IMAGE
+    elif isinstance(data, (zarr.Array, zarr.Group)):
+        # For now assume zarr.Group is an image
+        # In the future, once NGFF supports point sets fully
+        # We may need to do more introspection
+        return RenderType.IMAGE
     elif isinstance(data, np.ndarray):
         if input_type == 'point_sets':
             return RenderType.POINT_SET
@@ -83,29 +178,25 @@ def _detect_render_type(data, input_type) -> RenderType:
             return RenderType.IMAGE
         elif isinstance(data, vtk.vtkPolyData):
             return RenderType.POINT_SET
-    if HAVE_DASK:
-        import dask
-        if isinstance(data, dask.array.core.Array):
-            if input_type == 'point_sets':
-                return RenderType.POINT_SET
-            else:
-                return RenderType.IMAGE
+    if isinstance(data, dask.array.core.Array):
+        if input_type == 'point_sets':
+            return RenderType.POINT_SET
+        else:
+            return RenderType.IMAGE
     if HAVE_TORCH:
         import torch
         if isinstance(data, torch.Tensor):
             if input_type == 'point_sets':
                 return RenderType.POINT_SET
             else:
                 return RenderType.IMAGE
-    if HAVE_XARRAY:
-        import xarray
-        if isinstance(data, xarray.DataArray):
-            if input_type == 'point_sets':
-                return RenderType.POINT_SET
-            else:
-                return RenderType.IMAGE
-        if isinstance(data, xarray.Dataset):
-            if input_type == 'point_sets':
-                return RenderType.POINT_SET
-            else:
-                return RenderType.IMAGE
+    if isinstance(data, xr.DataArray):
+        if input_type == 'point_sets':
+            return RenderType.POINT_SET
+        else:
+            return RenderType.IMAGE
+    if isinstance(data, xr.Dataset):
+        if input_type == 'point_sets':
+            return RenderType.POINT_SET
+        else:
+            return RenderType.IMAGE
diff --git a/itkwidgets/integrations/dask.py b/itkwidgets/integrations/dask.py
diff --git a/itkwidgets/integrations/vtk.py b/itkwidgets/integrations/vtk.py
@@ -6,12 +6,23 @@
 except ImportError:
     pass
 
+from spatial_image import to_spatial_image
 
-def vtk_image_to_ndarray(image):
+
+def vtk_image_to_spatial_image(image):
     array = vtk_to_numpy(image.GetPointData().GetScalars())
-    dims = list(image.GetDimensions())
-    array.shape = dims[::-1]
-    return array
+    dimensions = list(image.GetDimensions())
+    array.shape = dimensions[::-1]
+
+    origin = image.GetOrigin()
+    translation = { 'x': origin[0], 'y': origin[1], 'z': origin[2] }
+
+    spacing = image.GetSpacing()
+    scale = { 'x': spacing[0], 'y': spacing[1], 'z': spacing[2] }
+
+    spatial_image = to_spatial_image(array, scale=scale, translation=translation)
+
+    return spatial_image
 
 def vtk_polydata_to_vtkjs(point_set):
     array = vtk_to_numpy(point_set.GetPoints().GetData())
diff --git a/pyproject.toml b/pyproject.toml
@@ -39,6 +39,7 @@ dependencies = [
     "imjoy-rpc >= 0.5.13",
     "imjoy-utils >= 0.1.2",
     "numcodecs",
+    "multiscale_spatial_image >= 0.10.1",
     "zarr",
 ]
 

Original file line number	Diff line number	Diff line change
`@@ -39,6 +39,7 @@ dependencies = [`
`39`	`39`	`"imjoy-rpc >= 0.5.13",`
`40`	`40`	`"imjoy-utils >= 0.1.2",`
`41`	`41`	`"numcodecs",`
	`42`	`+ "multiscale_spatial_image >= 0.10.1",`
`42`	`43`	`"zarr",`
`43`	`44`	`]`
`44`	`45`