Refactor generate3d function in inference.py to implement CCM-based UV assignment for improved texture mapping. Updated color map preparation and vertex projection to enhance mesh quality in exports, ensuring compatibility with various input formats.

inference.py

@@ -62,24 +62,36 @@ def generate3d(model, rgb, ccm, device):
     data_config['verts'] = torch.from_numpy(verts).contiguous()
     data_config['faces'] = torch.from_numpy(faces).contiguous()
 
-    # Generate UVs using xatlas (CPU)
-    mesh_v = data_config['verts'].cpu().numpy()
-    mesh_f = data_config['faces'].cpu().numpy()
-    vmapping, ft, vt = xatlas.parametrize(mesh_v, mesh_f)
-
-    # Load your color map (photorealistic texture)
-    # Assume you save the color map as a PNG somewhere, or pass it as a numpy array
-    # If rgb is your color map, convert it to the right format:
-    if rgb.max() > 1.0:
-        color_map = rgb.astype(np.uint8)
-    else:
-        color_map = (rgb * 255).astype(np.uint8)
-    # If rgb is already in HWC, skip transpose; otherwise, transpose as needed
+    # --- CCM-based UV assignment ---
+    mesh_v = data_config['verts'].cpu().numpy()  # [N, 3]
+    mesh_f = data_config['faces'].cpu().numpy()  # [M, 3]
+
+    # Prepare CCM and color map
+    ccm_img = ccm.astype(np.uint8) if ccm.max() > 1 else (ccm * 255).astype(np.uint8)
+    if ccm_img.shape[-1] != 3:
+        ccm_img = np.transpose(ccm_img, (1, 2, 0))
+    H, W, _ = ccm_img.shape
+
+    color_map = rgb.astype(np.uint8) if rgb.max() > 1 else (rgb * 255).astype(np.uint8)
     if color_map.shape[-1] != 3:
         color_map = np.transpose(color_map, (1, 2, 0))
     albedo = cv2.cvtColor(color_map, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
 
-    tex_res = (albedo.shape[0], albedo.shape[1])
+    # Project mesh vertices to CCM image space and get UVs
+    vt = []
+    for v in mesh_v:
+        # Assume mesh is in [-1,1] in x/y, project to CCM image
+        x, y, z = v
+        u_img = int((x + 1) / 2 * (W - 1))
+        v_img = int((y + 1) / 2 * (H - 1))
+        u_img = np.clip(u_img, 0, W-1)
+        v_img = np.clip(v_img, 0, H-1)
+        r, g, b = ccm_img[v_img, u_img]
+        u = r / 255.0
+        v_ = g / 255.0
+        vt.append([u, v_])
+    vt = np.array(vt, dtype=np.float32)
+    ft = mesh_f.copy()
 
     # Create Mesh and export .glb
     mesh = Mesh(