diff --git a/glue_ar/common/marching_cubes.py b/glue_ar/common/marching_cubes.py index d815804..5937568 100644 --- a/glue_ar/common/marching_cubes.py +++ b/glue_ar/common/marching_cubes.py @@ -11,7 +11,7 @@ from glue_ar.common.stl_builder import STLBuilder from glue_ar.common.usd_builder import USDBuilder from glue_ar.common.volume_export_options import ARIsosurfaceExportOptions -from glue_ar.gltf_utils import SHORT_MAX, add_points_to_bytearray, add_triangles_to_bytearray, \ +from glue_ar.gltf_utils import add_points_to_bytearray, add_triangles_to_bytearray, index_export_option, \ index_mins, index_maxes from glue_ar.utils import BoundsWithResolution, clip_sides, frb_for_layer, hex_to_components, isomin_for_layer, \ isomax_for_layer, layer_color @@ -61,8 +61,8 @@ def add_isosurface_layer_gltf(builder: GLTFBuilder, max_tri_index = int(max(idx for tri in triangles for idx in tri)) tri_maxes = [max_tri_index] - use_short = max_tri_index <= SHORT_MAX - add_triangles_to_bytearray(barr, triangles, short=use_short) + index_format = index_export_option(max_tri_index) + add_triangles_to_bytearray(barr, triangles, export_option=index_format) triangle_len = len(barr) - point_len builder.add_buffer(byte_length=len(barr), uri=level_bin) @@ -88,10 +88,9 @@ def add_isosurface_layer_gltf(builder: GLTFBuilder, byte_offset=point_len, 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, + component_type=index_format.component_type, count=len(triangles)*3, type=AccessorType.SCALAR, mins=tri_mins, diff --git a/glue_ar/common/scatter_gltf.py b/glue_ar/common/scatter_gltf.py index 65bfaa7..361f724 100644 --- a/glue_ar/common/scatter_gltf.py +++ b/glue_ar/common/scatter_gltf.py @@ -13,7 +13,7 @@ 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 SHORT_MAX, add_points_to_bytearray, add_triangles_to_bytearray, \ +from glue_ar.gltf_utils import add_points_to_bytearray, add_triangles_to_bytearray, index_export_option, \ index_mins, index_maxes from glue_ar.utils import Viewer3DState, iterable_has_nan, hex_to_components, \ layer_color, offset_triangles, unique_id, xyz_bounds, xyz_for_layer, Bounds @@ -54,10 +54,10 @@ def add_vectors_gltf(builder: GLTFBuilder, max_index = max(idx for tri in triangles for idx in tri) add_triangles_to_bytearray(barr, triangles) - use_short = max_index <= SHORT_MAX + index_format = index_export_option(max_index) if layer_state.vector_arrowhead: tip_triangles = cone_triangles(theta_resolution=tip_resolution, start_index=max_index + 1) - add_triangles_to_bytearray(barr, tip_triangles, short=use_short) + add_triangles_to_bytearray(barr, tip_triangles, export_option=index_format) max_index = max(idx for tri in tip_triangles for idx in tri) triangle_count += len(tip_triangles) @@ -70,10 +70,9 @@ def add_vectors_gltf(builder: GLTFBuilder, byte_offset=0, 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, + component_type=index_format.component_type, count=triangle_count*3, type=AccessorType.SCALAR, mins=[0], @@ -304,9 +303,9 @@ def add_scatter_layer_gltf(builder: GLTFBuilder, 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 + index_format = index_export_option(max_triangle_index) triangles_start = len(barr) - add_triangles_to_bytearray(barr, mesh_triangles, short=use_short) + add_triangles_to_bytearray(barr, mesh_triangles, export_option=index_format) triangles_len = len(barr) builder.add_buffer_view( buffer=buffer, @@ -314,10 +313,9 @@ def add_scatter_layer_gltf(builder: GLTFBuilder, 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, + component_type=index_format.component_type, count=len(mesh_triangles)*3, type=AccessorType.SCALAR, mins=[0], @@ -367,10 +365,9 @@ def add_scatter_layer_gltf(builder: GLTFBuilder, 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, + component_type=index_format.component_type, count=len(triangles)*3*count, type=AccessorType.SCALAR, mins=[0], @@ -422,9 +419,9 @@ def add_scatter_layer_gltf(builder: GLTFBuilder, 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 + index_format = index_export_option(max_triangle_index) triangles_start = len(barr) - add_triangles_to_bytearray(barr, mesh_triangles, short=use_short) + add_triangles_to_bytearray(barr, mesh_triangles, export_option=index_format) triangles_len = len(barr) builder.add_buffer_view( @@ -433,10 +430,9 @@ def add_scatter_layer_gltf(builder: GLTFBuilder, 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, + component_type=index_format.component_type, count=len(mesh_triangles)*3, type=AccessorType.SCALAR, mins=[0], @@ -486,10 +482,9 @@ def add_scatter_layer_gltf(builder: GLTFBuilder, 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, + component_type=index_format.component_type, count=len(triangles)*3*count, type=AccessorType.SCALAR, mins=[0], diff --git a/glue_ar/common/tests/gltf_helpers.py b/glue_ar/common/tests/gltf_helpers.py index 4903a12..994850e 100644 --- a/glue_ar/common/tests/gltf_helpers.py +++ b/glue_ar/common/tests/gltf_helpers.py @@ -4,12 +4,13 @@ from numbers import Number from struct import iter_unpack -from typing import List, Literal, Optional, Tuple, Union, cast +from typing import List, Literal, Optional, Tuple, cast +from glue_ar.gltf_utils import GLTFIndexExportOption from glue_ar.utils import iterator_count -BufferFormat = Union[Literal["f"], Literal["I"], Literal["H"]] +BufferFormat = Literal["f", "B", "H", "I"] def get_data(gltf: GLTF, buffer: Buffer, buffer_view: Optional[BufferView] = None) -> bytes: @@ -41,9 +42,8 @@ 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, use_short=False): - format = 'H' if use_short else 'I' - return count_points(gltf, buffer, buffer_view, format) +def count_indices(gltf: GLTF, buffer: Buffer, buffer_view: BufferView, export_option: GLTFIndexExportOption): + return count_points(gltf, buffer, buffer_view, export_option.format) def unpack_points(gltf: GLTF, @@ -71,5 +71,9 @@ def unpack_vertices(gltf: GLTF, buffer: Buffer, buffer_view: BufferView) -> List return unpack_points(gltf, buffer, buffer_view, 'f') -def unpack_indices(gltf: GLTF, buffer: Buffer, buffer_view: BufferView) -> List[Tuple[Number, Number, Number]]: - return unpack_points(gltf, buffer, buffer_view, 'I') +def unpack_indices(gltf: GLTF, + buffer: Buffer, + buffer_view: BufferView, + export_option: GLTFIndexExportOption = GLTFIndexExportOption.Int + ) -> List[Tuple[Number, Number, Number]]: + return unpack_points(gltf, buffer, buffer_view, export_option.format) diff --git a/glue_ar/common/tests/test_scatter_gltf.py b/glue_ar/common/tests/test_scatter_gltf.py index d6dc4b8..d1f8113 100644 --- a/glue_ar/common/tests/test_scatter_gltf.py +++ b/glue_ar/common/tests/test_scatter_gltf.py @@ -10,7 +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.gltf_utils import index_export_option from glue_ar.utils import export_label_for_layer, hex_to_components, layers_to_export, mask_for_bounds, \ xyz_bounds, xyz_for_layer @@ -69,8 +69,11 @@ 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) - use_short = points_count <= SHORT_MAX - assert count_indices(gltf, model.buffers[0], model.bufferViews[0], use_short=use_short) == triangles_count + index_format = index_export_option(points_count) + assert count_indices(gltf, + model.buffers[0], + model.bufferViews[0], + export_option=index_format) == triangles_count assert count_vertices(gltf, model.buffers[0], model.bufferViews[1]) == points_count assert model.bufferViews[0].target == BufferTarget.ELEMENT_ARRAY_BUFFER.value @@ -78,7 +81,7 @@ def test_basic_export(self, app_type: str, viewer_type: str): indices_accessor = model.accessors[0] assert indices_accessor.bufferView == 0 - expected_indices_type = ComponentType.UNSIGNED_SHORT if use_short else ComponentType.UNSIGNED_INT + expected_indices_type = index_format.component_type assert indices_accessor.componentType == expected_indices_type.value assert indices_accessor.count == triangles_count * 3 assert indices_accessor.type == AccessorType.SCALAR.value diff --git a/glue_ar/common/voxels.py b/glue_ar/common/voxels.py index 2d57e3c..7a88c7c 100644 --- a/glue_ar/common/voxels.py +++ b/glue_ar/common/voxels.py @@ -15,7 +15,7 @@ clip_sides, frb_for_layer, hex_to_components, isomin_for_layer, \ isomax_for_layer, layer_color, offset_triangles, unique_id, xyz_bounds -from glue_ar.gltf_utils import SHORT_MAX, add_points_to_bytearray, add_triangles_to_bytearray, \ +from glue_ar.gltf_utils import add_points_to_bytearray, add_triangles_to_bytearray, index_export_option, \ index_mins, index_maxes from glue_ar.common.shapes import rectangular_prism_points, rectangular_prism_triangulation @@ -112,10 +112,10 @@ def add_voxel_layers_gltf(builder: GLTFBuilder, triangles_count = len(tris) mesh_triangles = [tri for box in tris for tri in box] max_triangle_index = max(idx for tri in mesh_triangles for idx in tri) - use_short = max_triangle_index <= SHORT_MAX + index_format = index_export_option(max_triangle_index) triangles_barr = bytearray() - add_triangles_to_bytearray(triangles_barr, mesh_triangles, short=use_short) + add_triangles_to_bytearray(triangles_barr, mesh_triangles, export_option=index_format) triangles_len = len(triangles_barr) builder.add_buffer(byte_length=len(triangles_barr), uri=triangles_bin) @@ -130,10 +130,9 @@ def add_voxel_layers_gltf(builder: GLTFBuilder, 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, + component_type=index_format.component_type, count=len(mesh_triangles)*3, type=AccessorType.SCALAR, mins=[0], @@ -195,10 +194,9 @@ def add_voxel_layers_gltf(builder: GLTFBuilder, 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, + component_type=index_format.component_type, count=len(last_mesh_triangles)*3, type=AccessorType.SCALAR, mins=[0], diff --git a/glue_ar/gltf_utils.py b/glue_ar/gltf_utils.py index ceab78c..2dc5133 100644 --- a/glue_ar/gltf_utils.py +++ b/glue_ar/gltf_utils.py @@ -1,10 +1,13 @@ +from enum import Enum import operator import struct -from typing import Callable, Iterable, List, Optional, Type, TypeVar, Union +from typing import Callable, Iterable, List, Literal, Optional, Type, TypeVar, Union -from gltflib import Material, PBRMetallicRoughness +from gltflib import ComponentType, Material, PBRMetallicRoughness __all__ = [ + "GLTFIndexExportOption", + "index_export_option", "create_material_for_color", "add_points_to_bytearray", "add_triangles_to_bytearray", @@ -13,12 +16,28 @@ ] +class GLTFIndexExportOption(Enum): + Byte = ("B", ComponentType.UNSIGNED_BYTE, 2**8-1) + Short = ("H", ComponentType.UNSIGNED_SHORT, 2**16-1) + Int = ("I", ComponentType.UNSIGNED_INT, 2**32-1) + + def __init__(self, format: Literal["B", "H", "I"], component_type: ComponentType, max: int): + self.format = format + self.component_type = component_type + self.max = max + + GLTF_COMPRESSION_EXTENSIONS = { "draco": "KHR_draco_mesh_compression", "meshoptimizer": "EXT_meshopt_compression", } -SHORT_MAX = 65_535 + +def index_export_option(max_index: int) -> GLTFIndexExportOption: + for option in GLTFIndexExportOption: + if max_index <= option.max: + return option + return GLTFIndexExportOption.Int def create_material_for_color( @@ -44,11 +63,10 @@ def add_points_to_bytearray(arr: bytearray, points: Iterable[Iterable[Union[int, def add_triangles_to_bytearray(arr: bytearray, triangles: Iterable[Iterable[int]], - short: bool = False): - format = "H" if short else "I" + export_option: GLTFIndexExportOption = GLTFIndexExportOption.Int): for triangle in triangles: for index in triangle: - arr.extend(struct.pack(format, index)) + arr.extend(struct.pack(export_option.format, index)) T = TypeVar("T", bound=Union[int, float]) diff --git a/glue_ar/tests/test_gltf_utils.py b/glue_ar/tests/test_gltf_utils.py new file mode 100644 index 0000000..cbd0761 --- /dev/null +++ b/glue_ar/tests/test_gltf_utils.py @@ -0,0 +1,14 @@ +from ..gltf_utils import GLTFIndexExportOption, index_export_option + + +def test_index_export_option(): + assert index_export_option(3) == GLTFIndexExportOption.Byte + assert index_export_option(100) == GLTFIndexExportOption.Byte + assert index_export_option(255) == GLTFIndexExportOption.Byte + assert index_export_option(256) == GLTFIndexExportOption.Short + assert index_export_option(1234) == GLTFIndexExportOption.Short + assert index_export_option(10_000) == GLTFIndexExportOption.Short + assert index_export_option(65_535) == GLTFIndexExportOption.Short + assert index_export_option(65_536) == GLTFIndexExportOption.Int + assert index_export_option(100_000) == GLTFIndexExportOption.Int + assert index_export_option(1_000_000) == GLTFIndexExportOption.Int