update

asset3d_gen/data/backproject_v2.py

@@ -266,12 +266,13 @@ class TextureBacker:
 
     def get_mesh_np_attrs(
         self,
+        mesh: trimesh.Trimesh,
         scale: float = None,
         center: np.ndarray = None,
     ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
-        vertices = self.vertices.cpu().numpy()
-        faces = self.faces.cpu().numpy()
-        uv_map = self.uv_map.cpu().numpy()
+        vertices = mesh.vertices
+        faces = mesh.faces
+        uv_map = mesh.uv_map
         uv_map[:, 1] = 1.0 - uv_map[:, 1]
 
         if scale is not None:
@@ -400,9 +401,9 @@ class TextureBacker:
         return texture_merge, trust_map_merge > 1e-8
 
     def uv_inpaint(
-        self, texture: np.ndarray, mask: np.ndarray
+        self, mesh: trimesh.Trimesh, texture: np.ndarray, mask: np.ndarray
     ) -> np.ndarray:
-        vertices, faces, uv_map = self.get_mesh_np_attrs()
+        vertices, faces, uv_map = self.get_mesh_np_attrs(mesh)
 
         texture, mask = _texture_inpaint_smooth(
             texture, mask, vertices, faces, uv_map
@@ -424,7 +425,7 @@ class TextureBacker:
         mesh: trimesh.Trimesh,
     ) -> trimesh.Trimesh:
         self._lazy_init_render(self.camera_params, self.mask_thresh)
-        
+
         self.vertices = torch.from_numpy(mesh.vertices).to(self.device).float()
         self.faces = torch.from_numpy(mesh.faces).to(self.device).to(torch.int)
         self.uv_map = torch.from_numpy(mesh.visual.uv).to(self.device).float()
@@ -470,10 +471,10 @@ class TextureBacker:
 
         texture_np, mask_np = self.cuda_forward(colors, mesh)
         
-        texture_np = self.uv_inpaint(texture_np, mask_np)
+        texture_np = self.uv_inpaint(mesh, texture_np, mask_np)
         texture_np = post_process_texture(texture_np)
         vertices, faces, uv_map = self.get_mesh_np_attrs(
-            self.scale, self.center
+            mesh, self.scale, self.center
         )
         textured_mesh = save_mesh_with_mtl(
             vertices, faces, uv_map, texture_np, output_path