Add color assignment to new vertices in generate3d function

- Integrated cKDTree for efficient nearest neighbor search to assign colors to newly generated vertices.
- Created a new mesh using trimesh with updated vertex colors based on nearest old vertices.
- Enhanced overall mesh generation process by incorporating color information into the output.

inference.py

@@ -6,6 +6,8 @@ import tempfile
 from util.renderer import Renderer
 import os
 from PIL import Image
+import trimesh
+from scipy.spatial import cKDTree
 
 
 def generate3d(model, rgb, ccm, device):
@@ -77,4 +79,15 @@ def generate3d(model, rgb, ccm, device):
 
     model.export_mesh(data_config, base_path, tri_fea_2=triplane_feature2)  # writes .obj
 
+    # Build KDTree for nearest neighbor search
+    tree = cKDTree(data_config['verts'].squeeze().cpu().numpy())
+
+    # For each new vertex, find the nearest old vertex and copy its color
+    k = 3  # or 5, etc.
+    dists, idxs = tree.query(verts.cpu().numpy(), k=k)
+    new_colors = np.mean(color_tri[idxs], axis=1)
+
+    # Create the new mesh with colors
+    mesh = trimesh.Trimesh(vertices=verts.cpu().numpy(), faces=faces.cpu().numpy(), vertex_colors=new_colors)
+
     return obj_path