Merge pull request #340 from thewtex/vtk-js-zarr

ENH: Add vtkjs_to_zarr, zarr_to_vtkjs

Author: thewtex

itkwidgets/_transform_types.py

@@ -1,4 +1,4 @@
-__all__ = ['to_itk_image', 'to_point_set', 'to_geometry']
+__all__ = ['to_itk_image', 'to_point_set', 'to_geometry', 'vtkjs_to_zarr', 'zarr_to_vtkjs']
 
 import itk
 import numpy as np
@@ -41,6 +41,77 @@ def is_arraylike(arr):
     have_skan = True
 except ImportError:
     pass
+have_zarr = False
+try:
+    import zarr
+    have_zarr = True
+except ImportError:
+    pass
+
+def vtkjs_to_zarr(vtkjs, group, chunks=True):
+    """Convert a vtk.js-like Python object to a Zarr Group.
+
+    Parameters
+    ----------
+
+    vtkjs: dictionary, required
+        The vtk.js-like data structure to convert.
+
+    group: zarr.Group, required
+        The Zarr group to store the result.
+
+    chunks: bool or int or tuple of ints, optional
+        The chunk size passed to zarr.creation.create.
+    """
+    for key, value in vtkjs.items():
+        if key == 'vtkClass':
+            group.attrs[key] = value
+        elif key == 'arrays':
+            for index, arr in enumerate(value):
+                vtkjs_to_zarr(arr,
+                              group.create_group('arrays/' + str(index), True),
+                              chunks=chunks)
+        elif isinstance(value, dict):
+            vtkjs_to_zarr(value,
+                          group.create_group(key, True),
+                          chunks=chunks)
+        elif isinstance(value, np.ndarray):
+            group.array(key, value, chunks=chunks)
+        else:
+            group.attrs[key] = value
+    return group
+
+def zarr_to_vtkjs(group):
+    """Convert Zarr Group that contains vtk.js data structure to a Python-like object.
+
+    Parameters
+    ----------
+
+    group: zarr.Group, required
+        The Zarr group to convert.
+    """
+
+    def process_group(group, result):
+        for key, value in group.attrs.items():
+            result[key] = value
+        for name, value in group.arrays():
+            result[name] = np.asarray(value)
+        for name, value in group.groups():
+            if name == 'arrays':
+                nested = []
+                for index, subgroup in value.groups():
+                    subresult = dict()
+                    process_group(subgroup, subresult)
+                    nested.append(subresult)
+                result[name] = nested
+            else:
+                nested = dict()
+                process_group(value, nested)
+                result[name] = nested
+    result = dict()
+    process_group(group, result)
+    return result
+
 
 _itk_pixel_to_vtkjs_type_components = {
     itk.SC: ('Int8Array', 1),
@@ -53,7 +124,6 @@ def is_arraylike(arr):
     itk.D: ('Float64Array', 1),
 }
 
-
 def _vtk_to_vtkjs(data_array):
     from vtk.util.numpy_support import vtk_to_numpy
     # From vtkType.h
@@ -315,6 +385,8 @@ def to_point_set(point_set_like):  # noqa: C901
             point_set['cellData'] = vtkjs_cell_data
 
         return point_set
+    elif isinstance(point_set_like, zarr.Group):
+        return zarr_to_vtkjs(point_set_like)
 
     return None
 
@@ -542,5 +614,7 @@ def to_geometry(geometry_like):  # noqa: C901
         geometry_filter.Update()
         geometry = to_geometry(geometry_filter.GetOutput())
         return geometry
+    elif isinstance(geometry_like, zarr.Group):
+        return zarr_to_vtkjs(geometry_like)
 
     return None

tests/test_transform_types.py

@@ -2,7 +2,7 @@
 import numpy as np
 import pytest
 
-from itkwidgets._transform_types import to_point_set, to_geometry, to_itk_image
+from itkwidgets._transform_types import to_point_set, to_geometry, to_itk_image, vtkjs_to_zarr, zarr_to_vtkjs
 
 
 def test_mesh_to_geometry():
@@ -78,6 +78,71 @@ def test_mesh_to_geometry():
     # assert(np.alltrue(geometry['points']['values'] ==
     # geometry_array.astype(np.float32)))
 
+def test_vtkjs_to_zarr():
+    zarr = pytest.importorskip("zarr")
+
+    # 3D
+    Dimension = 3
+    PixelType = itk.ctype('double')
+    MeshType = itk.Mesh[PixelType, Dimension]
+    mesh = MeshType.New()
+    PointType = itk.Point[itk.F, Dimension]
+    point0 = PointType()
+    point0[0] = -1
+    point0[1] = -1
+    point0[2] = 0
+    mesh.SetPoint(0, point0)
+    mesh.SetPointData(0, 8.0)
+    point1 = PointType()
+    point1[0] = 1
+    point1[1] = -1
+    point1[2] = 0
+    mesh.SetPointData(1, 9.0)
+    mesh.SetPoint(1, point1)
+    point2 = PointType()
+    point2[0] = 1
+    point2[1] = 1
+    point2[2] = 0
+    mesh.SetPoint(2, point2)
+    mesh.SetPointData(2, 19.0)
+    point3 = PointType()
+    point3[0] = 1
+    point3[1] = 1
+    point3[2] = 0
+    mesh.SetPoint(3, point3)
+    mesh.SetPointData(3, 24.0)
+
+    vtkjs = to_geometry(mesh)
+
+    store = zarr.group()
+
+    vtkjs_to_zarr(vtkjs, store)
+    geometry = zarr_to_vtkjs(store)
+
+    points = mesh.GetPoints()
+    point_template = itk.template(points)
+    element_type = point_template[1][1]
+    point_values = itk.PyVectorContainer[element_type].array_from_vector_container(
+        points)
+
+    assert(geometry['vtkClass'] == 'vtkPolyData')
+    assert(geometry['points']['vtkClass'] == 'vtkPoints')
+    assert(geometry['points']['numberOfComponents'] == 3)
+    assert(geometry['points']['dataType'] == 'Float32Array')
+    assert(geometry['points']['size'] == 4 * 3)
+    assert(np.array_equal(geometry['points']['values'],
+                          point_values.astype(np.float32)))
+    assert(geometry['pointData']['vtkClass'] == 'vtkDataSetAttributes')
+    assert(geometry['pointData']['arrays'][0]
+           ['data']['vtkClass'] == 'vtkDataArray')
+    assert(geometry['pointData']['arrays'][0]['data']['name'] == 'Point Data')
+    assert(geometry['pointData']['arrays'][0]
+           ['data']['numberOfComponents'] == 1)
+    assert(geometry['pointData']['arrays'][0]['data']['size'] == 4)
+    assert(geometry['pointData']['arrays'][0]
+           ['data']['dataType'] == 'Float64Array')
+    assert(np.array_equal(geometry['pointData']['arrays'][0]['data']['values'],
+                          np.array([8.0, 9.0, 19.0, 24.0], dtype=np.float64)))
 
 def test_vtkpolydata_to_geometry():
     vtk = pytest.importorskip("vtk")