Refactor texture baking in vertex_color_to_uv_textured_glb function to use vectorized operations. This change improves performance by precomputing face data and applying vectorized barycentric tests, enhancing the efficiency of texture generation.

inference.py

@@ -21,54 +21,57 @@ def vertex_color_to_uv_textured_glb(obj_path, glb_path, texture_size=512):
     vertex_colors = vertex_colors[vmapping]
     mesh.vertices = vertices
     mesh.faces = indices
-    # Bake texture
+    # Bake texture (vectorized)
     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):
+    # Precompute face data
+    face_uvs = uvs[mesh.faces]
+    face_colors = vertex_colors[mesh.faces]
+    # Compute bounding boxes for all faces
+    min_xy = np.floor(np.min(face_uvs, axis=1) * (buffer_size - 1)).astype(int)
+    max_xy = np.ceil(np.max(face_uvs, axis=1) * (buffer_size - 1)).astype(int)
+    for i in range(len(mesh.faces)):
+        uv0, uv1, uv2 = face_uvs[i]
+        c0, c1, c2 = face_colors[i]
+        min_x, min_y = min_xy[i]
+        max_x, max_y = max_xy[i]
+        # Create a grid of pixel coordinates in the bounding box
+        xs = np.arange(min_x, max_x + 1)
+        ys = np.arange(min_y, max_y + 1)
+        xv, yv = np.meshgrid(xs, ys)
+        pts = np.stack([xv, yv], axis=-1).reshape(-1, 2) + 0.5
+        # Barycentric test (vectorized)
+        v0, v1, v2 = uv0 * (buffer_size - 1), uv1 * (buffer_size - 1), uv2 * (buffer_size - 1)
         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
+            return (p1[..., 0] - p3[0]) * (p2[1] - p3[1]) - (p2[0] - p3[0]) * (p1[..., 1] - p3[1])
+        d1 = sign(pts, v0, v1)
+        d2 = sign(pts, v1, v2)
+        d3 = sign(pts, v2, v0)
+        mask = ~((d1 < 0) | (d2 < 0) | (d3 < 0)) & ~((d1 > 0) & (d2 > 0) & (d3 > 0))
+        inside_pts = pts[mask]
+        if len(inside_pts) == 0:
+            continue
+        # Barycentric coordinates (vectorized)
+        def barycentric(p, v0, v1, v2):
+            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
+            v = (d11 * d20 - d01 * d21) / denom
+            w = (d00 * d21 - d01 * d20) / denom
+            u = 1.0 - v - w
+            return np.clip(u, 0, 1), np.clip(v, 0, 1), np.clip(w, 0, 1)
+        u, v, w = barycentric(inside_pts, v0, v1, v2)
+        colors = (u[:, None] * c0 + v[:, None] * c1 + w[:, None] * c2)
+        xi = inside_pts[:, 0].astype(int)
+        yi = inside_pts[:, 1].astype(int)
+        texture_buffer[yi, xi, :3] = np.clip(colors, 0, 255).astype(np.uint8)
+        texture_buffer[yi, xi, 3] = 255
     # Inpainting, filtering, and downsampling
     image_bgra = texture_buffer.copy()
     mask = (image_bgra[:, :, 3] == 0).astype(np.uint8) * 255