Merge pull request #12 from mkocabas/test

add joints3d to the output file

demo.py

@@ -146,7 +146,7 @@ def main(args):
 
         with torch.no_grad():
 
-            pred_cam, pred_verts, pred_pose, pred_betas, norm_joints2d = [], [], [], [], []
+            pred_cam, pred_verts, pred_pose, pred_betas, pred_joints3d, norm_joints2d = [], [], [], [], [], []
 
             for batch in dataloader:
                 if has_keypoints:
@@ -163,12 +163,14 @@ def main(args):
                 pred_verts.append(output['verts'].reshape(batch_size * seqlen, -1, 3))
                 pred_pose.append(output['theta'][:,:,3:75].reshape(batch_size * seqlen, -1))
                 pred_betas.append(output['theta'][:, :,75:].reshape(batch_size * seqlen, -1))
+                pred_joints3d.append(output['kp_3d'].reshape(batch_size * seqlen, -1, 3))
 
 
             pred_cam = torch.cat(pred_cam, dim=0)
             pred_verts = torch.cat(pred_verts, dim=0)
             pred_pose = torch.cat(pred_pose, dim=0)
             pred_betas = torch.cat(pred_betas, dim=0)
+            pred_joints3d = torch.cat(pred_joints3d, dim=0)
 
             del batch
 
@@ -180,7 +182,7 @@ def main(args):
 
             # Run Temporal SMPLify
             update, new_opt_vertices, new_opt_cam, new_opt_pose, new_opt_betas, \
-            new_opt_joint_loss, opt_joint_loss = smplify_runner(
+            new_opt_joints3d, new_opt_joint_loss, opt_joint_loss = smplify_runner(
                 pred_rotmat=pred_pose,
                 pred_betas=pred_betas,
                 pred_cam=pred_cam,
@@ -200,6 +202,7 @@ def main(args):
             pred_cam[update] = new_opt_cam[update]
             pred_pose[update] = new_opt_pose[update]
             pred_betas[update] = new_opt_betas[update]
+            pred_joints3d[update] = new_opt_joints3d[update]
 
         elif args.run_smplify and args.tracking_method == 'bbox':
             print('[WARNING] You need to enable pose tracking to run Temporal SMPLify algorithm!')
@@ -210,6 +213,7 @@ def main(args):
         pred_verts = pred_verts.cpu().numpy()
         pred_pose = pred_pose.cpu().numpy()
         pred_betas = pred_betas.cpu().numpy()
+        pred_joints3d = pred_joints3d.cpu().numpy()
 
         orig_cam = convert_crop_cam_to_orig_img(
             cam=pred_cam,
@@ -224,6 +228,7 @@ def main(args):
             'verts': pred_verts,
             'pose': pred_pose,
             'betas': pred_betas,
+            'joints3d': pred_joints3d,
             'joints2d': joints2d,
             'bboxes': bboxes,
             'frame_ids': frames,

doc/demo.md

@@ -62,7 +62,7 @@ We can inspect what this file contains by:
 ```python
 >>> import joblib # you may use native pickle here as well
 
->>> output = joblib.load('data/demo_results_kinetics_v2/group_dance/vibe_output.pkl') 
+>>> output = joblib.load('output/group_dance/vibe_output.pkl') 
 
 >>> print(output.keys())  
 
@@ -75,11 +75,13 @@ orig_cam (n_frames, 4)      # weak perspective camera parameters in original ima
 verts (n_frames, 6890, 3)   # SMPL mesh vertices
 pose (n_frames, 72)         # SMPL pose parameters
 betas (n_frames, 10)        # SMPL body shape parameters
+joints3d (n_frames, 49, 3)  # SMPL 3D joints
 joints2d (n_frames, 21, 3)  # 2D keypoint detections by STAF if pose tracking enabled otherwise None
 bboxes (n_frames, 4)        # bbox detections (cx,cy,w,h)
 frame_ids (n_frames,)       # frame ids in which subject with tracking id #1 appears
 
 ```
+You can find the names & order of 3d joints [here](https://github.com/mkocabas/VIBE/blob/master/lib/data_utils/kp_utils.py#L212) and 2D joints [here](https://github.com/mkocabas/VIBE/blob/master/lib/data_utils/kp_utils.py#L187).
 
 ## Runtime Performance
 Here is the breakdown of runtime speeds per step namely tracking and VIBE. This results are obtained by running VIBE

lib/utils/demo_utils.py

@@ -156,10 +156,11 @@ def smplify_runner(
     new_opt_cam_t = output['theta'][:,:3]
     new_opt_pose = output['theta'][:,3:75]
     new_opt_betas = output['theta'][:,75:]
+    new_opt_joints3d = output['kp_3d']
 
     return_val = [
         update, new_opt_vertices.cpu(), new_opt_cam_t.cpu(),
-        new_opt_pose.cpu(), new_opt_betas.cpu(),
+        new_opt_pose.cpu(), new_opt_betas.cpu(), new_opt_joints3d.cpu(),
         new_opt_joint_loss, opt_joint_loss,
     ]
 
@@ -275,4 +276,4 @@ def prepare_rendering_results(vibe_results, nframes):
             {list(frame_data.keys())[i]:frame_data[list(frame_data.keys())[i]] for i in sort_idx}
         )
 
-    return frame_results
+    return frame_results