collada/buffer: correctly handle <mesh> with more than one <triangles>

collada/buffer.py

@@ -1,5 +1,8 @@
 from collections import defaultdict
 from itertools import chain, islice
+from typing import Dict, Tuple, List
+from dataclasses import dataclass
+import dataclasses
 
 from collada import parse
 from collada import types
@@ -8,7 +11,7 @@ from collada.util import find_semantics
 def linearize_offset_table(by_offset, p_stride):
     for offset in range(p_stride):
         for input, source in by_offset[offset]:
-            yield offset, input, source
+            yield input, source
 
 mesh_semantic_names = ["NORMAL", "TEXCOORD"]
 
@@ -34,65 +37,85 @@ def build_offset_table(collada, triangles, p_stride):
     linearized_table = linearize_offset_table(by_offset, p_stride)
     return list(linearized_table), used_offsets
 
-def mesh_vertex_index_buffer(collada, mesh):
+def mesh_vertex_index_buffer_triangles(collada, mesh, triangles):
-    assert len(mesh.primitive_elements) == 1
-    triangles, = mesh.primitive_elements
-    assert type(triangles) is types.Triangles
-
     max_offset = max(i.offset for i in triangles.inputs)
     p_stride = max_offset + 1
     offset_table, used_offsets = build_offset_table(collada, triangles, p_stride)
 
-    ######################################################################
+from prettyprinter import pprint, install_extras
-    # generate the index and vertex buffers
+install_extras(include=["dataclasses"])
-    ######################################################################
 
-    vertex_buffer_stride = sum(
+@dataclass
-        source.technique_common.accessor.stride
+class MeshVertexIndexBufferState:
-        for offset, input, source in offset_table
+    index_buffer: List[int]
-    )
+
-    vertex_index_table = []
+    def __init__(self):
+        self.index_buffer = [] # a list of integers
+        self.vertex_buffer = [] # a list of floats
+
+        # vertex_index_table: input/collada vertex indices in the order written
+        #                     to the index buffer
+        #   indices: output vertex index (absolute)
+        #   values: collada vertex index
+        self.vertex_index_table = []
+
+        # one offset table per <triangles>
+        self.offset_tables = []
+
+def mesh_vertex_index_buffer(collada, mesh):
+    state = MeshVertexIndexBufferState()
+
+    # index_table:
+    #   keys: index_table_key
+    #   values: index integer
     index_table = {}
     next_output_index = 0
-    index_buffer = []
-    vertex_buffer = []
 
-    for vertex_ix in range(triangles.count * 3):
+    for triangles in mesh.primitive_elements:
-        index_table_key = tuple(triangles.p[vertex_ix * p_stride + offset] for offset in used_offsets)
+        assert type(triangles) is types.Triangles
-        if index_table_key in index_table:
-            index_buffer.append(index_table[index_table_key])
-            continue
 
-        index_table[index_table_key] = next_output_index
+        max_offset = max(i.offset for i in triangles.inputs)
-        index_buffer.append(next_output_index)
+        p_stride = max_offset + 1
-        next_output_index += 1
+        offset_table, used_offsets = build_offset_table(collada, triangles, p_stride)
+        state.offset_tables.append(offset_table)
+        input_key = tuple(dataclasses.astuple(input) for input, source in offset_table)
 
-        # emit vertex attributes for new output index in vertex buffer
+        for vertex_ix in range(triangles.count * 3):
-        for offset, input, source in offset_table:
+            offset_key = tuple(triangles.p[vertex_ix * p_stride + offset] for offset in used_offsets)
-            p_index = triangles.p[vertex_ix * p_stride + offset]
+            index_table_key = tuple((input_key, offset_key))
-            if input.semantic == "VERTEX":
-                vertex_index_table.append(p_index)
 
-            source_stride = source.technique_common.accessor.stride
+            ######################################################################
-            source_index = p_index * source_stride
+            # append to the index buffer
-            array_slice = source.array_element.floats[source_index:source_index+source_stride]
+            ######################################################################
-            vertex_buffer.extend(array_slice)
+            if index_table_key in index_table:
+                state.index_buffer.append(index_table[index_table_key])
+                continue
+            index_table[index_table_key] = next_output_index
+            state.index_buffer.append(next_output_index)
+            next_output_index += 1
 
-    assert len(index_buffer) == triangles.count * 3
+            ######################################################################
-    assert len(vertex_buffer) == len(index_table) * vertex_buffer_stride
+            # emit vertex attributes for new output index in vertex buffer
+            ######################################################################
+            for input, source in offset_table:
+                p_index = triangles.p[vertex_ix * p_stride + input.offset]
+                if input.semantic == "VERTEX":
+                    state.vertex_index_table.append(p_index)
 
-    input_source_table = [(input, source) for offset, input, source in offset_table]
+                source_stride = source.technique_common.accessor.stride
+                source_index = p_index * source_stride
+                assert type(source.array_element) is types.FloatArray
+                array_slice = source.array_element.floats[source_index:source_index+source_stride]
+                state.vertex_buffer.extend(array_slice)
 
-    # vertex_index_table: input/collada vertex indices in the order written to the index buffer
+    return state
-    # input_source_table: (input, source) in the order written to the vertex buffer
-    return vertex_buffer, index_buffer, vertex_index_table, input_source_table
 
 def skin_vertex_buffer(collada, skin, vertex_index_table):
     max_offset = max(i.offset for i in skin.vertex_weights.inputs)
-    weights_input, = find_semantics(skin.vertex_weights.inputs, "WEIGHT")
-    weights_source = collada.lookup(weights_input.source, types.SourceCore)
     joints_input, = find_semantics(skin.vertex_weights.inputs, "JOINT")
     joints_source = collada.lookup(joints_input.source, types.SourceCore)
+    weights_input, = find_semantics(skin.vertex_weights.inputs, "WEIGHT")
+    weights_source = collada.lookup(weights_input.source, types.SourceCore)
     assert weights_source.technique_common.accessor.stride == 1
     assert joints_source.technique_common.accessor.stride == 1
 
@@ -136,8 +159,8 @@ if __name__ == "__main__":
 
     mesh = collada.library_geometries[0].geometries[0].geometric_element
     assert type(mesh) is types.Mesh
-    vertex_buffer_pnt, index_buffer, vertex_index_table, input_source_table = mesh_vertex_index_buffer(collada, mesh)
+    state = mesh_vertex_index_buffer(collada, mesh)
 
-    skin = collada.library_controllers[0].controllers[0].control_element
+    #skin = collada.library_controllers[0].controllers[0].control_element
-    assert type(skin) is types.Skin
+    #assert type(skin) is types.Skin
-    vertex_buffer_jw = skin_vertex_buffer(collada, skin, vertex_index_table)
+    #vertex_buffer_jw = skin_vertex_buffer(collada, skin, vertex_index_table)

collada/types.py

@@ -9,14 +9,14 @@ Float3 = Tuple[float, float, float]
 Float4 = Tuple[float, float, float, float]
 Float7 = Tuple[float, float, float, float, float, float, float]
 
-@dataclass
+@dataclass(frozen=True)
 class Lookat:
     sid: Optional[str]
     eye: Float3
     at: Float3
     up: Float3
 
-@dataclass
+@dataclass(frozen=True)
 class Matrix:
     # collada is the transpose of a directx matrix
     # 1.0  0.0  0.0  0.0
@@ -26,35 +26,35 @@ class Matrix:
     sid: Optional[str]
     values: Tuple[Float4, Float4, Float4, Float4]
 
-@dataclass
+@dataclass(frozen=True)
 class Rotate:
     # x y z w ; quaternion
     sid: Optional[str]
     rotate: Float4
 
-@dataclass
+@dataclass(frozen=True)
 class Scale:
     sid: Optional[str]
     scale: Float3
 
-@dataclass
+@dataclass(frozen=True)
 class Skew:
     sid: Optional[str]
     skew: Float7
 
-@dataclass
+@dataclass(frozen=True)
 class Translate:
     sid: Optional[str]
     translate: Float3
 
-@dataclass
+@dataclass(frozen=True)
 class Param:
     name: Optional[str]
     sid: Optional[str]
     type: str
     semantic: Optional[str]
 
-@dataclass
+@dataclass(frozen=True)
 class Accessor:
     count: int
     offset: Optional[int]
@@ -65,33 +65,33 @@ class Accessor:
 
     sid_lookup: dict = field(repr=False)
 
-@dataclass
+@dataclass(frozen=True)
 class TechniqueCommon_SourceCore:
     accessor: Accessor
 
-@dataclass
+@dataclass(frozen=True)
 class Ambient:
     color: Float3
 
-@dataclass
+@dataclass(frozen=True)
 class Directional:
     color: Float3
 
-@dataclass
+@dataclass(frozen=True)
 class Point:
     color: Float3
 
-@dataclass
+@dataclass(frozen=True)
 class Spot:
     color: Float3
 
-@dataclass
+@dataclass(frozen=True)
 class BindVertexInput:
     semantic: str
     input_semantic: str
     input_set: int
 
-@dataclass
+@dataclass(frozen=True)
 class InstanceMaterial:
     sid: Optional[str]
     name: Optional[str]
@@ -100,17 +100,17 @@ class InstanceMaterial:
 
     bind_vertex_inputs: List[BindVertexInput]
 
-@dataclass
+@dataclass(frozen=True)
 class TechniqueCommon_BindMaterial:
     materials: List[InstanceMaterial] # one or more
 
     sid_lookup: dict = field(repr=False)
 
-@dataclass
+@dataclass(frozen=True)
 class BindMaterial:
     technique_common: TechniqueCommon_BindMaterial
 
-@dataclass
+@dataclass(frozen=True)
 class InstanceGeometry:
     sid: Optional[str]
     name: Optional[str]
@@ -119,19 +119,19 @@ class InstanceGeometry:
     # child element
     bind_material: Optional[BindMaterial]
 
-@dataclass
+@dataclass(frozen=True)
 class InstanceLight:
     sid: Optional[str]
     name: Optional[str]
     url: URI
 
-@dataclass
+@dataclass(frozen=True)
 class InstanceVisualScene:
     sid: Optional[str]
     name: Optional[str]
     url: URI
 
-@dataclass
+@dataclass(frozen=True)
 class Scene:
     instance_visual_scene: InstanceVisualScene
 
@@ -146,43 +146,43 @@ class FxSurfaceType(Enum):
     DEPTH = auto()
     RECT = auto()
 
-@dataclass
+@dataclass(frozen=True)
 class InitFrom:
     uri: str
 
-@dataclass
+@dataclass(frozen=True)
 class Surface:
     type: FxSurfaceType
 
     init_from: InitFrom
 
-@dataclass
+@dataclass(frozen=True)
 class SourceFX:
     sid: ID
 
-@dataclass
+@dataclass(frozen=True)
 class Sampler2D:
     source: SourceFX
 
-@dataclass
+@dataclass(frozen=True)
 class Newparam:
     sid: str # required
     parameter_type: Union[Surface, Sampler2D]
 
-@dataclass
+@dataclass(frozen=True)
 class Color:
     value: Float4
 
-@dataclass
+@dataclass(frozen=True)
 class Float:
     value: float
 
-@dataclass
+@dataclass(frozen=True)
 class Texture:
     texture: str
     texcoord:  str
 
-@dataclass
+@dataclass(frozen=True)
 class Blinn:
     emission: Optional[Union[Color, Texture]]
     ambient: Optional[Union[Color, Texture]]
@@ -195,7 +195,7 @@ class Blinn:
     transparency: Optional[Float]
     index_of_refraction: Optional[Float]
 
-@dataclass
+@dataclass(frozen=True)
 class Lambert:
     emission: Optional[Union[Color, Texture]]
     ambient: Optional[Union[Color, Texture]]
@@ -206,7 +206,7 @@ class Lambert:
     transparency: Optional[Float]
     index_of_refraction: Optional[Float]
 
-@dataclass
+@dataclass(frozen=True)
 class Phong:
     emission: Optional[Union[Color, Texture]]
     ambient: Optional[Union[Color, Texture]]
@@ -219,7 +219,7 @@ class Phong:
     transparency: Optional[Float]
     index_of_refraction: Optional[Float]
 
-@dataclass
+@dataclass(frozen=True)
 class Constant:
     emission: Optional[Color]
     reflective: Optional[Union[Color, Texture]]
@@ -228,13 +228,13 @@ class Constant:
     transparency: Optional[Float]
     index_of_refraction: Optional[Float]
 
-@dataclass
+@dataclass(frozen=True)
 class TechniqueFX:
     id: Optional[ID]
     sid: str # required
     shader: Union[Blinn, Lambert, Phong, Constant]
 
-@dataclass
+@dataclass(frozen=True)
 class ProfileCommon:
     id: Optional[ID]
 
@@ -243,27 +243,27 @@ class ProfileCommon:
 
     sid_lookup: dict = field(repr=False)
 
-@dataclass
+@dataclass(frozen=True)
 class Effect:
     id: str
     name: Optional[str]
 
     profile_common: List[ProfileCommon]
 
-@dataclass
+@dataclass(frozen=True)
 class LibraryEffects:
     id: Optional[ID]
     name: Optional[str]
 
     effects: List[Effect]
 
-@dataclass
+@dataclass(frozen=True)
 class InstanceEffect:
     sid: Optional[str]
     name: Optional[str]
     url: URI
 
-@dataclass
+@dataclass(frozen=True)
 class Material:
     id: Optional[ID]
     name: Optional[str]
@@ -272,14 +272,14 @@ class Material:
 
     sid_lookup: dict = field(repr=False)
 
-@dataclass
+@dataclass(frozen=True)
 class LibraryMaterials:
     id: Optional[ID]
     name: Optional[str]
 
     materials: List[Material]
 
-@dataclass
+@dataclass(frozen=True)
 class NameArray:
     count: int
     id: Optional[ID]
@@ -287,7 +287,7 @@ class NameArray:
 
     names: List[str]
 
-@dataclass
+@dataclass(frozen=True)
 class BoolArray:
     count: int
     id: Optional[ID]
@@ -295,7 +295,7 @@ class BoolArray:
 
     bools: List[bool]
 
-@dataclass
+@dataclass(frozen=True)
 class FloatArray:
     count: int
     id: Optional[ID]
@@ -305,7 +305,7 @@ class FloatArray:
 
     floats: List[float]
 
-@dataclass
+@dataclass(frozen=True)
 class IntArray:
     count: int
     id: Optional[ID]
@@ -315,7 +315,7 @@ class IntArray:
 
     ints: List[int]
 
-@dataclass
+@dataclass(frozen=True)
 class SourceCore:
     id: ID
     name: Optional[str]
@@ -323,14 +323,14 @@ class SourceCore:
     array_element: Union[NameArray, BoolArray, FloatArray, IntArray]
     technique_common: TechniqueCommon_SourceCore
 
-@dataclass
+@dataclass(frozen=True)
 class InputShared:
     offset: int
     semantic: str
     source: URI
     set: Optional[int]
 
-@dataclass
+@dataclass(frozen=True)
 class Triangles:
     name: Optional[str]
     count: int
@@ -338,57 +338,57 @@ class Triangles:
     inputs: List[InputShared]
     p: List[int]
 
-@dataclass
+@dataclass(frozen=True)
 class InputUnshared:
     semantic: str
     source: URI
 
-@dataclass
+@dataclass(frozen=True)
 class Vertices:
     id: ID
     name: Optional[str]
 
     inputs: List[InputUnshared] # 1 or more
 
-@dataclass
+@dataclass(frozen=True)
 class Mesh:
     sources: List[SourceCore]
     vertices: Vertices
     primitive_elements: List[Union[Triangles]]
 
-@dataclass
+@dataclass(frozen=True)
 class Geometry:
     id: Optional[ID]
     name: Optional[str]
 
     geometric_element: Union[Mesh]
 
-@dataclass
+@dataclass(frozen=True)
 class LibraryGeometries:
     id: Optional[ID]
     name: Optional[str]
 
     geometries: List[Geometry]
 
-@dataclass
+@dataclass(frozen=True)
 class TechniqueCommon_Light:
     light: Union[Ambient, Directional, Point, Spot]
 
-@dataclass
+@dataclass(frozen=True)
 class Light:
     id: Optional[ID]
     name: Optional[str]
 
     technique_common: TechniqueCommon_Light
 
-@dataclass
+@dataclass(frozen=True)
 class LibraryLights:
     id: Optional[ID]
     name: Optional[str]
 
     lights: List[Light]
 
-@dataclass
+@dataclass(frozen=True)
 class Image:
     id: Optional[ID]
     name: Optional[str]
@@ -399,7 +399,7 @@ class Image:
 
     image_source: Union[InitFrom]
 
-@dataclass
+@dataclass(frozen=True)
 class LibraryImages:
     id: Optional[ID]
     name: Optional[str]
@@ -412,7 +412,7 @@ class NodeType(Enum):
 
 TransformationElements = Union[Lookat, Matrix, Rotate, Scale, Skew, Translate]
 
-@dataclass
+@dataclass(frozen=True)
 class Node:
     id: Optional[ID]
     name: Optional[str]
@@ -427,7 +427,7 @@ class Node:
 
     sid_lookup: dict = field(repr=False)
 
-@dataclass
+@dataclass(frozen=True)
 class VisualScene:
     id: Optional[ID]
     name: Optional[str]
@@ -436,30 +436,30 @@ class VisualScene:
 
     sid_lookup: dict = field(repr=False)
 
-@dataclass
+@dataclass(frozen=True)
 class LibraryVisualScenes:
     id: Optional[ID]
     name: Optional[str]
 
     visual_scenes: List[VisualScene]
 
-@dataclass
+@dataclass(frozen=True)
 class BindShapeMatrix:
     # it is written in row-major order in the COLLADA document for
     # human readability.
     values: Tuple[Float4, Float4, Float4, Float4]
 
-@dataclass
+@dataclass(frozen=True)
 class Joints:
     inputs: List[InputUnshared] # 2 or more
 
-@dataclass
+@dataclass(frozen=True)
 class VertexWeights:
     inputs: List[InputShared] # 2 or more
     vcount: List[int]
     v: List[int]
 
-@dataclass
+@dataclass(frozen=True)
 class Skin:
     source: URI # required
 
@@ -468,31 +468,31 @@ class Skin:
     joints: Joints # 1
     vertex_weights: VertexWeights # 1
 
-@dataclass
+@dataclass(frozen=True)
 class Controller:
     id: Optional[ID]
     name: Optional[str]
 
     control_element: Union[Skin]
 
-@dataclass
+@dataclass(frozen=True)
 class LibraryControllers:
     id: Optional[ID]
     name: Optional[str]
 
     controllers: List[Controller]
 
-@dataclass
+@dataclass(frozen=True)
 class Sampler:
     id: Optional[ID]
     inputs: List[InputUnshared] # 1 or more
 
-@dataclass
+@dataclass(frozen=True)
 class Channel:
     source: URI
     target: URI
 
-@dataclass
+@dataclass(frozen=True)
 class Animation:
     id: Optional[ID]
     name: Optional[str]
@@ -502,7 +502,7 @@ class Animation:
     samplers: List[Sampler]
     channels: List[Channel]
 
-@dataclass
+@dataclass(frozen=True)
 class LibraryAnimations:
     id: Optional[ID]
     name: Optional[str]