Add vertex color support in Mesh class and update GLB export in CRM model. Enhance mesh attributes and material handling for vertex colors, ensuring proper integration in GLTF format.

mesh.py

@@ -611,20 +611,56 @@ class Mesh:
         blob = f_np_blob + v_np_blob
         byteOffset = len(blob)
 
+        # Create attributes dictionary
+        attributes = pygltflib.Attributes(POSITION=1)
+        accessor_count = 2  # Start after position (0) and indices (1)
+
+        # 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(
+                pygltflib.BufferView(
+                    buffer=0,
+                    byteOffset=byteOffset,
+                    byteLength=len(vc_np_blob),
+                    byteStride=12,  # vec3
+                    target=pygltflib.ARRAY_BUFFER,
+                )
+            )
+            
+            # Add vertex color accessor
+            gltf.accessors.append(
+                pygltflib.Accessor(
+                    bufferView=accessor_count,
+                    componentType=pygltflib.FLOAT,
+                    count=len(vc_np),
+                    type=pygltflib.VEC3,
+                    max=vc_np.max(axis=0).tolist(),
+                    min=vc_np.min(axis=0).tolist(),
+                )
+            )
+            
+            # Add to attributes
+            attributes.COLOR_0 = accessor_count
+            accessor_count += 1
+            
+            blob += vc_np_blob
+            byteOffset += len(vc_np_blob)
+
         # base mesh
         gltf = pygltflib.GLTF2(
             scene=0,
             scenes=[pygltflib.Scene(nodes=[0])],
             nodes=[pygltflib.Node(mesh=0)],
             meshes=[pygltflib.Mesh(primitives=[pygltflib.Primitive(
-                # indices to accessors (0 is triangles)
-                attributes=pygltflib.Attributes(
-                    POSITION=1,
-                ),
+                attributes=attributes,
                 indices=0,
             )])],
             buffers=[
-                pygltflib.Buffer(byteLength=len(f_np_blob) + len(v_np_blob))
+                pygltflib.Buffer(byteLength=byteOffset)
             ],
             # buffer view (based on dtype)
             bufferViews=[
@@ -632,22 +668,22 @@ class Mesh:
                 pygltflib.BufferView(
                     buffer=0,
                     byteLength=len(f_np_blob),
-                    target=pygltflib.ELEMENT_ARRAY_BUFFER, # GL_ELEMENT_ARRAY_BUFFER (34963)
+                    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, # GL_ARRAY_BUFFER (34962)
+                    byteStride=12,  # vec3
+                    target=pygltflib.ARRAY_BUFFER,
                 ),
             ],
             accessors=[
                 # 0 = triangles
                 pygltflib.Accessor(
                     bufferView=0,
-                    componentType=pygltflib.UNSIGNED_INT, # GL_UNSIGNED_INT (5125)
+                    componentType=pygltflib.UNSIGNED_INT,
                     count=f_np.size,
                     type=pygltflib.SCALAR,
                     max=[int(f_np.max())],
@@ -656,7 +692,7 @@ class Mesh:
                 # 1 = positions
                 pygltflib.Accessor(
                     bufferView=1,
-                    componentType=pygltflib.FLOAT, # GL_FLOAT (5126)
+                    componentType=pygltflib.FLOAT,
                     count=len(v_np),
                     type=pygltflib.VEC3,
                     max=v_np.max(axis=0).tolist(),
@@ -665,6 +701,20 @@ class Mesh:
             ],
         )
 
+        # Add material for vertex colors
+        if self.vc is not None:
+            gltf.materials.append(pygltflib.Material(
+                pbrMetallicRoughness=pygltflib.PbrMetallicRoughness(
+                    baseColorFactor=[1.0, 1.0, 1.0, 1.0],
+                    metallicFactor=0.0,
+                    roughnessFactor=1.0,
+                ),
+                alphaMode=pygltflib.OPAQUE,
+                alphaCutoff=None,
+                doubleSided=True,
+            ))
+            gltf.meshes[0].primitives[0].material = 0
+
         # append texture info
         if self.vt is not None:
 

model/crm/model.py

@@ -102,13 +102,21 @@ class CRM(nn.Module):
         verts[..., 0]*= -1
         verts[..., 2]*= -1
         verts = verts.squeeze().cpu().numpy()
-        faces = faces[..., [2, 1, 0]][..., [0, 2, 1]]#[..., [1, 0, 2]]
-        faces = faces.squeeze().cpu().numpy()#faces[..., [2, 1, 0]].squeeze().cpu().numpy()
+        faces = faces[..., [2, 1, 0]][..., [0, 2, 1]]
+        faces = faces.squeeze().cpu().numpy()
 
         # export the final mesh
         with torch.no_grad():
-            mesh = trimesh.Trimesh(verts, faces, vertex_colors=colors, process=False) # important, process=True leads to seg fault...
-            mesh.export(f'{out_dir}.obj')
+            # Create a Mesh object with the data
+            from mesh import Mesh
+            mesh = Mesh(
+                v=torch.from_numpy(verts).float(),
+                f=torch.from_numpy(faces).int(),
+                vc=torch.from_numpy(colors).float(),
+                device='cpu'
+            )
+            # Write as GLB
+            mesh.write(out_dir)
 
     def export_mesh_wt_uv(self, ctx, data, out_dir, ind, device, res, tri_fea_2=None):