Refactor GLTF export in Mesh class to streamline buffer views and accessors. Enhance texture handling and ensure proper integration of vertex attributes, improving overall GLB file generation.

mesh.py

@@ -595,13 +595,12 @@ class Mesh:
         Args:
             path (str): path to write.
         """
+        import pygltflib
 
         # assert self.v.shape[0] == self.vn.shape[0] and self.v.shape[0] == self.vt.shape[0]
         if self.vt is not None and self.v.shape[0] != self.vt.shape[0]:
             self.align_v_to_vt()
 
-        import pygltflib
-
         f_np = self.f.detach().cpu().numpy().astype(np.uint32)
         f_np_blob = f_np.flatten().tobytes()
 
@@ -615,13 +614,53 @@ class Mesh:
         attributes = pygltflib.Attributes(POSITION=1)
         accessor_count = 2  # Start after position (0) and indices (1)
 
+        # Initialize buffer views list
+        buffer_views = [
+            # triangles; as flatten (element) array
+            pygltflib.BufferView(
+                buffer=0,
+                byteLength=len(f_np_blob),
+                target=pygltflib.ELEMENT_ARRAY_BUFFER,
+            ),
+            # positions; as vec3 array
+            pygltflib.BufferView(
+                buffer=0,
+                byteOffset=len(f_np_blob),
+                byteLength=len(v_np_blob),
+                byteStride=12,  # vec3
+                target=pygltflib.ARRAY_BUFFER,
+            ),
+        ]
+
+        # Initialize accessors list
+        accessors = [
+            # 0 = triangles
+            pygltflib.Accessor(
+                bufferView=0,
+                componentType=pygltflib.UNSIGNED_INT,
+                count=f_np.size,
+                type=pygltflib.SCALAR,
+                max=[int(f_np.max())],
+                min=[int(f_np.min())],
+            ),
+            # 1 = positions
+            pygltflib.Accessor(
+                bufferView=1,
+                componentType=pygltflib.FLOAT,
+                count=len(v_np),
+                type=pygltflib.VEC3,
+                max=v_np.max(axis=0).tolist(),
+                min=v_np.min(axis=0).tolist(),
+            ),
+        ]
+
         # Add vertex colors if they exist
         if self.vc is not None:
             vc_np = self.vc.detach().cpu().numpy().astype(np.float32)
             vc_np_blob = vc_np.tobytes()
             
             # Add vertex color buffer view
-            gltf.bufferViews.append(
+            buffer_views.append(
                 pygltflib.BufferView(
                     buffer=0,
                     byteOffset=byteOffset,
@@ -632,7 +671,7 @@ class Mesh:
             )
             
             # Add vertex color accessor
-            gltf.accessors.append(
+            accessors.append(
                 pygltflib.Accessor(
                     bufferView=accessor_count,
                     componentType=pygltflib.FLOAT,
@@ -650,7 +689,7 @@ class Mesh:
             blob += vc_np_blob
             byteOffset += len(vc_np_blob)
 
-        # base mesh
+        # Create the GLTF object with all components
         gltf = pygltflib.GLTF2(
             scene=0,
             scenes=[pygltflib.Scene(nodes=[0])],
@@ -659,46 +698,9 @@ class Mesh:
                 attributes=attributes,
                 indices=0,
             )])],
-            buffers=[
-                pygltflib.Buffer(byteLength=byteOffset)
-            ],
-            # buffer view (based on dtype)
-            bufferViews=[
-                # triangles; as flatten (element) array
-                pygltflib.BufferView(
-                    buffer=0,
-                    byteLength=len(f_np_blob),
-                    target=pygltflib.ELEMENT_ARRAY_BUFFER,
-                ),
-                # positions; as vec3 array
-                pygltflib.BufferView(
-                    buffer=0,
-                    byteOffset=len(f_np_blob),
-                    byteLength=len(v_np_blob),
-                    byteStride=12,  # vec3
-                    target=pygltflib.ARRAY_BUFFER,
-                ),
-            ],
-            accessors=[
-                # 0 = triangles
-                pygltflib.Accessor(
-                    bufferView=0,
-                    componentType=pygltflib.UNSIGNED_INT,
-                    count=f_np.size,
-                    type=pygltflib.SCALAR,
-                    max=[int(f_np.max())],
-                    min=[int(f_np.min())],
-                ),
-                # 1 = positions
-                pygltflib.Accessor(
-                    bufferView=1,
-                    componentType=pygltflib.FLOAT,
-                    count=len(v_np),
-                    type=pygltflib.VEC3,
-                    max=v_np.max(axis=0).tolist(),
-                    min=v_np.min(axis=0).tolist(),
-                ),
-            ],
+            buffers=[pygltflib.Buffer(byteLength=byteOffset)],
+            bufferViews=buffer_views,
+            accessors=accessors,
         )
 
         # Add material for vertex colors
@@ -715,53 +717,26 @@ class Mesh:
             ))
             gltf.meshes[0].primitives[0].material = 0
 
-        # append texture info
+        # Handle textures if they exist
         if self.vt is not None:
-
             vt_np = self.vt.detach().cpu().numpy().astype(np.float32)
             vt_np_blob = vt_np.tobytes()
 
-            albedo = self.albedo.detach().cpu().numpy()
-            albedo = (albedo * 255).astype(np.uint8)
-            albedo = cv2.cvtColor(albedo, cv2.COLOR_RGB2BGR)
-            albedo_blob = cv2.imencode('.png', albedo)[1].tobytes()
-
-            # update primitive
-            gltf.meshes[0].primitives[0].attributes.TEXCOORD_0 = 2
-            gltf.meshes[0].primitives[0].material = 0
-
-            # update materials
-            gltf.materials.append(pygltflib.Material(
-                pbrMetallicRoughness=pygltflib.PbrMetallicRoughness(
-                    baseColorTexture=pygltflib.TextureInfo(index=0, texCoord=0),
-                    metallicFactor=0.0,
-                    roughnessFactor=1.0,
-                ),
-                alphaMode=pygltflib.OPAQUE,
-                alphaCutoff=None,
-                doubleSided=True,
-            ))
-
-            gltf.textures.append(pygltflib.Texture(sampler=0, source=0))
-            gltf.samplers.append(pygltflib.Sampler(magFilter=pygltflib.LINEAR, minFilter=pygltflib.LINEAR_MIPMAP_LINEAR, wrapS=pygltflib.REPEAT, wrapT=pygltflib.REPEAT))
-            gltf.images.append(pygltflib.Image(bufferView=3, mimeType="image/png"))
-
-            # update buffers
+            # Add texture coordinates buffer view
             gltf.bufferViews.append(
-                # index = 2, texcoords; as vec2 array
                 pygltflib.BufferView(
                     buffer=0,
                     byteOffset=byteOffset,
                     byteLength=len(vt_np_blob),
-                    byteStride=8, # vec2
+                    byteStride=8,  # vec2
                     target=pygltflib.ARRAY_BUFFER,
                 )
             )
 
+            # Add texture coordinates accessor
             gltf.accessors.append(
-                # 2 = texcoords
                 pygltflib.Accessor(
-                    bufferView=2,
+                    bufferView=len(gltf.bufferViews) - 1,
                     componentType=pygltflib.FLOAT,
                     count=len(vt_np),
                     type=pygltflib.VEC2,
@@ -770,59 +745,19 @@ class Mesh:
                 )
             )
 
-            blob += vt_np_blob 
-            byteOffset += len(vt_np_blob)
-
-            gltf.bufferViews.append(
-                # index = 3, albedo texture; as none target
-                pygltflib.BufferView(
-                    buffer=0,
-                    byteOffset=byteOffset,
-                    byteLength=len(albedo_blob),
-                )
-            )
+            # Add texture coordinates to attributes
+            gltf.meshes[0].primitives[0].attributes.TEXCOORD_0 = len(gltf.accessors) - 1
 
-            blob += albedo_blob
-            byteOffset += len(albedo_blob)
+            blob += vt_np_blob
+            byteOffset += len(vt_np_blob)
 
+            # Update buffer size
             gltf.buffers[0].byteLength = byteOffset
 
-            # append metllic roughness
-            if self.metallicRoughness is not None:
-                metallicRoughness = self.metallicRoughness.detach().cpu().numpy()
-                metallicRoughness = (metallicRoughness * 255).astype(np.uint8)
-                metallicRoughness = cv2.cvtColor(metallicRoughness, cv2.COLOR_RGB2BGR)
-                metallicRoughness_blob = cv2.imencode('.png', metallicRoughness)[1].tobytes()
-
-                # update texture definition
-                gltf.materials[0].pbrMetallicRoughness.metallicFactor = 1.0
-                gltf.materials[0].pbrMetallicRoughness.roughnessFactor = 1.0
-                gltf.materials[0].pbrMetallicRoughness.metallicRoughnessTexture = pygltflib.TextureInfo(index=1, texCoord=0)
-
-                gltf.textures.append(pygltflib.Texture(sampler=1, source=1))
-                gltf.samplers.append(pygltflib.Sampler(magFilter=pygltflib.LINEAR, minFilter=pygltflib.LINEAR_MIPMAP_LINEAR, wrapS=pygltflib.REPEAT, wrapT=pygltflib.REPEAT))
-                gltf.images.append(pygltflib.Image(bufferView=4, mimeType="image/png"))
-
-                # update buffers
-                gltf.bufferViews.append(
-                    # index = 4, metallicRoughness texture; as none target
-                    pygltflib.BufferView(
-                        buffer=0,
-                        byteOffset=byteOffset,
-                        byteLength=len(metallicRoughness_blob),
-                    )
-                )
-
-                blob += metallicRoughness_blob
-                byteOffset += len(metallicRoughness_blob)
-
-                gltf.buffers[0].byteLength = byteOffset
-
-            
-        # set actual data
+        # Set the binary blob
         gltf.set_binary_blob(blob)
 
-        # glb = b"".join(gltf.save_to_bytes())
+        # Save the GLB file
         gltf.save(path)