update

infer_api.py

@@ -619,30 +619,26 @@ def infer_refine(meshes, imgs):
 
         # my mesh flow weight by nearest vertexs
         if fixed_v is not None and fixed_f is not None and level == 1:
-            t = trimesh.Trimesh(vertices=mesh_v, faces=mesh_f)
-
             fixed_v_cpu = fixed_v.cpu().numpy()
             kdtree_anchor = KDTree(fixed_v_cpu)
             kdtree_mesh_v = KDTree(mesh_v)
             _, idx_anchor = kdtree_anchor.query(mesh_v, k=1)
             _, idx_mesh_v = kdtree_mesh_v.query(mesh_v, k=25)
             idx_anchor = idx_anchor.squeeze()
-            neighbors = torch.tensor(mesh_v).cuda()[idx_mesh_v]  # V, 25, 3
+            neighbors = torch.tensor(mesh_v)[idx_mesh_v]  # V, 25, 3
             # calculate the distances neighbors [V, 25, 3]; mesh_v [V, 3] -> [V, 25]
-            neighbor_dists = torch.norm(neighbors - torch.tensor(mesh_v).cuda()[:, None], dim=-1)
+            neighbor_dists = torch.norm(neighbors - torch.tensor(mesh_v)[:, None], dim=-1)
             neighbor_dists[neighbor_dists > 0.06] = 114514.
             neighbor_weights = torch.exp(-neighbor_dists * 1.)
             neighbor_weights = neighbor_weights / neighbor_weights.sum(dim=1, keepdim=True)
             anchors = fixed_v[idx_anchor]  # V, 3
             anchor_normals = calc_vertex_normals(fixed_v, fixed_f)[idx_anchor]  # V, 3
-            dis_anchor = torch.clamp(((anchors - torch.tensor(mesh_v).cuda()) * anchor_normals).sum(-1), min=0) + 0.01
+            dis_anchor = torch.clamp(((anchors - torch.tensor(mesh_v)) * anchor_normals).sum(-1), min=0) + 0.01
             vec_anchor = dis_anchor[:, None] * anchor_normals  # V, 3
             vec_anchor = vec_anchor[idx_mesh_v]  # V, 25, 3
             weighted_vec_anchor = (vec_anchor * neighbor_weights[:, :, None]).sum(1)  # V, 3
             mesh_v += weighted_vec_anchor.cpu().numpy()
 
-            t = trimesh.Trimesh(vertices=mesh_v, faces=mesh_f)
-
         mesh_v = torch.tensor(mesh_v, dtype=torch.float32)
         mesh_f = torch.tensor(mesh_f)
 

refine/mesh_refine.py

@@ -268,6 +268,11 @@ def run_mesh_refine(vertices, faces, pils: List[Image.Image], fixed_v=None, fixe
 
 def geo_refine(mesh_v, mesh_f, rgb_ls, normal_ls, expansion_weight=0.1, fixed_v=None, fixed_f=None,
                distract_mask=None, distract_bbox=None, thres=3e-6, no_decompose=False):
+    print(mesh_v.device, mesh_f.device)
+    if fixed_v is not None:
+        print('fixed_v', fixed_v.shape, fixed_v.device)
+    if fixed_f is not None:
+        print('fixed_f', fixed_f.shape, fixed_f.device)
     vertices, faces = geo_refine_1(mesh_v, mesh_f, rgb_ls, normal_ls, expansion_weight=expansion_weight, fixed_v=fixed_v, fixed_f=fixed_f,
                  distract_mask=distract_mask, distract_bbox=distract_bbox, thres=thres, no_decompose=no_decompose)
     vertices, faces = geo_refine_2(vertices, faces, fixed_v=fixed_v)