fix bug of avatarizer; add "-novis" mode

README.md

@@ -86,9 +86,12 @@ ECON is designed for "Human digitization from a color image", which combines the
 ## Demo
 
 ```bash
-# For single-person image-based reconstruction
+# For single-person image-based reconstruction (w/ all visualization steps, 1.8min)
 python -m apps.infer -cfg ./configs/econ.yaml -in_dir ./examples -out_dir ./results
 
+# For single-person image-based reconstruction (w/o any visualization steps, 1.5min)
+python -m apps.infer -cfg ./configs/econ.yaml -in_dir ./examples -out_dir ./results -novis
+
 # For multi-person image-based reconstruction (see config/econ.yaml)
 python -m apps.infer -cfg ./configs/econ.yaml -in_dir ./examples -out_dir ./results -multi
 
@@ -105,7 +108,7 @@ python -m apps.avatarizer -n {filename}
 
 - `use_ifnet: False`
   - True: use IF-Nets+ for mesh completion ( $\text{ECON}_\text{IF}$ - Better quality, **~2min / img**)
-  - False: use SMPL-X for mesh completion ( $\text{ECON}_\text{EX}$ - Faster speed, **~1.5min / img**)
+  - False: use SMPL-X for mesh completion ( $\text{ECON}_\text{EX}$ - Faster speed, **~1.8min / img**)
 - `use_smpl: ["hand", "face"]`
   - [ ]: don't use either hands or face parts from SMPL-X
   - ["hand"]: only use the **visible** hands from SMPL-X

apps/avatarizer.py

@@ -146,7 +146,7 @@ if not osp.exists(f"{prefix}_econ_da.obj") or not osp.exists(f"{prefix}_smpl_da.
     smpl_hand.update_faces(smplx_container.mano_vertex_mask.numpy()[smpl_hand.faces].all(axis=1))
     smpl_hand.remove_unreferenced_vertices()
     econ_da = sum([smpl_hand, smpl_da_body, econ_da_body])
-    econ_da = poisson(econ_da, f"{prefix}_econ_da.obj")
+    econ_da = poisson(econ_da, f"{prefix}_econ_da.obj", depth=10, decimation=False)
 else:
     econ_da = trimesh.load(f"{prefix}_econ_da.obj")
     smpl_da = trimesh.load(f"{prefix}_smpl_da.obj", maintain_orders=True, process=False)
@@ -156,16 +156,16 @@ dist, idx = smpl_tree.query(econ_da.vertices, k=5)
 knn_weights = np.exp(-dist**2)
 knn_weights /= knn_weights.sum(axis=1, keepdims=True)
 
-econ_J_regressor = (smpl_model.J_regressor[:, idx] * knn_weights[None]).sum(axis=-1)
-econ_lbs_weights = (smpl_model.lbs_weights.T[:, idx] * knn_weights[None]).sum(axis=-1).T
+econ_J_regressor = (smpl_model.J_regressor[:, idx] * knn_weights[None]).sum(dim=-1)
+econ_lbs_weights = (smpl_model.lbs_weights.T[:, idx] * knn_weights[None]).sum(dim=-1).T
 
 num_posedirs = smpl_model.posedirs.shape[0]
 econ_posedirs = (
     smpl_model.posedirs.view(num_posedirs, -1, 3)[:, idx, :] * knn_weights[None, ..., None]
-).sum(axis=-2).view(num_posedirs, -1).float()
+).sum(dim=-2).view(num_posedirs, -1).float()
 
-econ_J_regressor /= econ_J_regressor.sum(axis=1, keepdims=True)
-econ_lbs_weights /= econ_lbs_weights.sum(axis=1, keepdims=True)
+econ_J_regressor /= econ_J_regressor.sum(dim=1, keepdims=True).clip(min=1e-10)
+econ_lbs_weights /= econ_lbs_weights.sum(dim=1, keepdims=True)
 
 # re-compute da-pose rot_mat for ECON
 rot_mat_da = smpl_out_lst[1].vertex_transformation.detach()[0][idx[:, 0]]

apps/infer.py

@@ -54,12 +54,12 @@ if __name__ == "__main__":
     parser.add_argument("-gpu", "--gpu_device", type=int, default=0)
     parser.add_argument("-loop_smpl", "--loop_smpl", type=int, default=50)
     parser.add_argument("-patience", "--patience", type=int, default=5)
-    parser.add_argument("-vis_freq", "--vis_freq", type=int, default=1000)
-    parser.add_argument("-multi", action="store_false")
     parser.add_argument("-in_dir", "--in_dir", type=str, default="./examples")
     parser.add_argument("-out_dir", "--out_dir", type=str, default="./results")
     parser.add_argument("-seg_dir", "--seg_dir", type=str, default=None)
     parser.add_argument("-cfg", "--config", type=str, default="./configs/econ.yaml")
+    parser.add_argument("-multi", action="store_false")
+    parser.add_argument("-novis", action="store_true")
 
     args = parser.parse_args()
 
@@ -319,8 +319,8 @@ if __name__ == "__main__":
                 pbar_desc += colored(f"| loose:{loose_str}, occluded:{occlude_str}", "yellow")
                 loop_smpl.set_description(pbar_desc)
 
-                # save intermediate results / vis_freq and final_step
-                if (i % args.vis_freq == 0) or (i == args.loop_smpl - 1):
+                # save intermediate results
+                if (i == args.loop_smpl - 1) and (not args.novis):
 
                     per_loop_lst.extend(
                         [
@@ -348,26 +348,32 @@ if __name__ == "__main__":
 
             in_tensor["smpl_verts"] = smpl_verts * torch.tensor([1.0, 1.0, -1.0]).to(device)
             in_tensor["smpl_faces"] = in_tensor["smpl_faces"][:, :, [0, 2, 1]]
-            per_data_lst[-1].save(osp.join(args.out_dir, cfg.name, f"png/{data['name']}_smpl.png"))
-
-        img_crop_path = osp.join(args.out_dir, cfg.name, "png", f"{data['name']}_crop.png")
-        torchvision.utils.save_image(
-            torch.cat(
-                [
-                    data["img_crop"][:, :3], (in_tensor['normal_F'].detach().cpu() + 1.0) * 0.5,
-                    (in_tensor['normal_B'].detach().cpu() + 1.0) * 0.5
-                ],
-                dim=3
-            ), img_crop_path
-        )
 
-        rgb_norm_F = blend_rgb_norm(in_tensor["normal_F"], data)
-        rgb_norm_B = blend_rgb_norm(in_tensor["normal_B"], data)
+            if not args.novis:
+                per_data_lst[-1].save(
+                    osp.join(args.out_dir, cfg.name, f"png/{data['name']}_smpl.png")
+                )
 
-        img_overlap_path = osp.join(args.out_dir, cfg.name, f"png/{data['name']}_overlap.png")
-        torchvision.utils.save_image(
-            torch.cat([data["img_raw"], rgb_norm_F, rgb_norm_B], dim=-1) / 255., img_overlap_path
-        )
+        if not args.novis:
+            img_crop_path = osp.join(args.out_dir, cfg.name, "png", f"{data['name']}_crop.png")
+            torchvision.utils.save_image(
+                torch.cat(
+                    [
+                        data["img_crop"][:, :3], (in_tensor['normal_F'].detach().cpu() + 1.0) * 0.5,
+                        (in_tensor['normal_B'].detach().cpu() + 1.0) * 0.5
+                    ],
+                    dim=3
+                ), img_crop_path
+            )
+
+            rgb_norm_F = blend_rgb_norm(in_tensor["normal_F"], data)
+            rgb_norm_B = blend_rgb_norm(in_tensor["normal_B"], data)
+
+            img_overlap_path = osp.join(args.out_dir, cfg.name, f"png/{data['name']}_overlap.png")
+            torchvision.utils.save_image(
+                torch.cat([data["img_raw"], rgb_norm_F, rgb_norm_B], dim=-1) / 255.,
+                img_overlap_path
+            )
 
         smpl_obj_lst = []
 
@@ -618,12 +624,13 @@ if __name__ == "__main__":
                 final_mesh = sum(full_lst)
                 final_mesh.export(final_path)
 
-            dataset.render.load_meshes(final_mesh.vertices, final_mesh.faces)
-            rotate_recon_lst = dataset.render.get_image(cam_type="four")
-            per_loop_lst.extend([in_tensor['image'][idx:idx + 1]] + rotate_recon_lst)
+            if not args.novis:
+                dataset.render.load_meshes(final_mesh.vertices, final_mesh.faces)
+                rotate_recon_lst = dataset.render.get_image(cam_type="four")
+                per_loop_lst.extend([in_tensor['image'][idx:idx + 1]] + rotate_recon_lst)
 
             if cfg.bni.texture_src == 'image':
-                
+
                 # coloring the final mesh (front: RGB pixels, back: normal colors)
                 final_colors = query_color(
                     torch.tensor(final_mesh.vertices).float(),
@@ -633,22 +640,24 @@ if __name__ == "__main__":
                 )
                 final_mesh.visual.vertex_colors = final_colors
                 final_mesh.export(final_path)
-                
+
             elif cfg.bni.texture_src == 'SD':
-                
+
                 # !TODO: add texture from Stable Diffusion
                 pass
 
-            # for video rendering
-            in_tensor["BNI_verts"].append(torch.tensor(final_mesh.vertices).float())
-            in_tensor["BNI_faces"].append(torch.tensor(final_mesh.faces).long())
-
-        if len(per_loop_lst) > 0:
+        if len(per_loop_lst) > 0 and (not args.novis):
 
             per_data_lst.append(get_optim_grid_image(per_loop_lst, None, nrow=5, type="cloth"))
             per_data_lst[-1].save(osp.join(args.out_dir, cfg.name, f"png/{data['name']}_cloth.png"))
 
-        os.makedirs(osp.join(args.out_dir, cfg.name, "vid"), exist_ok=True)
-        in_tensor["uncrop_param"] = data["uncrop_param"]
-        in_tensor["img_raw"] = data["img_raw"]
-        torch.save(in_tensor, osp.join(args.out_dir, cfg.name, f"vid/{data['name']}_in_tensor.pt"))
+            # for video rendering
+            in_tensor["BNI_verts"].append(torch.tensor(final_mesh.vertices).float())
+            in_tensor["BNI_faces"].append(torch.tensor(final_mesh.faces).long())
+
+            os.makedirs(osp.join(args.out_dir, cfg.name, "vid"), exist_ok=True)
+            in_tensor["uncrop_param"] = data["uncrop_param"]
+            in_tensor["img_raw"] = data["img_raw"]
+            torch.save(
+                in_tensor, osp.join(args.out_dir, cfg.name, f"vid/{data['name']}_in_tensor.pt")
+            )

lib/dataset/mesh_util.py

@@ -384,7 +384,7 @@ def remesh_laplacian(mesh, obj_path):
     return mesh
 
 
-def poisson(mesh, obj_path, depth=10):
+def poisson(mesh, obj_path, depth=10, decimation=True):
 
     pcd_path = obj_path[:-4] + ".ply"
     assert (mesh.vertex_normals.shape[1] == 3)
@@ -400,10 +400,12 @@ def poisson(mesh, obj_path, depth=10):
     largest_mesh = keep_largest(trimesh.Trimesh(np.array(mesh.vertices), np.array(mesh.triangles)))
     largest_mesh.export(obj_path)
 
-    # mesh decimation for faster rendering
-    low_res_mesh = largest_mesh.simplify_quadratic_decimation(50000)
-
-    return low_res_mesh
+    if decimation:
+        # mesh decimation for faster rendering
+        low_res_mesh = largest_mesh.simplify_quadratic_decimation(50000)
+        return low_res_mesh
+    else:
+        return largest_mesh
 
 
 # Losses to smooth / regularize the mesh shape