Add vertex color to UV-textured GLB conversion in inference.py

Implemented a new function to convert .obj files with vertex colors to UV-mapped textured .glb files. This includes UV parameterization using xatlas, texture baking with inpainting and filtering, and assigning the generated texture to the mesh. Updated the generate3d function to utilize this new conversion method, enhancing texture quality in mesh exports.

inference.py

@@ -8,6 +8,88 @@ from mesh import Mesh
 import zipfile
 from util.renderer import Renderer
 import trimesh
+import xatlas
+import cv2
+from PIL import Image, ImageFilter
+
+def vertex_color_to_uv_textured_glb(obj_path, glb_path, texture_size=1024):
+    mesh = trimesh.load(obj_path, process=False)
+    vertex_colors = mesh.visual.vertex_colors[:, :3]  # (N, 3), uint8
+    # Generate UVs
+    vmapping, indices, uvs = xatlas.parametrize(mesh.vertices, mesh.faces)
+    vertices = mesh.vertices[vmapping]
+    vertex_colors = vertex_colors[vmapping]
+    mesh.vertices = vertices
+    mesh.faces = indices
+    # Bake texture
+    buffer_size = texture_size * 2
+    texture_buffer = np.zeros((buffer_size, buffer_size, 4), dtype=np.uint8)
+    def barycentric_interpolate(v0, v1, v2, c0, c1, c2, p):
+        v0v1 = v1 - v0
+        v0v2 = v2 - v0
+        v0p = p - v0
+        d00 = np.dot(v0v1, v0v1)
+        d01 = np.dot(v0v1, v0v2)
+        d11 = np.dot(v0v2, v0v2)
+        d20 = np.dot(v0p, v0v1)
+        d21 = np.dot(v0p, v0v2)
+        denom = d00 * d11 - d01 * d01
+        if abs(denom) < 1e-8:
+            return (c0 + c1 + c2) / 3
+        v = (d11 * d20 - d01 * d21) / denom
+        w = (d00 * d21 - d01 * d20) / denom
+        u = 1.0 - v - w
+        u = np.clip(u, 0, 1)
+        v = np.clip(v, 0, 1)
+        w = np.clip(w, 0, 1)
+        return u * c0 + v * c1 + w * c2
+    def is_point_in_triangle(p, v0, v1, v2):
+        def sign(p1, p2, p3):
+            return (p1[0] - p3[0]) * (p2[1] - p3[1]) - (p2[0] - p3[0]) * (p1[1] - p3[1])
+        d1 = sign(p, v0, v1)
+        d2 = sign(p, v1, v2)
+        d3 = sign(p, v2, v0)
+        has_neg = (d1 < 0) or (d2 < 0) or (d3 < 0)
+        has_pos = (d1 > 0) or (d2 > 0) or (d3 > 0)
+        return not (has_neg and has_pos)
+    for face in mesh.faces:
+        uv0, uv1, uv2 = uvs[face]
+        c0, c1, c2 = vertex_colors[face]
+        uv0 = (uv0 * (buffer_size - 1)).astype(int)
+        uv1 = (uv1 * (buffer_size - 1)).astype(int)
+        uv2 = (uv2 * (buffer_size - 1)).astype(int)
+        min_x = max(int(np.floor(min(uv0[0], uv1[0], uv2[0]))), 0)
+        max_x = min(int(np.ceil(max(uv0[0], uv1[0], uv2[0]))), buffer_size - 1)
+        min_y = max(int(np.floor(min(uv0[1], uv1[1], uv2[1]))), 0)
+        max_y = min(int(np.ceil(max(uv0[1], uv1[1], uv2[1]))), buffer_size - 1)
+        for y in range(min_y, max_y + 1):
+            for x in range(min_x, max_x + 1):
+                p = np.array([x + 0.5, y + 0.5])
+                if is_point_in_triangle(p, uv0, uv1, uv2):
+                    color = barycentric_interpolate(uv0, uv1, uv2, c0, c1, c2, p)
+                    texture_buffer[y, x, :3] = np.clip(color, 0, 255).astype(np.uint8)
+                    texture_buffer[y, x, 3] = 255
+    # Inpainting, filtering, and downsampling
+    image_bgra = texture_buffer.copy()
+    mask = (image_bgra[:, :, 3] == 0).astype(np.uint8) * 255
+    image_bgr = cv2.cvtColor(image_bgra, cv2.COLOR_BGRA2BGR)
+    inpainted_bgr = cv2.inpaint(image_bgr, mask, inpaintRadius=3, flags=cv2.INPAINT_TELEA)
+    inpainted_bgra = cv2.cvtColor(inpainted_bgr, cv2.COLOR_BGR2BGRA)
+    texture_buffer = inpainted_bgra[::-1]
+    image_texture = Image.fromarray(texture_buffer)
+    image_texture = image_texture.filter(ImageFilter.MedianFilter(size=3))
+    image_texture = image_texture.filter(ImageFilter.GaussianBlur(radius=1))
+    image_texture = image_texture.resize((texture_size, texture_size), Image.LANCZOS)
+    # Assign UVs and texture to mesh
+    material = trimesh.visual.material.PBRMaterial(
+        baseColorFactor=[1.0, 1.0, 1.0, 1.0],
+        baseColorTexture=image_texture,
+        metallicFactor=0.0,
+        roughnessFactor=1.0,
+    )
+    visuals = trimesh.visual.TextureVisuals(uv=uvs, material=material)
+    mesh.visual = visuals
+    mesh.export(glb_path)
 
 def generate3d(model, rgb, ccm, device):
 
@@ -82,9 +164,8 @@ def generate3d(model, rgb, ccm, device):
     elapsed_time = end_time - start_time
     print(f"uv takes {elapsed_time}s")
 
-    # Convert .obj (with vertex colors) to .glb
+    # Convert .obj (with vertex colors) to UV-mapped textured .glb
     obj_path = mesh_path_glb + ".obj"
     glb_path = mesh_path_glb + ".glb"
-    mesh = trimesh.load(obj_path, process=False)
-    mesh.export(glb_path)
+    vertex_color_to_uv_textured_glb(obj_path, glb_path)
     return glb_path