diff --git a/.github/workflows/ci-workflows.yml b/.github/workflows/ci-workflows.yml index e99fa01..6e88069 100644 --- a/.github/workflows/ci-workflows.yml +++ b/.github/workflows/ci-workflows.yml @@ -29,7 +29,7 @@ jobs: apt: - '^libxcb.*-dev' - libxkbcommon-x11-0 - - libegl1-mesa + - libegl1-mesa-dev - libhdf5-dev envs: | - linux: py39-test-all diff --git a/glue_ar/__init__.py b/glue_ar/__init__.py index a6f3850..ac5c27f 100644 --- a/glue_ar/__init__.py +++ b/glue_ar/__init__.py @@ -11,7 +11,7 @@ def setup_common(): from .common.usd_builder import USDBuilder # noqa: F401 from .common.stl_builder import STLBuilder # noqa: F401 - from .compression import compress_gltfpack, compress_gltf_pipeline # noqa: F401 + from .compression import compress_draco, compress_meshoptimizer # noqa: F401 def setup_qt(): diff --git a/glue_ar/common/scatter_gltf.py b/glue_ar/common/scatter_gltf.py index 1cdb2c9..b394035 100644 --- a/glue_ar/common/scatter_gltf.py +++ b/glue_ar/common/scatter_gltf.py @@ -1,3 +1,4 @@ +from collections import defaultdict from gltflib import AccessorType, BufferTarget, ComponentType, PrimitiveMode from glue_vispy_viewers.common.viewer_state import Vispy3DViewerState from glue_vispy_viewers.scatter.layer_state import ScatterLayerState @@ -11,9 +12,10 @@ from glue_ar.common.scatter_export_options import ARIpyvolumeScatterExportOptions, ARVispyScatterExportOptions from glue_ar.common.shapes import cone_triangles, cone_points, cylinder_points, cylinder_triangles, \ normalize, rectangular_prism_triangulation, sphere_triangles -from glue_ar.gltf_utils import add_points_to_bytearray, add_triangles_to_bytearray, index_mins, index_maxes -from glue_ar.utils import Viewer3DState, get_stretches, iterable_has_nan, hex_to_components, layer_color, \ - unique_id, xyz_bounds, xyz_for_layer, Bounds +from glue_ar.gltf_utils import SHORT_MAX, add_points_to_bytearray, add_triangles_to_bytearray, \ + index_mins, index_maxes +from glue_ar.utils import Viewer3DState, get_stretches, iterable_has_nan, hex_to_components, \ + layer_color, offset_triangles, unique_id, xyz_bounds, xyz_for_layer, Bounds from glue_ar.common.gltf_builder import GLTFBuilder from glue_ar.common.scatter import Scatter3DLayerState, ScatterLayerState3D, \ PointsGetter, box_points_getter, IPYVOLUME_POINTS_GETTERS, \ @@ -208,7 +210,8 @@ def add_scatter_layer_gltf(builder: GLTFBuilder, points_getter: PointsGetter, triangles: List[Tuple[int, int, int]], bounds: Bounds, - clip_to_bounds: bool = True): + clip_to_bounds: bool = True, + points_per_mesh: Optional[int] = None): if layer_state is None: return @@ -227,34 +230,6 @@ def add_scatter_layer_gltf(builder: GLTFBuilder, scaled=True) data = data[:, [1, 2, 0]] - barr = bytearray() - add_triangles_to_bytearray(barr, triangles) - triangles_len = len(barr) - max_index = max(idx for tri in triangles for idx in tri) - - buffer = builder.buffer_count - builder.add_buffer_view( - buffer=buffer, - byte_length=triangles_len, - byte_offset=0, - target=BufferTarget.ELEMENT_ARRAY_BUFFER, - ) - builder.add_accessor( - buffer_view=builder.buffer_view_count-1, - component_type=ComponentType.UNSIGNED_INT, - count=len(triangles)*3, - type=AccessorType.SCALAR, - mins=[0], - maxes=[max_index], - ) - sphere_triangles_accessor = builder.accessor_count - 1 - - first_material_index = builder.material_count - if fixed_color: - color = layer_color(layer_state) - color_components = hex_to_components(color) - builder.add_material(color=color_components, opacity=layer_state.alpha) - buffer = builder.buffer_count cmap = layer_state.cmap cmap_attr = "cmap_attribute" if vispy_layer_state else "cmap_att" @@ -264,43 +239,242 @@ def add_scatter_layer_gltf(builder: GLTFBuilder, uri = f"layer_{unique_id()}.bin" sizes = sizes_for_scatter_layer(layer_state, bounds, mask) - for i, point in enumerate(data): - prev_len = len(barr) - size = radius if fixed_size else sizes[i] - pts = points_getter(point, size) - add_points_to_bytearray(barr, pts) - point_mins = index_mins(pts) - point_maxes = index_maxes(pts) + barr = bytearray() + n_points = len(data) - if not fixed_color: - cval = cmap_vals[i] - normalized = max(min((cval - layer_state.cmap_vmin) / crange, 1), 0) - cindex = int(normalized * 255) - color = cmap(cindex) - builder.add_material(color, layer_state.alpha) + # If points per mesh is not specified, + # we don't do any mesh chunking. + # Note that this will work for colormapping as well + if points_per_mesh is None: + points_per_mesh = n_points + first_material_index = builder.material_count + if fixed_color: + points = [] + tris = [] + + color = layer_color(layer_state) + color_components = hex_to_components(color) + builder.add_material(color=color_components, opacity=layer_state.alpha) + + for i, point in enumerate(data): + size = radius if fixed_size else sizes[i] + pts = points_getter(point, size) + points.append(pts) + + # If n_points is less than our designated chunk size, we only want + # to make triangles for that many points (and put everything in one mesh). + # This is both more space-efficient and necessary to be glTF spec-compliant + triangle_offset = 0 + pts_count = len(points_getter((0, 0, 0), 1)) + for _ in range(min(points_per_mesh, n_points)): + pt_triangles = offset_triangles(triangles, triangle_offset) + triangle_offset += pts_count + tris.append(pt_triangles) + + mesh_triangles = [tri for sphere in tris for tri in sphere] + max_triangle_index = max(idx for tri in mesh_triangles for idx in tri) + use_short = max_triangle_index <= SHORT_MAX + triangles_start = len(barr) + add_triangles_to_bytearray(barr, mesh_triangles, short=use_short) + triangles_len = len(barr) builder.add_buffer_view( buffer=buffer, - byte_length=len(barr)-prev_len, - byte_offset=prev_len, - target=BufferTarget.ARRAY_BUFFER, + byte_length=triangles_len-triangles_start, + byte_offset=triangles_start, + target=BufferTarget.ELEMENT_ARRAY_BUFFER, ) + component_type = ComponentType.UNSIGNED_SHORT if use_short else ComponentType.UNSIGNED_INT builder.add_accessor( buffer_view=builder.buffer_view_count-1, - component_type=ComponentType.FLOAT, - count=len(pts), - type=AccessorType.VEC3, - mins=point_mins, - maxes=point_maxes, + component_type=component_type, + count=len(mesh_triangles)*3, + type=AccessorType.SCALAR, + mins=[0], + maxes=[max_triangle_index], ) - material_index = builder.material_count - 1 - builder.add_mesh( - position_accessor=builder.accessor_count-1, - indices_accessor=sphere_triangles_accessor, - material=material_index, - ) + start = 0 + triangles_accessor = builder.accessor_count - 1 + while start < n_points: + mesh_points = [pt for pts in points[start:start+points_per_mesh] for pt in pts] + barr_offset = len(barr) + add_points_to_bytearray(barr, mesh_points) + point_mins = index_mins(mesh_points) + point_maxes = index_maxes(mesh_points) + + builder.add_buffer_view( + buffer=buffer, + byte_length=len(barr)-barr_offset, + byte_offset=barr_offset, + target=BufferTarget.ARRAY_BUFFER, + ) + builder.add_accessor( + buffer_view=builder.buffer_view_count-1, + component_type=ComponentType.FLOAT, + count=len(mesh_points), + type=AccessorType.VEC3, + mins=point_mins, + maxes=point_maxes, + ) + points_accessor = builder.accessor_count - 1 + + # This should only happen on the final iteration + # or not at all, if points_per_mesh is a divisor of count + # But in this case we do need a separate accessor as the + # byte length is different. + # Note that if we're on the first chunk (start == 0) + # there's no need to do this - we can use the buffer view that we just created + count = n_points - start + if start != 0 and count < points_per_mesh: + triangles_count = len(tris) + byte_length = count * triangles_len // triangles_count + mesh_triangles = [tri for sphere in tris[:count] for tri in sphere] + max_triangle_index = max(idx for tri in mesh_triangles for idx in tri) + builder.add_buffer_view( + buffer=buffer, + byte_length=byte_length, + byte_offset=triangles_start, + target=BufferTarget.ELEMENT_ARRAY_BUFFER, + ) + component_type = ComponentType.UNSIGNED_SHORT if use_short else ComponentType.UNSIGNED_INT + builder.add_accessor( + buffer_view=builder.buffer_view_count-1, + component_type=component_type, + count=len(triangles)*3*count, + type=AccessorType.SCALAR, + mins=[0], + maxes=[max_triangle_index], + ) + triangles_accessor = builder.accessor_count - 1 + + builder.add_mesh( + position_accessor=points_accessor, + indices_accessor=triangles_accessor, + material=builder.material_count-1, + ) + start += points_per_mesh + + else: + points_by_color = defaultdict(list) + color_materials = defaultdict(int) + + for i, point in enumerate(data): + cval = cmap_vals[i] + normalized = max(min((cval - layer_state.cmap_vmin) / crange, 1), 0) + cindex = int(normalized * 255) + color = cmap(cindex) + + material_index = color_materials.get(cindex, None) + if material_index is None: + builder.add_material(color, layer_state.alpha) + material_index = builder.material_count - 1 + color_materials[cindex] = material_index + + size = radius if fixed_size else sizes[i] + pts = points_getter(point, size) + points_by_color[cindex].append(pts) + + for cindex, points in points_by_color.items(): + + # If the maximum number of points in any one color is less than our designated chunk size, + # we only want to make triangles for that many points (and put everything in one mesh). + # This is both more space-efficient and necessary to be glTF spec-compliant + triangle_offset = 0 + tris = [] + pts_count = len(points_getter((0, 0, 0), 1)) + for _ in range(min(points_per_mesh, len(points))): + pt_triangles = offset_triangles(triangles, triangle_offset) + triangle_offset += pts_count + tris.append(pt_triangles) + + mesh_points = [pt for pts in points for pt in pts] + mesh_triangles = [tri for sphere in tris for tri in sphere] + max_triangle_index = max(idx for tri in mesh_triangles for idx in tri) + use_short = max_triangle_index <= SHORT_MAX + triangles_start = len(barr) + add_triangles_to_bytearray(barr, mesh_triangles, short=use_short) + triangles_len = len(barr) + + builder.add_buffer_view( + buffer=buffer, + byte_length=triangles_len-triangles_start, + byte_offset=triangles_start, + target=BufferTarget.ELEMENT_ARRAY_BUFFER, + ) + component_type = ComponentType.UNSIGNED_SHORT if use_short else ComponentType.UNSIGNED_INT + builder.add_accessor( + buffer_view=builder.buffer_view_count-1, + component_type=component_type, + count=len(mesh_triangles)*3, + type=AccessorType.SCALAR, + mins=[0], + maxes=[max_triangle_index], + ) + + start = 0 + triangles_accessor = builder.accessor_count - 1 + n_points = len(points) + while start < n_points: + mesh_points = [pt for pts in points[start:start+points_per_mesh] for pt in pts] + barr_offset = len(barr) + add_points_to_bytearray(barr, mesh_points) + point_mins = index_mins(mesh_points) + point_maxes = index_maxes(mesh_points) + + builder.add_buffer_view( + buffer=buffer, + byte_length=len(barr)-barr_offset, + byte_offset=barr_offset, + target=BufferTarget.ARRAY_BUFFER, + ) + builder.add_accessor( + buffer_view=builder.buffer_view_count-1, + component_type=ComponentType.FLOAT, + count=len(mesh_points), + type=AccessorType.VEC3, + mins=point_mins, + maxes=point_maxes, + ) + points_accessor = builder.accessor_count - 1 + + # This should only happen on the final iteration + # or not at all, if points_per_mesh is a divisor of count + # But in this case we do need a separate accessor as the + # byte length is different. + # Note that if we're on the first chunk (start == 0) + # there's no need to do this - we can use the buffer view that we just created + count = n_points - start + if start != 0 and count < points_per_mesh: + triangles_count = len(tris) + byte_length = count * triangles_len // triangles_count + mesh_triangles = [tri for sphere in tris[:count] for tri in sphere] + max_triangle_index = max(idx for tri in mesh_triangles for idx in tri) + builder.add_buffer_view( + buffer=buffer, + byte_length=byte_length, + byte_offset=triangles_start, + target=BufferTarget.ELEMENT_ARRAY_BUFFER, + ) + component_type = ComponentType.UNSIGNED_SHORT if use_short else ComponentType.UNSIGNED_INT + builder.add_accessor( + buffer_view=builder.buffer_view_count-1, + component_type=component_type, + count=len(triangles)*3*count, + type=AccessorType.SCALAR, + mins=[0], + maxes=[max_triangle_index], + ) + triangles_accessor = builder.accessor_count - 1 + + material = color_materials[cindex] + builder.add_mesh( + position_accessor=points_accessor, + indices_accessor=triangles_accessor, + material=material, + ) + start += points_per_mesh builder.add_buffer(byte_length=len(barr), uri=uri) builder.add_file_resource(uri, data=barr) @@ -358,6 +532,7 @@ def add_vispy_scatter_layer_gltf(builder: GLTFBuilder, points_getter = sphere_points_getter(theta_resolution=theta_resolution, phi_resolution=phi_resolution) + points_per_mesh = 1 add_scatter_layer_gltf(builder=builder, viewer_state=viewer_state, @@ -365,7 +540,8 @@ def add_vispy_scatter_layer_gltf(builder: GLTFBuilder, points_getter=points_getter, triangles=triangles, bounds=bounds, - clip_to_bounds=clip_to_bounds) + clip_to_bounds=clip_to_bounds, + points_per_mesh=points_per_mesh) @ar_layer_export(Scatter3DLayerState, "Scatter", ARIpyvolumeScatterExportOptions, ("gltf", "glb")) @@ -380,6 +556,7 @@ def add_ipyvolume_scatter_layer_gltf(builder: GLTFBuilder, triangle_getter = IPYVOLUME_TRIANGLE_GETTERS.get(geometry, rectangular_prism_triangulation) triangles = triangle_getter() points_getter = IPYVOLUME_POINTS_GETTERS.get(geometry, box_points_getter) + points_per_mesh = 1 add_scatter_layer_gltf(builder=builder, viewer_state=viewer_state, @@ -387,4 +564,5 @@ def add_ipyvolume_scatter_layer_gltf(builder: GLTFBuilder, points_getter=points_getter, triangles=triangles, bounds=bounds, - clip_to_bounds=clip_to_bounds) + clip_to_bounds=clip_to_bounds, + points_per_mesh=points_per_mesh) diff --git a/glue_ar/common/tests/gltf_helpers.py b/glue_ar/common/tests/gltf_helpers.py index 85476bf..4903a12 100644 --- a/glue_ar/common/tests/gltf_helpers.py +++ b/glue_ar/common/tests/gltf_helpers.py @@ -9,7 +9,7 @@ from glue_ar.utils import iterator_count -BufferFormat = Union[Literal["f"], Literal["I"]] +BufferFormat = Union[Literal["f"], Literal["I"], Literal["H"]] def get_data(gltf: GLTF, buffer: Buffer, buffer_view: Optional[BufferView] = None) -> bytes: @@ -41,8 +41,9 @@ def count_vertices(gltf: GLTF, buffer: Buffer, buffer_view: BufferView): return count_points(gltf, buffer, buffer_view, 'f') -def count_indices(gltf: GLTF, buffer: Buffer, buffer_view: BufferView): - return count_points(gltf, buffer, buffer_view, 'I') +def count_indices(gltf: GLTF, buffer: Buffer, buffer_view: BufferView, use_short=False): + format = 'H' if use_short else 'I' + return count_points(gltf, buffer, buffer_view, format) def unpack_points(gltf: GLTF, diff --git a/glue_ar/common/tests/test_scatter_gltf.py b/glue_ar/common/tests/test_scatter_gltf.py index dabc985..d6dc4b8 100644 --- a/glue_ar/common/tests/test_scatter_gltf.py +++ b/glue_ar/common/tests/test_scatter_gltf.py @@ -10,6 +10,7 @@ from glue_ar.common.tests.gltf_helpers import count_indices, count_vertices, unpack_vertices from glue_ar.common.tests.helpers import APP_VIEWER_OPTIONS from glue_ar.common.tests.test_scatter import BaseScatterTest +from glue_ar.gltf_utils import SHORT_MAX from glue_ar.utils import export_label_for_layer, hex_to_components, layers_to_export, mask_for_bounds, \ xyz_bounds, xyz_for_layer @@ -68,8 +69,8 @@ def test_basic_export(self, app_type: str, viewer_type: str): phi_resolution=phi_resolution) points_count = sphere_points_count(theta_resolution=theta_resolution, phi_resolution=phi_resolution) - - assert count_indices(gltf, model.buffers[0], model.bufferViews[0]) == triangles_count + use_short = points_count <= SHORT_MAX + assert count_indices(gltf, model.buffers[0], model.bufferViews[0], use_short=use_short) == triangles_count assert count_vertices(gltf, model.buffers[0], model.bufferViews[1]) == points_count assert model.bufferViews[0].target == BufferTarget.ELEMENT_ARRAY_BUFFER.value @@ -77,7 +78,8 @@ def test_basic_export(self, app_type: str, viewer_type: str): indices_accessor = model.accessors[0] assert indices_accessor.bufferView == 0 - assert indices_accessor.componentType == ComponentType.UNSIGNED_INT.value + expected_indices_type = ComponentType.UNSIGNED_SHORT if use_short else ComponentType.UNSIGNED_INT + assert indices_accessor.componentType == expected_indices_type.value assert indices_accessor.count == triangles_count * 3 assert indices_accessor.type == AccessorType.SCALAR.value assert indices_accessor.min == [0] diff --git a/glue_ar/compression.py b/glue_ar/compression.py index 1e49686..0941fcc 100644 --- a/glue_ar/compression.py +++ b/glue_ar/compression.py @@ -6,15 +6,14 @@ NODE_MODULES_DIR = join(PACKAGE_DIR, "js", "node_modules") -GLTF_PIPELINE_FILEPATH = join(NODE_MODULES_DIR, "gltf-pipeline", "bin", "gltf-pipeline.js") -GLTFPACK_FILEPATH = join(NODE_MODULES_DIR, "gltfpack", "cli.js") +GLTF_TRANSFORM_FILEPATH = join(NODE_MODULES_DIR, "@gltf-transform", "cli", "bin", "cli.js") @compressor("draco") -def compress_gltf_pipeline(filepath: str): - run(["node", GLTF_PIPELINE_FILEPATH, "-i", filepath, "-o", filepath, "-d"], capture_output=True) +def compress_draco(filepath: str): + run([GLTF_TRANSFORM_FILEPATH, "optimize", filepath, filepath, "--compress", "draco"], capture_output=True) @compressor("meshoptimizer") -def compress_gltfpack(filepath: str): - run(["node", GLTFPACK_FILEPATH, "-i", filepath, "-o", filepath], capture_output=True) +def compress_meshoptimizer(filepath: str): + run([GLTF_TRANSFORM_FILEPATH, "optimize", filepath, filepath], capture_output=True) diff --git a/glue_ar/gltf_utils.py b/glue_ar/gltf_utils.py index a0e592e..ceab78c 100644 --- a/glue_ar/gltf_utils.py +++ b/glue_ar/gltf_utils.py @@ -18,6 +18,8 @@ "meshoptimizer": "EXT_meshopt_compression", } +SHORT_MAX = 65_535 + def create_material_for_color( color: List[int], @@ -40,10 +42,13 @@ def add_points_to_bytearray(arr: bytearray, points: Iterable[Iterable[Union[int, arr.extend(struct.pack('f', coordinate)) -def add_triangles_to_bytearray(arr: bytearray, triangles: Iterable[Iterable[int]]): +def add_triangles_to_bytearray(arr: bytearray, + triangles: Iterable[Iterable[int]], + short: bool = False): + format = "H" if short else "I" for triangle in triangles: for index in triangle: - arr.extend(struct.pack('I', index)) + arr.extend(struct.pack(format, index)) T = TypeVar("T", bound=Union[int, float]) diff --git a/js/package.json b/js/package.json index 527f9eb..f97edf2 100644 --- a/js/package.json +++ b/js/package.json @@ -1,7 +1,6 @@ { "dependencies": { - "gltf-pipeline": "^4.1.0", - "gltfpack": "^0.21.0" + "@gltf-transform/cli": "^4.1.1" }, "scripts": { "glue-ar-export": "npm install && node glue-ar-export.js"