Merge pull request #112 from Carifio24/gltf-index-export-options

glTF index export options

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,

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,20 +303,19 @@ 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,
             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,
+            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],

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)

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,16 +69,19 @@ 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
         assert model.bufferViews[1].target == BufferTarget.ARRAY_BUFFER.value
 
         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

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],

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])

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