Refactor texture baking logic in generate3d function of inference.py to utilize per-vertex colors for improved texture quality. Updated the method for filling the texture image to enhance the visual output in mesh exports.

inference.py

@@ -68,29 +68,19 @@ def generate3d(model, rgb, ccm, device):
     mesh_f = data_config['faces'].cpu().numpy()
     vmapping, ft, vt = xatlas.parametrize(mesh_v, mesh_f)
 
-    # Bake texture using model's decoder and rgbMlp (CPU-only)
+    # Use per-vertex colors if available, else fallback to white
+    vertex_colors = np.ones((mesh_v.shape[0], 3), dtype=np.float32)  # fallback: white
+    # If you have per-vertex color, you can assign here, e.g.:
+    # vertex_colors = ...
+
+    # Bake vertex colors to texture in UV space
     tex_res = (1024, 1024)
-    # Generate a grid of UV coordinates
-    uv_grid = np.stack(np.meshgrid(
-        np.linspace(0, 1, tex_res[0]),
-        np.linspace(0, 1, tex_res[1])
-    ), -1).reshape(-1, 2)  # (H*W, 2)
-    # Map UVs to 3D positions using barycentric interpolation of mesh faces
-    # For simplicity, we'll sample random points on the mesh surface and use their UVs
-    # (A more advanced approach would rasterize each face, but this is a CPU-friendly approximation)
-    mesh_trimesh = trimesh.Trimesh(vertices=mesh_v, faces=mesh_f, process=False)
-    points, face_indices = trimesh.sample.sample_surface(mesh_trimesh, tex_res[0]*tex_res[1])
-    # Use the model's decoder and rgbMlp to get color for each sampled point
-    points_tensor = torch.from_numpy(points).float().unsqueeze(0).cpu()  # (1, N, 3) on CPU
-    triplane_feature2 = triplane_feature2.cpu()  # Ensure on CPU
-    with torch.no_grad():
-        dec_verts = model.decoder(triplane_feature2, points_tensor)
-        colors = model.rgbMlp(dec_verts).squeeze().cpu().numpy()  # (N, 3)
-        colors = (colors * 0.5 + 0.5).clip(0, 1)
-    # Fill the texture image
-    texture = np.zeros((tex_res[1]*tex_res[0], 3), dtype=np.float32)
-    texture[:colors.shape[0]] = colors
-    texture = texture.reshape(tex_res[1], tex_res[0], 3)
+    texture = np.zeros((tex_res[1], tex_res[0], 3), dtype=np.float32)
+    vt_img = (vt * np.array(tex_res)).astype(np.int32)
+    for face, uv_idx in zip(mesh_f, ft):
+        pts = vt_img[uv_idx]
+        color = vertex_colors[face].mean(axis=0)
+        cv2.fillPoly(texture, [pts], color.tolist())
     texture = np.clip(texture, 0, 1)
 
     # Create Mesh and export .glb