feat(Python): Add Pyodide PolyData support

InsightSoftwareConsortium · Apr 19, 2023 · c379d1e · c379d1e
1 parent fc9f404
commit c379d1e
Show file tree

Hide file tree

Showing 3 changed files with 74 additions and 3 deletions.
diff --git a/packages/core/python/itkwasm/itkwasm/polydata.py b/packages/core/python/itkwasm/itkwasm/polydata.py
@@ -27,7 +27,7 @@ def _default_points() -> ArrayLike:
 @dataclass
 class PolyData:
     polyDataType: Union[PolyDataType, Dict] = field(default_factory=PolyDataType)
-    name: str = 'polydata'
+    name: str = 'PolyData'
 
     numberOfPoints: int = 0
     points: ArrayLike = field(default_factory=_default_points)

diff --git a/packages/core/python/itkwasm/itkwasm/pyodide.py b/packages/core/python/itkwasm/itkwasm/pyodide.py
@@ -102,6 +102,27 @@ def to_py(js_proxy):
         if mesh_dict['cellData'] is not None:
             mesh_dict['cellData'] = _to_numpy_array(cell_pixel_component_type, mesh_dict['cellData'])
         return Mesh(**mesh_dict)
+    elif hasattr(js_proxy, "polyDataType"):
+        polydata_dict = js_proxy.to_py()
+        polydata_type = PolyDataType(**polydata_dict['polyDataType'])
+        polydata_dict['polyDataType'] = polydata_type
+        point_pixel_component_type = polydata_type.pointPixelComponentType
+        cell_pixel_component_type = polydata_type.cellPixelComponentType
+        if polydata_dict['points'] is not None:
+            polydata_dict['points'] = _to_numpy_array(str(FloatTypes.Float32), polydata_dict['points']).reshape((-1, 3))
+        if polydata_dict['vertices'] is not None:
+            polydata_dict['vertices'] = _to_numpy_array(str(IntTypes.UInt32), polydata_dict['vertices'])
+        if polydata_dict['lines'] is not None:
+            polydata_dict['lines'] = _to_numpy_array(str(IntTypes.UInt32), polydata_dict['lines'])
+        if polydata_dict['polygons'] is not None:
+            polydata_dict['polygons'] = _to_numpy_array(str(IntTypes.UInt32), polydata_dict['polygons'])
+        if polydata_dict['triangleStrips'] is not None:
+            polydata_dict['triangleStrips'] = _to_numpy_array(str(IntTypes.UInt32), polydata_dict['triangleStrips'])
+        if polydata_dict['pointData'] is not None:
+            polydata_dict['pointData'] = _to_numpy_array(point_pixel_component_type, polydata_dict['pointData'])
+        if polydata_dict['cellData'] is not None:
+            polydata_dict['cellData'] = _to_numpy_array(cell_pixel_component_type, polydata_dict['cellData'])
+        return PolyData(**polydata_dict)
     return js_proxy.to_py()
 
 def to_js(py):
@@ -132,5 +153,22 @@ def to_js(py):
         if mesh_dict['cellData'] is not None:
             mesh_dict['cellData'] = mesh_dict['cellData'].ravel()
         return pyodide.ffi.to_js(mesh_dict, dict_converter=js.Object.fromEntries)
+    elif isinstance(py, PolyData):
+        polydata_dict = asdict(py)
+        if polydata_dict['points'] is not None:
+            polydata_dict['points'] = polydata_dict['points'].ravel()
+        if polydata_dict['vertices'] is not None:
+            polydata_dict['vertices'] = polydata_dict['vertices'].ravel()
+        if polydata_dict['lines'] is not None:
+            polydata_dict['lines'] = polydata_dict['lines'].ravel()
+        if polydata_dict['polygons'] is not None:
+            polydata_dict['polygons'] = polydata_dict['polygons'].ravel()
+        if polydata_dict['triangleStrips'] is not None:
+            polydata_dict['triangleStrips'] = polydata_dict['triangleStrips'].ravel()
+        if polydata_dict['pointData'] is not None:
+            polydata_dict['pointData'] = polydata_dict['pointData'].ravel()
+        if polydata_dict['cellData'] is not None:
+            polydata_dict['cellData'] = polydata_dict['cellData'].ravel()
+        return pyodide.ffi.to_js(polydata_dict, dict_converter=js.Object.fromEntries)
 
     return py
diff --git a/packages/core/python/itkwasm/test/test_pyodide.py b/packages/core/python/itkwasm/test/test_pyodide.py
@@ -98,7 +98,7 @@ async def test_mesh_conversion(selenium, package_wheel):
     import micropip
     await micropip.install(package_wheel)
 
-    from itkwasm import Mesh, MeshType, FloatTypes
+    from itkwasm import Mesh, MeshType
     from itkwasm.pyodide import to_js, to_py
     import numpy as np
 
@@ -126,4 +126,37 @@ async def test_mesh_conversion(selenium, package_wheel):
     assert mesh.numberOfPoints == mesh_py.numberOfPoints
     assert np.array_equal(mesh.points, mesh_py.points)
     assert mesh.numberOfPointPixels == mesh_py.numberOfPointPixels
-    assert np.array_equal(mesh.pointData, mesh_py.pointData)
+    assert np.array_equal(mesh.pointData, mesh_py.pointData)
+
+@run_in_pyodide(packages=['micropip', 'numpy'])
+async def test_polydata_conversion(selenium, package_wheel):
+    import micropip
+    await micropip.install(package_wheel)
+
+    from itkwasm import PolyData, PolyDataType
+    from itkwasm.pyodide import to_js, to_py
+    import numpy as np
+
+    n_points = 5
+    dimension = 3
+
+    polydata_type = PolyDataType()
+
+    points = np.random.random((n_points, dimension)).astype(np.float32)
+    point_data = np.random.random((n_points,)).astype(np.float32)
+
+    polydata = PolyData(polydata_type, points=points, numberOfPoints=n_points, pointData=point_data, numberOfPointPixels=n_points)
+
+    polydata_js = to_js(polydata)
+    polydata_py = to_py(polydata_js)
+
+    polydata_type_py = polydata_py.polyDataType
+    assert polydata_type.pointPixelComponentType == polydata_type_py.pointPixelComponentType
+    assert polydata_type.pointPixelType == polydata_type_py.pointPixelType
+    assert polydata_type.pointPixelComponents == polydata_type_py.pointPixelComponents
+
+    assert polydata.name == polydata_py.name
+    assert polydata.numberOfPoints == polydata_py.numberOfPoints
+    assert np.array_equal(polydata.points, polydata_py.points)
+    assert polydata.numberOfPointPixels == polydata_py.numberOfPointPixels
+    assert np.array_equal(polydata.pointData, polydata_py.pointData)