update smpl-h wrapper

bharat-b7 · Mar 16, 2022 · 1a069e4 · 1a069e4
1 parent d22d45e
commit 1a069e4
Show file tree

Hide file tree

Showing 11 changed files with 361 additions and 128 deletions.
diff --git a/docs/prep_smpl.md b/docs/prep_smpl.md
@@ -1,34 +1,55 @@
+[comment]: <> (### Organizing model files)
+
+[comment]: <> (```)
+
+[comment]: <> (|model root)
+
+[comment]: <> (|--SMPL_male.pkl)
+
+[comment]: <> (|--SMPL_female.pkl)
+
+[comment]: <> (|--SMPL_neutral.pkl)
+
+[comment]: <> (|--priors)
+
+[comment]: <> (|----body_prior.pkl )
+
+[comment]: <> (|--regressors)
+
+[comment]: <> (|----body_25_openpose_joints.pkl)
+
+[comment]: <> (```)
+
 ### Organizing SMPL model files
+To run the registration scripts, different files are required depending on whether you use SMPL or SMPL-H model:
 ```
 |model root
-|--SMPL_male.pkl
+|--SMPL_male.pkl # SMPL body male model, required if you use SMPL
 |--SMPL_female.pkl
 |--SMPL_neutral.pkl
-|--priors
-|----body_prior.pkl 
-|--regressors
-|----body_25_openpose_joints.pkl
-```
-
-### Organizing SMPL+H model files
-To run the registration scripts, the following file structures are required:
-```
-|model root
-|--SMPLH_female.pkl # SMPLH female model
+|--SMPLH_female.pkl # SMPLH body models, required if you use SMPL-H  
 |--SMPLH_male.pkl
 |--priors             # folder containing body and hand pose priors
-|----body_prior.pkl 
-|----lh_prior.pkl
+|----body_prior.pkl  # body prior, required both in SMPL and SMPL-H
+|----lh_prior.pkl # hand priors, required if you use SMPL-H model 
 |----rh_prior.pkl
-|--regressors       # folder for body, face, and hand regressors
+|--regressors       # folder for body, face, and hand regressors, required both for SMPL and SMPL-H
 |----body_25_openpose_joints.pkl
 |----face_70_openpose_joints.pkl
 |----hands_42_openpose_joints.pkl
 ```
-#### SMPL+H body models
-Please download the body model files from the [official website](https://mano.is.tue.mpg.de/index.html) and place them to your model root accordingly.
+
+
+#### SMPL or SMPL-H body models
+For SMPL body model files `SMPL_*.pkl`, you can download from the [SMPL website](https://smpl.is.tue.mpg.de/download.php) and rename the model files accordingly. 
+
+
+For SMPL-H body model files `SMPLH_*.pkl`, you can download from the [SMPL-H website](https://mano.is.tue.mpg.de/index.html) and place them to your model root accordingly.
+
 #### Priors
-You can download our prebuilt priors in `assets/priors`. The body prior was built from a subset of [AMASS](https://amass.is.tue.mpg.de/) dataset and hand priors were built from [GRAB](https://grab.is.tue.mpg.de/) dataset. 
+You can download our prebuilt priors in `assets/priors`. The body prior `body_prior.pkl` is required whether you use SMPL or SMPL-H. The hand priors `lh_prior.pkl` and `rh_prior.pkl` are required if you use SMPL-H model. 
+
+The body prior was built from a subset of [AMASS](https://amass.is.tue.mpg.de/) dataset and hand priors were built from [GRAB](https://grab.is.tue.mpg.de/) dataset. 
 
 Alternatively you can build your own priors from another dataset using the script `utils/build_prior.py`:
 ```angular2html
@@ -38,4 +59,4 @@ python utils/build_prior.py data_path out_path
 #### Joint regressors
 You can download these files from `assets/regressors`. 
 
-Once these files are ready, you can change `SMPLH_MODELS_PATH` in `config.yml` accordingly.
+Once these files are ready, you can change `SMPL_MODELS_PATH` in `config.yml` accordingly.
diff --git a/lib/smpl/const.py b/lib/smpl/const.py
@@ -15,5 +15,6 @@
 HAND_POSE_NUM = 90
 TOP_BETA_NUM = 2
 
-
+# split smpl
+SMPL_HAND_POSE_NUM=6
 
diff --git a/lib/smpl/smplpytorch/smplpytorch/pytorch/smpl_layer.py b/lib/smpl/smplpytorch/smplpytorch/pytorch/smpl_layer.py
@@ -21,22 +21,27 @@ def __init__(self,
                  center_idx=None,
                  gender='neutral',
                  model_root='smpl/native/models',
-                 num_betas=300):
+                 num_betas=300,
+                 hands=False):
         """
         Args:
             center_idx: index of center joint in our computations,
             model_root: path to pkl files for the model
-            gender: 'neutral' (default) or 'female' or 'male'
+            gender: 'neutral' (default) or 'female' or 'male', for smplh, only supports male or female
         """
         super().__init__()
 
         self.center_idx = center_idx
         self.gender = gender
 
         self.model_root = model_root
-        self.model_folder = os.path.join(model_root, "models_v1.0.0/models")
-
-        self.model_path = os.path.join(model_root, f"SMPL_{self.gender}.pkl")
+        self.hands = hands
+        if self.hands:
+            assert self.gender in ['male', 'female'], 'SMPL-H model only supports male or female, not {}'.format(self.gender)
+            self.model_path = os.path.join(model_root, f"SMPLH_{self.gender}.pkl")
+        else:
+            self.model_folder = os.path.join(model_root, "models_v1.0.0/models")
+            self.model_path = os.path.join(model_root, f"SMPL_{self.gender}.pkl")
         # if gender == 'neutral':
         #     self.model_path = os.path.join(self.model_folder, 'basicmodel_neutral_lbs_10_207_0_v1.0.0.pkl')
         # elif gender == 'female':
@@ -92,7 +97,7 @@ def forward(self,
         root_rot = th_pose_rotmat[:, :9].view(batch_size, 3, 3)
         # Take out the remaining rotmats (23 joints)
         th_pose_rotmat = th_pose_rotmat[:, 9:]
-        th_pose_map = subtract_flat_id(th_pose_rotmat)
+        th_pose_map = subtract_flat_id(th_pose_rotmat, self.hands)
 
         # Below does: v_shaped = v_template + shapedirs * betas
         # If shape parameters are not provided
@@ -160,15 +165,19 @@ def forward(self,
         th_jtr *= scale
 
         # Shift to new root
-        if self.center_idx is not None:
-            center_joint = th_jtr[:, self.center_idx].unsqueeze(1)
-            th_jtr = th_jtr - center_joint
-            th_verts = th_verts - center_joint
-
-        # If translation is provided
-        if not(th_trans is None or bool(torch.norm(th_trans) == 0)):
-            th_jtr = th_jtr + th_trans.unsqueeze(1)
-            th_verts = th_verts + th_trans.unsqueeze(1)
+        # if self.center_idx is not None:
+        #     center_joint = th_jtr[:, self.center_idx].unsqueeze(1)
+        #     th_jtr = th_jtr - center_joint
+        #     th_verts = th_verts - center_joint
+        #
+        # # If translation is provided
+        # if not(th_trans is None or bool(torch.norm(th_trans) == 0)):
+        #     th_jtr = th_jtr + th_trans.unsqueeze(1)
+        #     th_verts = th_verts + th_trans.unsqueeze(1)
+
+        # XH: not doing shift, apply translation instead
+        th_jtr = th_jtr + th_trans.unsqueeze(1)
+        th_verts = th_verts + th_trans.unsqueeze(1)
 
         # Vertices and joints in meters
         return th_verts, th_jtr, th_v_posed, naked