NICP for SMPL-X completion

README.md

@@ -4,7 +4,7 @@
 
   <h1 align="center">ECON: Explicit Clothed humans Obtained from Normals</h1>
   <p align="center">
-    <a href="https://ps.is.tuebingen.mpg.de/person/yxiu"><strong>Yuliang Xiu</strong></a>
+    <a href="http://xiuyuliang.cn/"><strong>Yuliang Xiu</strong></a>
     ·
     <a href="https://ps.is.tuebingen.mpg.de/person/jyang"><strong>Jinlong Yang</strong></a>
     ·
@@ -28,7 +28,7 @@
       <img src='https://img.shields.io/badge/Paper-PDF (coming soon)-green?style=for-the-badge&logo=arXiv&logoColor=green' alt='Paper PDF'>
     </a>
     <a href='https://xiuyuliang.cn/econ/'>
-      <img src='https://img.shields.io/badge/ECON-Page-orange?style=for-the-badge&logo=Google%20chrome&logoColor=orange' alt='Project Page'></a>
+      <img src='https://img.shields.io/badge/ECON-Page-orange?style=for-the-badge&logo=Google%20chrome&logoColor=white' alt='Project Page'></a>
     <a href="https://discord.gg/Vqa7KBGRyk"><img src="https://img.shields.io/discord/940240966844035082?color=7289DA&labelColor=4a64bd&logo=discord&logoColor=white&style=for-the-badge"></a>
     <a href="https://youtu.be/j5hw4tsWpoY"><img alt="youtube views" title="Subscribe to my YouTube channel" src="https://img.shields.io/youtube/views/j5hw4tsWpoY?logo=youtube&labelColor=ce4630&style=for-the-badge"/></a>
   </p>

apps/avatarizer.py

@@ -3,12 +3,27 @@ import trimesh
 import torch
 import os.path as osp
 import lib.smplx as smplx
+from pytorch3d.ops import SubdivideMeshes
+from pytorch3d.structures import Meshes
+
+from lib.smplx.lbs import general_lbs
+from lib.dataset.mesh_util import keep_largest, poisson
+from scipy.spatial import cKDTree
 from lib.dataset.mesh_util import SMPLX
+from lib.common.local_affine import register
 
 smplx_container = SMPLX()
+device = torch.device("cuda:0")
+
+prefix = "./results/github/econ/obj/304e9c4798a8c3967de7c74c24ef2e38"
+smpl_path = f"{prefix}_smpl_00.npy"
+econ_path = f"{prefix}_0_full.obj"
 
-smpl_npy = "./results/github/econ/obj/304e9c4798a8c3967de7c74c24ef2e38_smpl_00.npy"
-smplx_param = np.load(smpl_npy, allow_pickle=True).item()
+smplx_param = np.load(smpl_path, allow_pickle=True).item()
+econ_obj = trimesh.load(econ_path)
+econ_obj.vertices *= np.array([1.0, -1.0, -1.0])
+econ_obj.vertices /= smplx_param["scale"].cpu().numpy()
+econ_obj.vertices -= smplx_param["transl"].cpu().numpy()
 
 for key in smplx_param.keys():
     smplx_param[key] = smplx_param[key].cpu().view(1, -1)
@@ -28,20 +43,106 @@ smpl_model = smplx.create(
 smpl_out = smpl_model(
     body_pose=smplx_param["body_pose"],
     global_orient=smplx_param["global_orient"],
-    # transl=smplx_param["transl"],
     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]
-inv_mat = torch.inverse(smpl_out.vertex_transformation.detach()[0])
-homo_coord = torch.ones_like(smpl_verts)[..., :1]
-smpl_verts = inv_mat @ torch.cat([smpl_verts, homo_coord], dim=1).unsqueeze(-1)
-smpl_verts = smpl_verts[:, :3, 0].cpu()
+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"):
+
+    # canonicalize for ECON
+    econ_verts = torch.tensor(econ_obj.vertices).float()
+    inv_mat = torch.inverse(smpl_out.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 = 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")
+
+    # remove hands from ECON for next registeration
+    econ_cano_body = econ_cano.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)
+
+    # remove SMPL-X hand and face
+    register_mask = ~np.isin(
+        np.arange(smpl_cano_verts.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()],
+    ).to(device)
+    sm = SubdivideMeshes(smpl_cano_body)
+    smpl_cano_body = register(econ_cano_body, sm(smpl_cano_body), device)
+
+    # remove over-streched+hand faces from ECON
+    econ_cano_body = econ_cano.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_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()
+
+    # 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()
+
+    smpl_hand = smpl_cano.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")
+else:
+    econ_cano = trimesh.load(f"{prefix}_econ_cano.obj")
+    smpl_cano = trimesh.load(f"{prefix}_smpl_cano.obj", maintain_orders=True, process=False)
+
+smpl_tree = cKDTree(smpl_cano.vertices)
+dist, idx = smpl_tree.query(econ_cano.vertices, k=2)
+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 /= econ_J_regressor.sum(axis=1, keepdims=True)
+econ_lbs_weights /= econ_lbs_weights.sum(axis=1, keepdims=True)
+
+posed_econ_verts, _ = general_lbs(
+    pose=smpl_out.full_pose,
+    v_template=torch.tensor(econ_cano.vertices).unsqueeze(0),
+    J_regressor=econ_J_regressor,
+    parents=smpl_model.parents,
+    lbs_weights=econ_lbs_weights)
 
-trimesh.Trimesh(smpl_verts, smpl_model.faces).show()
+econ_pose = trimesh.Trimesh(posed_econ_verts[0].detach(), econ_cano.faces)
+econ_pose.export(f"{prefix}_econ_pose.obj")

apps/infer.py

@@ -37,6 +37,7 @@ from lib.common.train_util import init_loss, load_normal_networks, load_networks
 from lib.common.BNI import BNI
 from lib.common.BNI_utils import save_normal_tensor
 from lib.dataset.TestDataset import TestDataset
+from lib.common.local_affine import register
 from lib.net.geometry import rot6d_to_rotmat, rotation_matrix_to_angle_axis
 from lib.dataset.mesh_util import *
 from lib.common.voxelize import VoxelGrid
@@ -156,8 +157,8 @@ if __name__ == "__main__":
 
         N_body, N_pose = optimed_pose.shape[:2]
 
-        smpl_path = osp.join(args.out_dir, "econ", f"png/{data['name']}_smpl.png")
-
+        smpl_path = f"{args.out_dir}/{cfg.name}/obj/{data['name']}_smpl_00.obj"
+        
         if osp.exists(smpl_path):
 
             smpl_verts_lst = []
@@ -182,6 +183,7 @@ if __name__ == "__main__":
 
             in_tensor["smpl_verts"] = batch_smpl_verts * torch.tensor([1., -1., 1.]).to(device)
             in_tensor["smpl_faces"] = batch_smpl_faces[:, :, [0, 2, 1]]
+            
         else:
             # smpl optimization
             loop_smpl = tqdm(range(args.loop_smpl))
@@ -447,15 +449,15 @@ if __name__ == "__main__":
                     (SMPLX_object.front_flame_vertex_mask + SMPLX_object.mano_vertex_mask +
                      SMPLX_object.eyeball_vertex_mask).eq(0).float(),
                 )
-
-                # upsample the side mesh
-                side_sub_mesh = Meshes(
+                
+                #register side_mesh to BNI surfaces
+                side_mesh = Meshes(
                     verts=[torch.tensor(side_mesh.vertices).float()],
                     faces=[torch.tensor(side_mesh.faces).long()],
-                )
-                sm = SubdivideMeshes(side_sub_mesh)
-                new_mesh = sm(side_sub_mesh)
-                side_mesh = trimesh.Trimesh(new_mesh.verts_padded().squeeze(), new_mesh.faces_padded().squeeze())
+                ).to(device)
+                sm = SubdivideMeshes(side_mesh)
+                side_mesh = register(BNI_object.F_B_trimesh, sm(side_mesh), device)
+
 
             side_verts = torch.tensor(side_mesh.vertices).float().to(device)
             side_faces = torch.tensor(side_mesh.faces).long().to(device)

lib/common/local_affine.py

@@ -0,0 +1,136 @@
+# Copyright 2021 by Haozhe Wu, Tsinghua University, Department of Computer Science and Technology.
+# All rights reserved.
+# This file is part of the pytorch-nicp,
+# and is released under the "MIT License Agreement". Please see the LICENSE
+# file that should have been included as part of this package.
+
+import torch
+import trimesh
+import torch.nn as nn
+from tqdm import tqdm
+from pytorch3d.structures import Meshes
+from pytorch3d.loss import chamfer_distance
+from lib.dataset.mesh_util import update_mesh_shape_prior_losses
+from lib.common.train_util import init_loss
+
+
+# reference: https://github.com/wuhaozhe/pytorch-nicp
+class LocalAffine(nn.Module):
+
+    def __init__(self, num_points, batch_size=1, edges=None):
+        '''
+            specify the number of points, the number of points should be constant across the batch
+            and the edges torch.Longtensor() with shape N * 2
+            the local affine operator supports batch operation
+            batch size must be constant
+            add additional pooling on top of w matrix
+        '''
+        super(LocalAffine, self).__init__()
+        self.A = nn.Parameter(torch.eye(3).unsqueeze(0).unsqueeze(0).repeat(batch_size, num_points, 1, 1))
+        self.b = nn.Parameter(torch.zeros(3).unsqueeze(0).unsqueeze(0).unsqueeze(3).repeat(batch_size, num_points, 1, 1))
+        self.edges = edges
+        self.num_points = num_points
+
+    def stiffness(self):
+        '''
+            calculate the stiffness of local affine transformation
+            f norm get infinity gradient when w is zero matrix, 
+        '''
+        if self.edges is None:
+            raise Exception("edges cannot be none when calculate stiff")
+        idx1 = self.edges[:, 0]
+        idx2 = self.edges[:, 1]
+        affine_weight = torch.cat((self.A, self.b), dim=3)
+        w1 = torch.index_select(affine_weight, dim=1, index=idx1)
+        w2 = torch.index_select(affine_weight, dim=1, index=idx2)
+        w_diff = (w1 - w2)**2
+        w_rigid = (torch.linalg.det(self.A) - 1.0)**2
+        return w_diff, w_rigid
+
+    def forward(self, x):
+        '''
+            x should have shape of B * N * 3
+        '''
+        x = x.unsqueeze(3)
+        out_x = torch.matmul(self.A, x)
+        out_x = out_x + self.b
+        stiffness, rigid = self.stiffness()
+        out_x.squeeze_(3)
+        return out_x, stiffness, rigid
+
+
+def trimesh2meshes(mesh):
+    '''
+        convert trimesh mesh to pytorch3d mesh
+    '''
+    verts = torch.from_numpy(mesh.vertices).float()
+    faces = torch.from_numpy(mesh.faces).long()
+    mesh = Meshes(verts.unsqueeze(0), faces.unsqueeze(0))
+    return mesh
+
+
+def register(target_mesh, src_mesh, device):
+
+    # define local_affine deform verts
+    tgt_mesh = trimesh2meshes(target_mesh).to(device)
+    src_verts = src_mesh.verts_padded().clone()
+
+    local_affine_model = LocalAffine(src_mesh.verts_padded().shape[1],
+                                     src_mesh.verts_padded().shape[0], src_mesh.edges_packed()).to(device)
+
+    optimizer_cloth = torch.optim.Adam([{'params': local_affine_model.parameters()}], lr=1e-2, amsgrad=True)
+    scheduler_cloth = torch.optim.lr_scheduler.ReduceLROnPlateau(
+        optimizer_cloth,
+        mode="min",
+        factor=0.1,
+        verbose=0,
+        min_lr=1e-5,
+        patience=5,
+    )
+
+    losses = init_loss()
+
+    loop_cloth = tqdm(range(200))
+
+    for i in loop_cloth:
+
+        optimizer_cloth.zero_grad()
+
+        deformed_verts, stiffness, rigid = local_affine_model(src_verts)
+        src_mesh = src_mesh.update_padded(deformed_verts)
+
+        # losses for laplacian, edge, normal consistency
+        update_mesh_shape_prior_losses(src_mesh, losses)
+
+        losses["cloth"]["value"] = chamfer_distance(
+            x=src_mesh.verts_padded(),
+            y=tgt_mesh.verts_padded())[0]
+        
+        losses["stiffness"]["value"] = torch.mean(stiffness)
+        losses["rigid"]["value"] = torch.mean(rigid)
+
+        # Weighted sum of the losses
+        cloth_loss = torch.tensor(0.0, requires_grad=True).to(device)
+        pbar_desc = "Register SMPL-X towards ECON --- "
+
+        for k in losses.keys():
+            if losses[k]["weight"] > 0.0 and losses[k]["value"] != 0.0:
+                cloth_loss = cloth_loss + \
+                    losses[k]["value"] * losses[k]["weight"]
+                pbar_desc += f"{k}:{losses[k]['value']* losses[k]['weight']:.3f} | "
+
+        pbar_desc += f"Total: {cloth_loss:.5f}"
+        loop_cloth.set_description(pbar_desc)
+
+        # update params
+        cloth_loss.backward(retain_graph=True)
+        optimizer_cloth.step()
+        scheduler_cloth.step(cloth_loss)
+
+    final = trimesh.Trimesh(
+        src_mesh.verts_packed().detach().squeeze(0).cpu(),
+        src_mesh.faces_packed().detach().squeeze(0).cpu(),
+        process=False,
+        maintains_order=True)
+
+    return final

lib/common/train_util.py

@@ -32,7 +32,7 @@ def init_loss():
     losses = {
         # Cloth: Normal_recon - Normal_pred
         "cloth": {
-            "weight": 1e1,
+            "weight": 1e3,
             "value": 0.0
         },
         # Cloth: [RT]_v1 - [RT]_v2 (v1-edge-v2)

lib/dataset/mesh_util.py

@@ -552,7 +552,7 @@ def poisson_remesh(obj_path):
     ms.meshing_decimation_quadric_edge_collapse(targetfacenum=50000)
     # ms.apply_coord_laplacian_smoothing()
     ms.save_current_mesh(obj_path)
-    ms.save_current_mesh(obj_path.replace(".obj", ".ply"))
+    # ms.save_current_mesh(obj_path.replace(".obj", ".ply"))
     polished_mesh = trimesh.load_mesh(obj_path)
 
     return polished_mesh
@@ -1013,6 +1013,15 @@ def clean_floats(mesh):
     return sum(clean_mesh_lst)
 
 
+def keep_largest(mesh):
+    mesh_lst = mesh.split(only_watertight=False)
+    keep_mesh = mesh_lst[0]
+    for mesh in mesh_lst:
+        if mesh.vertices.shape[0] > keep_mesh.vertices.shape[0]:
+            keep_mesh = mesh
+    return keep_mesh
+
+
 def mesh_move(mesh_lst, step, scale=1.0):
 
     trans = np.array([1.0, 0.0, 0.0]) * step
@@ -1036,3 +1045,16 @@ def rescale_smpl(fitted_path, scale=100, translate=(0, 0, 0)):
     fitted_body.apply_transform(resize_matrix)
 
     return np.array(fitted_body.vertices)
+
+
+def get_joint_mesh(joints, radius=2.0):
+
+    ball = trimesh.creation.icosphere(radius=radius)
+    combined = None
+    for joint in joints:
+        ball_new = trimesh.Trimesh(vertices=ball.vertices + joint, faces=ball.faces, process=False)
+        if combined is None:
+            combined = ball_new
+        else:
+            combined = sum([combined, ball_new])
+    return combined

lib/smplx/lbs.py

@@ -194,6 +194,7 @@ def lbs(
     # 3. 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])
 
@@ -229,6 +230,77 @@ def lbs(
     return verts, J_transformed
 
 
+def general_lbs(
+    pose: Tensor,
+    v_template: Tensor,
+    J_regressor: Tensor,
+    parents: Tensor,
+    lbs_weights: Tensor,
+    pose2rot: bool = True,
+) -> Tuple[Tensor, Tensor, Optional[Tensor], Optional[Tensor]]:
+    """Performs Linear Blend Skinning with the given shape and pose parameters
+
+    Parameters
+    ----------
+    pose : torch.tensor Bx(J + 1) * 3
+        The pose parameters in axis-angle format
+    v_template torch.tensor BxVx3
+        The template mesh that will be deformed
+    J_regressor : torch.tensor JxV
+        The regressor array that is used to calculate the joints from
+        the position of the vertices
+    parents: torch.tensor J
+        The array that describes the kinematic tree for the model
+    lbs_weights: torch.tensor N x V x (J + 1)
+        The linear blend skinning weights that represent how much the
+        rotation matrix of each part affects each vertex
+    pose2rot: bool, optional
+        Flag on whether to convert the input pose tensor to rotation
+        matrices. The default value is True. If False, then the pose tensor
+        should already contain rotation matrices and have a size of
+        Bx(J + 1)x9
+    dtype: torch.dtype, optional
+
+    Returns
+    -------
+    verts: torch.tensor BxVx3
+        The vertices of the mesh after applying the shape and pose
+        displacements.
+    joints: torch.tensor BxJx3
+        The joints of the model
+    """
+
+    batch_size = pose.shape[0]
+    device, dtype = pose.device, pose.dtype
+
+    # Get the joints
+    # NxJx3 array
+    J = vertices2joints(J_regressor, v_template)
+
+    if pose2rot:
+        rot_mats = batch_rodrigues(pose.view(-1, 3)).view([batch_size, -1, 3, 3])
+    else:
+        rot_mats = pose.view(batch_size, -1, 3, 3)
+
+    # 4. Get the global joint location
+    J_transformed, A = batch_rigid_transform(rot_mats, J, parents, dtype=dtype)
+
+    # 5. Do skinning:
+    # W is N x V x (J + 1)
+    W = lbs_weights.unsqueeze(dim=0).expand([batch_size, -1, -1])
+    # (N x V x (J + 1)) x (N x (J + 1) x 16)
+    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)
+    v_homo = torch.matmul(T, torch.unsqueeze(v_posed_homo, dim=-1))
+
+    verts = v_homo[:, :, :3, 0]
+
+    return verts, J
+
+
 def vertices2joints(J_regressor: Tensor, vertices: Tensor) -> Tensor:
     """Calculates the 3D joint locations from the vertices