update

app.py

@@ -47,8 +47,8 @@ This is official demo for our CVPR 2025 paper <a href="">StdGEN: Semantic-Decomp
 
 Code: <a href='https://github.com/hyz317/StdGEN' target='_blank'>GitHub</a>. Paper: <a href='https://arxiv.org/abs/2411.05738' target='_blank'>ArXiv</a>.
 
-❗️❗️❗️**Important Notes:**
-1. Refinement stage takes about ~3.5min, and the mesh result may possibly delayed due to the server load, please wait patiently.
+❗️❗️❗️**Important Notes:** This is only a **PREVIEW** version with lower quality. We only perform color back-projection to clothes and hair. Please refer to GitHub repo for complete version.
+1. Refinement stage takes about ~2.5min, and the mesh result may possibly delayed due to the server load, please wait patiently.
 
 2. You can upload any reference image (with or without background), A-pose images are also supported (white bkg required). If the image has an alpha channel (transparency), background segmentation will be automatically performed. Alternatively, you can pre-segment the background using other tools and upload the result directly.  
 

infer_api.py

@@ -542,19 +542,19 @@ def get_distract_mask(generator, color_0, color_1, normal_0=None, normal_1=None,
     return distract_mask, distract_bbox, random_sampled_points, final_mask
 
 
-infer_refine_sam = sam_model_registry["vit_h"](checkpoint="./ckpt/sam_vit_h_4b8939.pth").cuda()
-infer_refine_generator = SamAutomaticMaskGenerator(
-    model=infer_refine_sam,
-    points_per_side=64,
-    pred_iou_thresh=0.80,
-    stability_score_thresh=0.92,
-    crop_n_layers=1,
-    crop_n_points_downscale_factor=2,
-    min_mask_region_area=100,
-)
+# infer_refine_sam = sam_model_registry["vit_h"](checkpoint="./ckpt/sam_vit_h_4b8939.pth").cuda()
+# infer_refine_generator = SamAutomaticMaskGenerator(
+#     model=infer_refine_sam,
+#     points_per_side=64,
+#     pred_iou_thresh=0.80,
+#     stability_score_thresh=0.92,
+#     crop_n_layers=1,
+#     crop_n_points_downscale_factor=2,
+#     min_mask_region_area=100,
+# )
 infer_refine_outside_ratio = 0.20
 
-@spaces.GPU(duration=150)
+
 def infer_refine(meshes, imgs):
     fixed_v, fixed_f, fixed_t = None, None, None
     flow_vert, flow_vector = None, None
@@ -564,7 +564,6 @@ def infer_refine(meshes, imgs):
 
     mv, proj = make_star_cameras_orthographic(8, 1, r=1.2)
     mv = mv[[4, 3, 2, 0, 6, 5]]        
-    renderer = NormalsRenderer(mv,proj,(1024,1024))
 
     results = []
 
@@ -576,12 +575,17 @@ def infer_refine(meshes, imgs):
         mesh_v, mesh_f = mesh.vertices, mesh.faces
 
         if last_colors is None:
-            images = renderer.render(
-                torch.tensor(mesh_v, device='cuda').float(),
-                torch.ones_like(torch.from_numpy(mesh_v), device='cuda').float(),
-                torch.tensor(mesh_f, device='cuda'),
-            )
-            mask = (images[..., 3] < 0.9).cpu().numpy()
+            @spaces.GPU()
+            def get_mask():
+                renderer = NormalsRenderer(mv,proj,(1024,1024))
+                images = renderer.render(
+                    torch.tensor(mesh_v, device='cuda').float(),
+                    torch.ones_like(torch.from_numpy(mesh_v), device='cuda').float(),
+                    torch.tensor(mesh_f, device='cuda'),
+                )
+                mask = (images[..., 3] < 0.9).cpu().numpy()
+                return mask
+            mask = get_mask()
 
         colors, normals = [], []
         for i in range(6):
@@ -604,18 +608,15 @@ def infer_refine(meshes, imgs):
             colors.append(color)
             normals.append(normal)
 
-        if last_front_color is not None and level == 0:
-            original_mask, distract_bbox, _, distract_mask = get_distract_mask(infer_refine_generator, last_front_color, np.array(colors[0]).astype(np.float32) / 255.0, outside_ratio=infer_refine_outside_ratio)
-        else:  
-            distract_mask = None
-            distract_bbox = None
-
-        last_front_color = np.array(colors[0]).astype(np.float32) / 255.0
-        last_front_normal = np.array(normals[0]).astype(np.float32) / 255.0
+        # if last_front_color is not None and level == 0:
+        #     original_mask, distract_bbox, _, distract_mask = get_distract_mask(infer_refine_generator, last_front_color, np.array(colors[0]).astype(np.float32) / 255.0, outside_ratio=infer_refine_outside_ratio)
+        # else:  
+        distract_mask = None
+        distract_bbox = None
 
         if last_colors is None:
             from copy import deepcopy
-            last_colors, last_normals = deepcopy(colors), deepcopy(normals)
+            last_colors = deepcopy(colors)
 
         # my mesh flow weight by nearest vertexs
         if fixed_v is not None and fixed_f is not None and level == 1:
@@ -643,8 +644,8 @@ def infer_refine(meshes, imgs):
 
             t = trimesh.Trimesh(vertices=mesh_v, faces=mesh_f)
 
-        mesh_v = torch.tensor(mesh_v, device='cuda', dtype=torch.float32)
-        mesh_f = torch.tensor(mesh_f, device='cuda')
+        mesh_v = torch.tensor(mesh_v, dtype=torch.float32)
+        mesh_f = torch.tensor(mesh_f)
 
         new_mesh, simp_v, simp_f = geo_refine(mesh_v, mesh_f, colors, normals, fixed_v=fixed_v, fixed_f=fixed_f, distract_mask=distract_mask, distract_bbox=distract_bbox)
 
@@ -659,22 +660,22 @@ def infer_refine(meshes, imgs):
                 _, idx_anchor = kdtree_anchor.query(new_mesh_v, k=1)
                 _, idx_mesh_v = kdtree_mesh_v.query(new_mesh_v, k=25)
                 idx_anchor = idx_anchor.squeeze()
-                neighbors = torch.tensor(new_mesh_v).cuda()[idx_mesh_v]  # V, 25, 3
+                neighbors = torch.tensor(new_mesh_v)[idx_mesh_v]  # V, 25, 3
                 # calculate the distances neighbors [V, 25, 3]; new_mesh_v [V, 3] -> [V, 25]
-                neighbor_dists = torch.norm(neighbors - torch.tensor(new_mesh_v).cuda()[:, None], dim=-1)
+                neighbor_dists = torch.norm(neighbors - torch.tensor(new_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(new_mesh_v).cuda()) * anchor_normals).sum(-1), min=0) + 0.01
+                dis_anchor = torch.clamp(((anchors - torch.tensor(new_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
                 new_mesh_v += weighted_vec_anchor.cpu().numpy()
 
                 # replace new_mesh verts with new_mesh_v
-                new_mesh = Meshes(verts=[torch.tensor(new_mesh_v, device='cuda')], faces=new_mesh.faces_list(), textures=new_mesh.textures)
+                new_mesh = Meshes(verts=[torch.tensor(new_mesh_v)], faces=new_mesh.faces_list(), textures=new_mesh.textures)
 
         except Exception as e:
             pass

refine/mesh_refine.py

@@ -1,4 +1,5 @@
 import torch
+import spaces
 import numpy as np
 import trimesh
 from PIL import Image
@@ -267,6 +268,15 @@ 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):
+    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)
+    return geo_refine_3(vertices, faces, rgb_ls, fixed_v=fixed_v, fixed_f=fixed_f, distract_mask=distract_mask)
+
+@spaces.GPU()
+def geo_refine_1(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):
+    
     rm_normals = simple_remove(normal_ls)
 
     # transfer the alpha channel of rm_normals to img_list
@@ -282,6 +292,9 @@ def geo_refine(mesh_v, mesh_f, rgb_ls, normal_ls, expansion_weight=0.1, fixed_v=
     if no_decompose:
         stage1_lr = 0.03
         stage1_remesh_interval = 30
+    
+    if fixed_v is not None:
+        return mesh_v, mesh_f
 
     vertices, faces = reconstruct_stage1(rm_normals, steps=200, vertices=mesh_v, faces=mesh_f, fixed_v=fixed_v, fixed_f=fixed_f,
                                          lr=stage1_lr, remesh_interval=stage1_remesh_interval, start_edge_len=0.02,
@@ -290,20 +303,18 @@ def geo_refine(mesh_v, mesh_f, rgb_ls, normal_ls, expansion_weight=0.1, fixed_v=
 
     vertices, faces = run_mesh_refine(vertices, faces, rm_normals, fixed_v=fixed_v, fixed_f=fixed_f, steps=100, start_edge_len=0.005, end_edge_len=0.0002,
                                       decay=0.99, update_normal_interval=20, update_warmup=5, process_inputs=False, process_outputs=False, remesh_interval=1)
-    meshes = simple_clean_mesh(to_pyml_mesh(vertices, faces), apply_smooth=True, stepsmoothnum=2, apply_sub_divide=False, sub_divide_threshold=0.25).to("cuda")
-    # subdivide meshes
+    return vertices, faces
+
+def geo_refine_2(vertices, faces, fixed_v=None):
+    meshes = simple_clean_mesh(to_pyml_mesh(vertices, faces), apply_smooth=True, stepsmoothnum=2, apply_sub_divide=False, sub_divide_threshold=0.25)
     simp_vertices, simp_faces = meshes.verts_packed(), meshes.faces_packed()
     vertices, faces = simp_vertices.detach().cpu().numpy(), simp_faces.detach().cpu().numpy()
+    if fixed_v is not None:
+        vertices, faces = trimesh.remesh.subdivide(vertices, faces)
+    return vertices, faces
 
-    mesh = trimesh.Trimesh(vertices=vertices, faces=faces, process=False)
-    mesh = merge_small_faces(mesh, thres=thres)
-    new_mesh = mesh.split(only_watertight=False)
-
-    new_mesh = [ j for j in new_mesh if len(j.vertices) >= 200 ]
-    mesh = trimesh.Scene(new_mesh).dump(concatenate=True)
-    vertices, faces = mesh.vertices.astype('float32'), mesh.faces
-
-    vertices, faces = trimesh.remesh.subdivide(vertices, faces)
+@spaces.GPU()
+def geo_refine_3(vertices, faces, rgb_ls, fixed_v=None, fixed_f=None, distract_mask=None):
     origin_len_v, origin_len_f = len(vertices), len(faces)
     # concatenate fixed_v and fixed_f
     if fixed_v is not None and fixed_f is not None:
@@ -316,4 +327,4 @@ def geo_refine(mesh_v, mesh_f, rgb_ls, normal_ls, expansion_weight=0.1, fixed_v=
     if fixed_v is not None and fixed_f is not None:
         new_meshes = Meshes(verts=[new_meshes.verts_packed()[:origin_len_v]], faces=[new_meshes.faces_packed()[:origin_len_f]],
                             textures=pytorch3d.renderer.mesh.textures.TexturesVertex([new_meshes.textures.verts_features_packed()[:origin_len_v]]))
-    return new_meshes, simp_vertices, simp_faces
+    return new_meshes.to("cpu"), vertices, faces