Skip to content

Commit

Permalink
test(mesh-io): python read and write point set wasi tests
Browse files Browse the repository at this point in the history
  • Loading branch information
thewtex committed Oct 4, 2024
1 parent 04c79b0 commit fd4d5fe
Show file tree
Hide file tree
Showing 29 changed files with 666 additions and 9 deletions.
2 changes: 1 addition & 1 deletion packages/core/python/itkwasm/itkwasm/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@

from .interface_types import InterfaceTypes
from .image import Image, ImageType, ImageRegion
from .pointset import PointSet, PointSetType
from .point_set import PointSet, PointSetType
from .mesh import Mesh, MeshType
from .transform import Transform, TransformType, TransformParameterizations
from .polydata import PolyData, PolyDataType
Expand Down
1 change: 1 addition & 0 deletions packages/core/python/itkwasm/itkwasm/interface_types.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ class InterfaceTypes(str, Enum):
BinaryStream = "BinaryStream"
Image = "Image"
Mesh = "Mesh"
PointSet = "PointSet"
PolyData = "PolyData"
Transform = "Transform"
JsonCompatible = "JsonCompatible"
47 changes: 47 additions & 0 deletions packages/core/python/itkwasm/itkwasm/pipeline.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
from .binary_file import BinaryFile
from .image import Image
from .mesh import Mesh
from .point_set import PointSet
from .polydata import PolyData
from .json_compatible import JsonCompatible
from .int_types import IntTypes
Expand Down Expand Up @@ -277,6 +278,27 @@ def run(
"cellData": f"data:application/vnd.itk.address,0:{cell_data_ptr}",
}
ri.set_input_json(mesh_json, index)
elif input_.type == InterfaceTypes.PointSet:
point_set = input_.data
if point_set.numberOfPoints:
pv = array_like_to_bytes(point_set.points)
else:
pv = bytes([])
points_ptr = ri.set_input_array(pv, index, 0)
if point_set.numberOfPointPixels:
pdv = array_like_to_bytes(point_set.pointData)
else:
pdv = bytes([])
point_data_ptr = ri.set_input_array(pdv, index, 1)
point_set_json = {
"pointSetType": asdict(point_set.pointSetType),
"name": point_set.name,
"numberOfPoints": point_set.numberOfPoints,
"points": f"data:application/vnd.itk.address,0:{points_ptr}",
"numberOfPointPixels": point_set.numberOfPointPixels,
"pointData": f"data:application/vnd.itk.address,0:{point_data_ptr}",
}
ri.set_input_json(point_set_json, index)
elif input_.type == InterfaceTypes.PolyData:
polydata = input_.data
if polydata.numberOfPoints:
Expand Down Expand Up @@ -442,6 +464,31 @@ def run(
mesh.cellData = buffer_to_numpy_array(mesh.meshType.cellPixelComponentType, bytes([]))

output_data = PipelineOutput(InterfaceTypes.Mesh, mesh)
elif output.type == InterfaceTypes.PointSet:
point_set_json = ri.get_output_json(index)
point_set = PointSet(**point_set_json)

if point_set.numberOfPoints > 0:
data_ptr = ri.get_output_array_address(0, index, 0)
data_size = ri.get_output_array_size(0, index, 0)
point_set.points = buffer_to_numpy_array(
point_set.pointSetType.pointComponentType,
ri.wasmtime_lift(data_ptr, data_size),
)
else:
point_set.points = buffer_to_numpy_array(point_set.pointSetType.pointComponentType, bytes([]))

if point_set.numberOfPointPixels > 0:
data_ptr = ri.get_output_array_address(0, index, 1)
data_size = ri.get_output_array_size(0, index, 1)
point_set.pointData = buffer_to_numpy_array(
point_set.pointSetType.pointPixelComponentType,
ri.wasmtime_lift(data_ptr, data_size),
)
else:
point_set.pointData = buffer_to_numpy_array(point_set.pointSetType.pointPixelComponentType, bytes([]))

output_data = PipelineOutput(InterfaceTypes.PointSet, point_set)
elif output.type == InterfaceTypes.PolyData:
polydata_json = ri.get_output_json(index)
polydata = PolyData(**polydata_json)
Expand Down
2 changes: 1 addition & 1 deletion packages/core/python/itkwasm/itkwasm/pyodide.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
from typing import Optional

from .image import Image, ImageType
from .pointset import PointSet, PointSetType
from .point_set import PointSet, PointSetType
from .mesh import Mesh, MeshType
from .polydata import PolyData, PolyDataType
from .binary_file import BinaryFile
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,9 @@
from .read_mesh_async import read_mesh_async
from .write_mesh_async import write_mesh_async

from .read_point_set_async import read_point_set_async
from .write_point_set_async import write_point_set_async

from .byu_read_mesh_async import byu_read_mesh_async
from .byu_write_mesh_async import byu_write_mesh_async
from .free_surfer_ascii_read_mesh_async import free_surfer_ascii_read_mesh_async
Expand All @@ -24,4 +27,15 @@
from .wasm_zstd_read_mesh_async import wasm_zstd_read_mesh_async
from .wasm_zstd_write_mesh_async import wasm_zstd_write_mesh_async

from .obj_read_mesh_async import obj_read_mesh_async
from .obj_write_mesh_async import obj_write_mesh_async
from .off_read_mesh_async import off_read_mesh_async
from .off_write_mesh_async import off_write_mesh_async
from .vtk_poly_data_read_mesh_async import vtk_poly_data_read_mesh_async
from .vtk_poly_data_write_mesh_async import vtk_poly_data_write_mesh_async
from .wasm_read_mesh_async import wasm_read_mesh_async
from .wasm_write_mesh_async import wasm_write_mesh_async
from .wasm_zstd_read_mesh_async import wasm_zstd_read_mesh_async
from .wasm_zstd_write_mesh_async import wasm_zstd_write_mesh_async

from ._version import __version__
Original file line number Diff line number Diff line change
Expand Up @@ -12,5 +12,4 @@
('.iwm', 'wasm'),
('.iwm.cbor', 'wasm'),
('.iwm.cbor.zst', 'wasmZstd'),
('.bmp', 'bmp'),
])
Original file line number Diff line number Diff line change
@@ -0,0 +1,10 @@
from collections import OrderedDict

extension_to_point_set_io = OrderedDict([
('.vtk', 'vtkPolyData'),
('.obj', 'obj'),
('.off', 'off'),
('.iwm', 'wasm'),
('.iwm.cbor', 'wasm'),
('.iwm.cbor.zst', 'wasmZstd'),
])
Original file line number Diff line number Diff line change
@@ -0,0 +1,8 @@
point_set_io_index = [
'vtk_poly_data',
'obj',
'off',
'wasm',
'wasm_zstd',
]

Original file line number Diff line number Diff line change
@@ -0,0 +1,87 @@
import os
from typing import Optional, Union
from pathlib import Path

from itkwasm import (
PointSet,
BinaryFile,
)

from .js_package import js_package

from itkwasm.pyodide import (
to_js,
to_py,
js_resources
)

from .extension_to_point_set_io import extension_to_point_set_io
from .point_set_io_index import point_set_io_index

async def read_point_set_async(
serialized_point_set: os.PathLike,
information_only: bool = False,
) -> PointSet:
"""Read an point set file format and convert it to the ITK-Wasm file format.
:param serialized_point_set: Input point set serialized in the file format
:type serialized_point_set: os.PathLike
:param information_only: Only read point set metadata -- do not read pixel data.
:type information_only: bool
:return: Output point set
:rtype: PointSet
"""
js_module = await js_package.js_module
web_worker = js_resources.web_worker

kwargs = {}
if information_only:
kwargs["informationOnly"] = to_js(information_only)

extension = ''.join(Path(serialized_point_set).suffixes)

io = None
if extension in extension_to_point_set_io:
func = f"{extension_to_point_set_io[extension]}ReadPointSet"
io = getattr(js_module, func)
else:
for ioname in point_set_io_index:
func = f"{ioname}ReadPointSet"
io = getattr(js_module, func)
outputs = await io(to_js(BinaryFile(serialized_point_set)), webWorker=web_worker, noCopy=True, **kwargs)
outputs_object_map = outputs.as_object_map()
web_worker = outputs_object_map['webWorker']
js_resources.web_worker = web_worker
could_read = to_py(outputs_object_map['couldRead'])
if could_read:
point_set = to_py(outputs_object_map['point_set'])
return point_set

if io is None:
raise RuntimeError(f"Could not find an point_set reader for {extension}")

outputs = await io(to_js(BinaryFile(serialized_point_set)), webWorker=web_worker, noCopy=True, **kwargs)
outputs_object_map = outputs.as_object_map()
web_worker = outputs_object_map['webWorker']
could_read = to_py(outputs_object_map['couldRead'])

if not could_read:
raise RuntimeError(f"Could not read {serialized_point_set}")

js_resources.web_worker = web_worker

point_set = to_py(outputs_object_map['point_set'])

return point_set

async def pointsetread_async(
serialized_point_set: os.PathLike,
information_only: bool = False,
) -> PointSet:
return await read_point_set_async(serialized_point_set, information_only=information_only)

pointsetread_async.__doc__ = f"""{read_point_set_async.__doc__}
Alias for read_point_set_async.
"""
Original file line number Diff line number Diff line change
@@ -0,0 +1,94 @@
import os
import importlib
from pathlib import Path
from typing import Optional, Union

from itkwasm import PointSet, PixelTypes, IntTypes, FloatTypes, BinaryFile

from itkwasm.pyodide import (
to_js,
to_py,
js_resources
)

from .js_package import js_package

from .extension_to_point_set_io import extension_to_point_set_io
from .point_set_io_index import point_set_io_index

async def write_point_set_async(
point_set: PointSet,
serialized_point_set: os.PathLike,
information_only: bool = False,
use_compression: bool = False,
) -> None:
"""Write an ITK-Wasm PointSet to a point set file format.
:param point_set: Input point set
:type point_set: PointSet
:param serialized_point_set: Output point set serialized in the file format.
:type serialized_point_set: str
:param information_only: Only write point set metadata -- do not write pixel data.
:type information_only: bool
:param use_compression: Use compression in the written file
:type use_compression: bool
:param serialized_point_set: Input point set serialized in the file format
:type serialized_point_set: os.PathLike
"""
js_module = await js_package.js_module
web_worker = js_resources.web_worker

kwargs = {}
if information_only:
kwargs["informationOnly"] = to_js(information_only)
if use_compression:
kwargs["useCompression"] = to_js(use_compression)

extension = ''.join(Path(serialized_point_set).suffixes)

io = None
if extension in extension_to_point_set_io:
func = f"{extension_to_point_set_io[extension]}WritePointSet"
io = getattr(js_module, func)
else:
for ioname in point_set_io_index:
func = f"{ioname}WritePointSet"
io = getattr(js_module, func)
outputs = await io(to_js(point_set), to_js(serialized_point_set), webWorker=web_worker, noCopy=True, **kwargs)
outputs_object_map = outputs.as_object_map()
web_worker = outputs_object_map['webWorker']
js_resources.web_worker = web_worker
could_write = to_py(outputs_object_map['couldWrite'])
if could_write:
to_py(outputs_object_map['serializedPointSet'])
return

if io is None:
raise RuntimeError(f"Could not find an point_set writer for {extension}")

outputs = await io(to_js(point_set), to_js(serialized_point_set), webWorker=web_worker, noCopy=True, **kwargs)
outputs_object_map = outputs.as_object_map()
web_worker = outputs_object_map['webWorker']
js_resources.web_worker = web_worker
could_write = to_py(outputs_object_map['couldWrite'])

if not could_write:
raise RuntimeError(f"Could not write {serialized_point_set}")

to_py(outputs_object_map['serializedPointSet'])

async def point_setwrite_async(
point_set: PointSet,
serialized_point_set: os.PathLike,
information_only: bool = False,
use_compression: bool = False,
) -> None:
return write_point_set_async(point_set, serialized_point_set, information_only=information_only, use_compression=use_compression)

point_setwrite_async.__doc__ = f"""{write_point_set_async.__doc__}
Alias for write_point_set.
"""
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,9 @@
from .read_mesh import read_mesh, meshread
from .write_mesh import write_mesh, meshwrite

from .read_point_set import read_point_set, pointsetread
from .write_point_set import write_point_set, pointsetwrite

from .byu_read_mesh import byu_read_mesh
from .byu_write_mesh import byu_write_mesh
from .free_surfer_ascii_read_mesh import free_surfer_ascii_read_mesh
Expand All @@ -24,4 +27,15 @@
from .wasm_zstd_read_mesh import wasm_zstd_read_mesh
from .wasm_zstd_write_mesh import wasm_zstd_write_mesh

from .obj_read_point_set import obj_read_point_set
from .obj_write_point_set import obj_write_point_set
from .off_read_point_set import off_read_point_set
from .off_write_point_set import off_write_point_set
from .vtk_poly_data_read_point_set import vtk_poly_data_read_point_set
from .vtk_poly_data_write_point_set import vtk_poly_data_write_point_set
from .wasm_read_point_set import wasm_read_point_set
from .wasm_write_point_set import wasm_write_point_set
from .wasm_zstd_read_point_set import wasm_zstd_read_point_set
from .wasm_zstd_write_point_set import wasm_zstd_write_point_set

from ._version import __version__
Original file line number Diff line number Diff line change
Expand Up @@ -12,5 +12,4 @@
('.iwm', 'wasm'),
('.iwm.cbor', 'wasm'),
('.iwm.cbor.zst', 'wasm_zstd'),
('.bmp', 'bmp'),
])
Original file line number Diff line number Diff line change
@@ -0,0 +1,10 @@
from collections import OrderedDict

extension_to_point_set_io = OrderedDict([
('.vtk', 'vtk_poly_data'),
('.obj', 'obj'),
('.off', 'off'),
('.iwm', 'wasm'),
('.iwm.cbor', 'wasm'),
('.iwm.cbor.zst', 'wasm_zstd'),
])
Original file line number Diff line number Diff line change
@@ -0,0 +1,8 @@
point_set_io_index = [
'vtk_poly_data',
'obj',
'off',
'wasm',
'wasm_zstd',
]

Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ def read_mesh(
serialized_mesh: os.PathLike,
information_only: bool = False,
) -> Mesh:
"""Read an mesh file format and convert it to the itk-wasm file format.
"""Read a mesh file format and convert it to the ITK-Wasm file format.
:param serialized_mesh: Input mesh serialized in the file format
:type serialized_mesh: os.PathLike
Expand Down
Loading

0 comments on commit fd4d5fe

Please sign in to comment.