SMPL-X based Animatable Avatar

.gitignore

@@ -16,3 +16,4 @@ build
 dist
 *egg-info
 *.so
+run.sh

README.md

@@ -23,6 +23,8 @@
   <br>
     <a href="https://pytorch.org/get-started/locally/"><img alt="PyTorch" src="https://img.shields.io/badge/PyTorch-ee4c2c?logo=pytorch&logoColor=white"></a>
     <a href="https://pytorchlightning.ai/"><img alt="Lightning" src="https://img.shields.io/badge/-Lightning-792ee5?logo=pytorchlightning&logoColor=white"></a>
+    <a href="https://cupy.dev/"><img alt="cupy" src="https://img.shields.io/badge/-Cupy-46C02B?logo=numpy&logoColor=white"></a>
+    <a href="https://twitter.com/yuliangxiu"><img alt='Twitter' src="https://img.shields.io/twitter/follow/yuliangxiu?label=%40yuliangxiu"></a>
     <br></br>
     <a href=''>
       <img src='https://img.shields.io/badge/Paper-PDF (coming soon)-green?style=for-the-badge&logo=arXiv&logoColor=green' alt='Paper PDF'>
@@ -36,7 +38,7 @@
 
 <br/>
 
-ECON is designed for **"Human digitization from a color image"**, which combines the best properties of implicit and explicit representations, to infer high-fidelity 3D clothed humans from in-the-wild images, even with **loose clothing** or in **challenging poses**. ECON also supports batch reconstruction from **multi-person** photos.
+ECON is designed for "Human digitization from a color image", which combines the best properties of implicit and explicit representations, to infer high-fidelity 3D clothed humans from in-the-wild images, even with **loose clothing** or in **challenging poses**. ECON also supports **multi-person reconstruction** and **SMPL-X based animation**.
 <br/>
 <br/>
 
@@ -61,6 +63,9 @@ ECON is designed for **"Human digitization from a color image"**, which combines
     <li>
       <a href="#demo">Demo</a>
     </li>
+    <li>
+      <a href="#applications">Applications</a>
+    </li>
     <li>
       <a href="#tricks">Tricks</a>
     </li>
@@ -87,6 +92,9 @@ python -m apps.infer -cfg ./configs/econ.yaml -in_dir ./examples -out_dir ./resu
 
 # To generate the demo video of reconstruction results
 python -m apps.multi_render -n {filename}
+
+# To animate the reconstruction with SMPL-X pose parameters
+python -m apps.avatarizer -n {filename}
 ```
 
 ## Tricks
@@ -101,24 +109,28 @@ python -m apps.multi_render -n {filename}
   - ["hand"]: only use the **visible** hands from SMPL-X
   - ["hand", "face"]: use both **visible** hands and face from SMPL-X
 - `thickness: 2cm`
-  - could be increased accordingly in case **xx_full.obj** looks flat
-- `hps_type: pixie`
+  - could be increased accordingly in case final reconstruction **xx_full.obj** looks flat
+- `hps_type: PIXIE`
   - "pixie": more accurate for face and hands
   - "pymafx": more robust for challenging poses
+- `k: 4`
+  - could be reduced accordingly in case the surface of **xx_full.obj** has discontinous artifacts
 
 <br/>
 
 ## More Qualitative Results
 
-|                         ![OOD Poses](assets/OOD-poses.jpg)                         |
-| :--------------------------------------------------------------------------------: |
-|                                _Challenging Poses_                                 |
-|                       ![OOD Clothes](assets/OOD-outfits.jpg)                       |
-|                                  _Loose Clothes_                                   |
-|                              ![SHHQ](assets/SHHQ.gif)                              |
-| _ECON Results on [SHHQ Dataset](https://github.com/stylegan-human/StyleGAN-Human)_ |
-|                             ![crowd](assets/crowd.gif)                             |
-|                        _ECON Results on Multi-Person Image_                        |
+|   ![OOD Poses](assets/OOD-poses.jpg)   |
+| :------------------------------------: |
+|          _Challenging Poses_           |
+| ![OOD Clothes](assets/OOD-outfits.jpg) |
+|            _Loose Clothes_             |
+
+## Applications
+
+|                                        ![SHHQ](assets/SHHQ.gif)                                        |         ![crowd](assets/crowd.gif)          |
+| :----------------------------------------------------------------------------------------------------: | :-----------------------------------------: |
+| _ECON could provide pseudo 3D GT for [SHHQ Dataset](https://github.com/stylegan-human/StyleGAN-Human)_ | _ECON supports multi-person reconstruction_ |
 
 <br/>
 <br/>
@@ -127,7 +139,7 @@ python -m apps.multi_render -n {filename}
 
 ```bibtex
 @misc{xiu2022econ,
-    title={ECON: Explicit Clothed humans Obtained from Normals},
+    title={{ECON: Explicit Clothed humans Obtained from Normals}},
     author={Xiu, Yuliang and Yang, Jinlong and Cao, Xu and Tzionas, Dimitrios and Black, Michael J.},
     year={2022}
     publisher={arXiv},
@@ -146,6 +158,7 @@ Here are some great resources we benefit from:
 - [ICON](https://github.com/YuliangXiu/ICON) for Body Fitting
 - [MonoPortDataset](https://github.com/Project-Splinter/MonoPortDataset) for Data Processing
 - [rembg](https://github.com/danielgatis/rembg) for Human Segmentation
+- [PyTorch-NICP](https://github.com/wuhaozhe/pytorch-nicp) for non-rigid registration
 - [smplx](https://github.com/vchoutas/smplx), [PyMAF-X](https://www.liuyebin.com/pymaf-x/), [PIXIE](https://github.com/YadiraF/PIXIE) for Human Pose & Shape Estimation
 - [CAPE](https://github.com/qianlim/CAPE) and [THuman](https://github.com/ZhengZerong/DeepHuman/tree/master/THUmanDataset) for Dataset
 - [PyTorch3D](https://github.com/facebookresearch/pytorch3d) for Differential Rendering
@@ -171,4 +184,3 @@ MJB has received research gift funds from Adobe, Intel, Nvidia, Meta/Facebook, a
 For technical questions, please contact [email protected]
 
 For commercial licensing, please contact [email protected]
-

apps/avatarizer.py

@@ -1,6 +1,7 @@
 import numpy as np
 import trimesh
 import torch
+import argparse
 import os.path as osp
 import lib.smplx as smplx
 from pytorch3d.ops import SubdivideMeshes
@@ -12,10 +13,16 @@ from scipy.spatial import cKDTree
 from lib.dataset.mesh_util import SMPLX
 from lib.common.local_affine import register
 
+# loading cfg file
+parser = argparse.ArgumentParser()
+parser.add_argument("-n", "--name", type=str, default="")
+parser.add_argument("-g", "--gpu", type=int, default=0)
+args = parser.parse_args()
+
 smplx_container = SMPLX()
-device = torch.device("cuda:0")
+device = torch.device(f"cuda:{args.gpu}")
 
-prefix = "./results/github/econ/obj/304e9c4798a8c3967de7c74c24ef2e38"
+prefix = f"./results/econ/obj/{args.name}"
 smpl_path = f"{prefix}_smpl_00.npy"
 econ_path = f"{prefix}_0_full.obj"
 
@@ -27,7 +34,6 @@ econ_obj.vertices -= smplx_param["transl"].cpu().numpy()
 
 for key in smplx_param.keys():
     smplx_param[key] = smplx_param[key].cpu().view(1, -1)
-    # print(key, smplx_param[key].device, smplx_param[key].shape)
 
 smpl_model = smplx.create(
     smplx_container.model_dir,
@@ -40,109 +46,135 @@ smpl_model = smplx.create(
     num_expression_coeffs=50,
     ext='pkl')
 
-smpl_out = smpl_model(
-    body_pose=smplx_param["body_pose"],
-    global_orient=smplx_param["global_orient"],
-    betas=smplx_param["betas"],
-    expression=smplx_param["expression"],
-    jaw_pose=smplx_param["jaw_pose"],
-    left_hand_pose=smplx_param["left_hand_pose"],
-    right_hand_pose=smplx_param["right_hand_pose"],
-    return_verts=True,
-    return_full_pose=True,
-    return_joint_transformation=True,
-    return_vertex_transformation=True)
-
-smpl_verts = smpl_out.vertices.detach()[0]
+smpl_out_lst = []
+
+for pose_type in ["t-pose", "da-pose", "pose"]:
+    smpl_out_lst.append(
+        smpl_model(
+            body_pose=smplx_param["body_pose"],
+            global_orient=smplx_param["global_orient"],
+            betas=smplx_param["betas"],
+            expression=smplx_param["expression"],
+            jaw_pose=smplx_param["jaw_pose"],
+            left_hand_pose=smplx_param["left_hand_pose"],
+            right_hand_pose=smplx_param["right_hand_pose"],
+            return_verts=True,
+            return_full_pose=True,
+            return_joint_transformation=True,
+            return_vertex_transformation=True,
+            pose_type=pose_type))
+
+smpl_verts = smpl_out_lst[2].vertices.detach()[0]
 smpl_tree = cKDTree(smpl_verts.cpu().numpy())
 dist, idx = smpl_tree.query(econ_obj.vertices, k=5)
 
-if not osp.exists(f"{prefix}_econ_cano.obj") or not osp.exists(f"{prefix}_smpl_cano.obj"):
+if not osp.exists(f"{prefix}_econ_da.obj") or not osp.exists(f"{prefix}_smpl_da.obj"):
 
-    # canonicalize for ECON
+    # t-pose for ECON
     econ_verts = torch.tensor(econ_obj.vertices).float()
-    inv_mat = torch.inverse(smpl_out.vertex_transformation.detach()[0][idx[:, 0]])
+    rot_mat_t = smpl_out_lst[2].vertex_transformation.detach()[0][idx[:, 0]]
     homo_coord = torch.ones_like(econ_verts)[..., :1]
-    econ_cano_verts = inv_mat @ torch.cat([econ_verts, homo_coord], dim=1).unsqueeze(-1)
+    econ_cano_verts = torch.inverse(rot_mat_t) @ torch.cat([econ_verts, homo_coord], dim=1).unsqueeze(-1)
     econ_cano_verts = econ_cano_verts[:, :3, 0].cpu()
     econ_cano = trimesh.Trimesh(econ_cano_verts, econ_obj.faces)
 
-    # canonicalize for SMPL-X
-    inv_mat = torch.inverse(smpl_out.vertex_transformation.detach()[0])
-    homo_coord = torch.ones_like(smpl_verts)[..., :1]
-    smpl_cano_verts = inv_mat @ torch.cat([smpl_verts, homo_coord], dim=1).unsqueeze(-1)
-    smpl_cano_verts = smpl_cano_verts[:, :3, 0].cpu()
-    smpl_cano = trimesh.Trimesh(smpl_cano_verts, smpl_model.faces, maintain_orders=True, process=False)
-    smpl_cano.export(f"{prefix}_smpl_cano.obj")
+    # da-pose for ECON
+    rot_mat_da = smpl_out_lst[1].vertex_transformation.detach()[0][idx[:, 0]]
+    econ_da_verts = rot_mat_da @ torch.cat([econ_cano_verts, homo_coord], dim=1).unsqueeze(-1)
+    econ_da = trimesh.Trimesh(econ_da_verts[:, :3, 0].cpu(), econ_obj.faces)
+
+    # da-pose for SMPL-X
+    smpl_da = trimesh.Trimesh(smpl_out_lst[1].vertices.detach()[0], smpl_model.faces, maintain_orders=True, process=False)
+    smpl_da.export(f"{prefix}_smpl_da.obj")
 
     # remove hands from ECON for next registeration
-    econ_cano_body = econ_cano.copy()
+    econ_da_body = econ_da.copy()
     mano_mask = ~np.isin(idx[:, 0], smplx_container.smplx_mano_vid)
-    econ_cano_body.update_faces(mano_mask[econ_cano.faces].all(axis=1))
-    econ_cano_body.remove_unreferenced_vertices()
-    econ_cano_body = keep_largest(econ_cano_body)
+    econ_da_body.update_faces(mano_mask[econ_da.faces].all(axis=1))
+    econ_da_body.remove_unreferenced_vertices()
+    econ_da_body = keep_largest(econ_da_body)
 
     # remove SMPL-X hand and face
     register_mask = ~np.isin(
-        np.arange(smpl_cano_verts.shape[0]),
+        np.arange(smpl_da.vertices.shape[0]),
         np.concatenate([smplx_container.smplx_mano_vid, smplx_container.smplx_front_flame_vid]))
     register_mask *= ~smplx_container.eyeball_vertex_mask.bool().numpy()
-    smpl_cano_body = smpl_cano.copy()
-    smpl_cano_body.update_faces(register_mask[smpl_cano.faces].all(axis=1))
-    smpl_cano_body.remove_unreferenced_vertices()
-    smpl_cano_body = keep_largest(smpl_cano_body)
-
-    # upsample the smpl_cano_body and do registeration
-    smpl_cano_body = Meshes(
-        verts=[torch.tensor(smpl_cano_body.vertices).float()],
-        faces=[torch.tensor(smpl_cano_body.faces).long()],
+    smpl_da_body = smpl_da.copy()
+    smpl_da_body.update_faces(register_mask[smpl_da.faces].all(axis=1))
+    smpl_da_body.remove_unreferenced_vertices()
+    smpl_da_body = keep_largest(smpl_da_body)
+
+    # upsample the smpl_da_body and do registeration
+    smpl_da_body = Meshes(
+        verts=[torch.tensor(smpl_da_body.vertices).float()],
+        faces=[torch.tensor(smpl_da_body.faces).long()],
     ).to(device)
-    sm = SubdivideMeshes(smpl_cano_body)
-    smpl_cano_body = register(econ_cano_body, sm(smpl_cano_body), device)
+    sm = SubdivideMeshes(smpl_da_body)
+    smpl_da_body = register(econ_da_body, sm(smpl_da_body), device)
 
     # remove over-streched+hand faces from ECON
-    econ_cano_body = econ_cano.copy()
+    econ_da_body = econ_da.copy()
     edge_before = np.sqrt(
         ((econ_obj.vertices[econ_cano.edges[:, 0]] - econ_obj.vertices[econ_cano.edges[:, 1]])**2).sum(axis=1))
-    edge_after = np.sqrt(
-        ((econ_cano.vertices[econ_cano.edges[:, 0]] - econ_cano.vertices[econ_cano.edges[:, 1]])**2).sum(axis=1))
+    edge_after = np.sqrt(((econ_da.vertices[econ_cano.edges[:, 0]] - econ_da.vertices[econ_cano.edges[:, 1]])**2).sum(axis=1))
     edge_diff = edge_after / edge_before.clip(1e-2)
     streched_mask = np.unique(econ_cano.edges[edge_diff > 6])
     mano_mask = ~np.isin(idx[:, 0], smplx_container.smplx_mano_vid)
     mano_mask[streched_mask] = False
-    econ_cano_body.update_faces(mano_mask[econ_cano.faces].all(axis=1))
-    econ_cano_body.remove_unreferenced_vertices()
+    econ_da_body.update_faces(mano_mask[econ_cano.faces].all(axis=1))
+    econ_da_body.remove_unreferenced_vertices()
 
     # stitch the registered SMPL-X body and floating hands to ECON
-    econ_cano_tree = cKDTree(econ_cano.vertices)
-    dist, idx = econ_cano_tree.query(smpl_cano_body.vertices, k=1)
-    smpl_cano_body.update_faces((dist > 0.02)[smpl_cano_body.faces].all(axis=1))
-    smpl_cano_body.remove_unreferenced_vertices()
+    econ_da_tree = cKDTree(econ_da.vertices)
+    dist, idx = econ_da_tree.query(smpl_da_body.vertices, k=1)
+    smpl_da_body.update_faces((dist > 0.02)[smpl_da_body.faces].all(axis=1))
+    smpl_da_body.remove_unreferenced_vertices()
 
-    smpl_hand = smpl_cano.copy()
+    smpl_hand = smpl_da.copy()
     smpl_hand.update_faces(smplx_container.mano_vertex_mask.numpy()[smpl_hand.faces].all(axis=1))
     smpl_hand.remove_unreferenced_vertices()
-    econ_cano = sum([smpl_hand, smpl_cano_body, econ_cano_body])
-    econ_cano = poisson(econ_cano, f"{prefix}_econ_cano.obj")
+    econ_da = sum([smpl_hand, smpl_da_body, econ_da_body])
+    econ_da = poisson(econ_da, f"{prefix}_econ_da.obj")
 else:
-    econ_cano = trimesh.load(f"{prefix}_econ_cano.obj")
-    smpl_cano = trimesh.load(f"{prefix}_smpl_cano.obj", maintain_orders=True, process=False)
+    econ_da = trimesh.load(f"{prefix}_econ_da.obj")
+    smpl_da = trimesh.load(f"{prefix}_smpl_da.obj", maintain_orders=True, process=False)
 
-smpl_tree = cKDTree(smpl_cano.vertices)
-dist, idx = smpl_tree.query(econ_cano.vertices, k=2)
+smpl_tree = cKDTree(smpl_da.vertices)
+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
+
+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()
+
 econ_J_regressor /= econ_J_regressor.sum(axis=1, keepdims=True)
 econ_lbs_weights /= econ_lbs_weights.sum(axis=1, keepdims=True)
 
+# re-compute da-pose rot_mat for ECON
+rot_mat_da = smpl_out_lst[1].vertex_transformation.detach()[0][idx[:, 0]]
+econ_da_verts = torch.tensor(econ_da.vertices).float()
+econ_cano_verts = torch.inverse(rot_mat_da) @ torch.cat([econ_da_verts, torch.ones_like(econ_da_verts)[..., :1]],
+                                                        dim=1).unsqueeze(-1)
+econ_cano_verts = econ_cano_verts[:, :3, 0].double()
+
+# ----------------------------------------------------
+# use any SMPL-X pose to animate ECON reconstruction
+# ----------------------------------------------------
+
+new_pose = smpl_out_lst[2].full_pose
+new_pose[:, :3] = 0.
+
 posed_econ_verts, _ = general_lbs(
-    pose=smpl_out.full_pose,
-    v_template=torch.tensor(econ_cano.vertices).unsqueeze(0),
+    pose=new_pose,
+    v_template=econ_cano_verts.unsqueeze(0),
+    posedirs=econ_posedirs,
     J_regressor=econ_J_regressor,
     parents=smpl_model.parents,
     lbs_weights=econ_lbs_weights)
 
-econ_pose = trimesh.Trimesh(posed_econ_verts[0].detach(), econ_cano.faces)
+econ_pose = trimesh.Trimesh(posed_econ_verts[0].detach(), econ_da.faces)
 econ_pose.export(f"{prefix}_econ_pose.obj")

apps/infer.py

@@ -100,7 +100,7 @@ if __name__ == "__main__":
         print(colored("Use SMPL-X (Explicit) for completion", "green"))
 
     dataset = TestDataset(dataset_param, device)
-
+    
     print(colored(f"Dataset Size: {len(dataset)}", "green"))
 
     pbar = tqdm(dataset)

docs/installation.md

@@ -27,7 +27,7 @@ conda activate econ
 pip install -r requirements.txt
 
 # install libmesh & libvoxelize
-cd lib/commmon/libmesh
+cd lib/common/libmesh
 python setup.py build_ext --inplace
 cd ../libvoxelize
 python setup.py build_ext --inplace

lib/smplx/body_models.py

@@ -1151,6 +1151,7 @@ class SMPLX(SMPLH):
         pose2rot: bool = True,
         return_joint_transformation: bool = False,
         return_vertex_transformation: bool = False,
+        pose_type: str = 'posed',
         **kwargs,
     ) -> SMPLXOutput:
         """
@@ -1240,9 +1241,30 @@ class SMPLX(SMPLH):
             dim=1,
         )
 
+        if pose_type == "t-pose":
+            full_pose *= 0.0
+        elif pose_type == "da-pose":
+            body_pose = torch.zeros_like(body_pose).view(body_pose.shape[0], -1, 3)
+            body_pose[:, 0] = torch.tensor([0., 0., 30 * np.pi / 180.])
+            body_pose[:, 1] = torch.tensor([0., 0., -30 * np.pi / 180.])
+            body_pose = body_pose.view(body_pose.shape[0], -1)
+
+            full_pose = torch.cat(
+                [
+                    global_orient * 0.,
+                    body_pose,
+                    jaw_pose * 0.,
+                    leye_pose * 0.,
+                    reye_pose * 0.,
+                    left_hand_pose * 0.,
+                    right_hand_pose * 0.,
+                ],
+                dim=1,
+            )
+
         # Add the mean pose of the model. Does not affect the body, only the
         # hands when flat_hand_mean == False
-        full_pose += self.pose_mean
+        # full_pose += self.pose_mean
 
         batch_size = max(betas.shape[0], global_orient.shape[0], body_pose.shape[0])
         # Concatenate the shape and expression coefficients

lib/smplx/lbs.py

@@ -233,6 +233,7 @@ def lbs(
 def general_lbs(
     pose: Tensor,
     v_template: Tensor,
+    posedirs: Tensor,
     J_regressor: Tensor,
     parents: Tensor,
     lbs_weights: Tensor,
@@ -246,6 +247,8 @@ def general_lbs(
         The pose parameters in axis-angle format
     v_template torch.tensor BxVx3
         The template mesh that will be deformed
+    posedirs : torch.tensor Px(V * 3)
+        The pose PCA coefficients
     J_regressor : torch.tensor JxV
         The regressor array that is used to calculate the joints from
         the position of the vertices
@@ -277,10 +280,21 @@ def general_lbs(
     # NxJx3 array
     J = vertices2joints(J_regressor, v_template)
 
+    # Add pose blend shapes
+    # N x J x 3 x 3
+    ident = torch.eye(3, dtype=dtype, device=device)
+
     if pose2rot:
         rot_mats = batch_rodrigues(pose.view(-1, 3)).view([batch_size, -1, 3, 3])
+        pose_feature = (rot_mats[:, 1:, :, :] - ident).view([batch_size, -1])
+        # (N x P) x (P, V * 3) -> N x V x 3
+        pose_offsets = torch.matmul(pose_feature, posedirs).view(batch_size, -1, 3)
     else:
         rot_mats = pose.view(batch_size, -1, 3, 3)
+        pose_feature = pose[:, 1:].view(batch_size, -1, 3, 3) - ident
+        pose_offsets = torch.matmul(pose_feature.view(batch_size, -1), posedirs).view(batch_size, -1, 3)
+
+    v_posed = pose_offsets + v_template
 
     # 4. Get the global joint location
     J_transformed, A = batch_rigid_transform(rot_mats, J, parents, dtype=dtype)
@@ -292,13 +306,13 @@ def general_lbs(
     num_joints = J_regressor.shape[0]
     T = torch.matmul(W, A.view(batch_size, num_joints, 16)).view(batch_size, -1, 4, 4)
 
-    homogen_coord = torch.ones([batch_size, v_template.shape[1], 1], dtype=dtype, device=device)
-    v_posed_homo = torch.cat([v_template, homogen_coord], dim=2)
+    homogen_coord = torch.ones([batch_size, v_posed.shape[1], 1], dtype=dtype, device=device)
+    v_posed_homo = torch.cat([v_posed, homogen_coord], dim=2)
     v_homo = torch.matmul(T, torch.unsqueeze(v_posed_homo, dim=-1))
 
     verts = v_homo[:, :, :3, 0]
 
-    return verts, J
+    return verts, J_transformed
 
 
 def vertices2joints(J_regressor: Tensor, vertices: Tensor) -> Tensor: