minor fix

configs/smpler_x/config_ft_3dpw.py

@@ -90,7 +90,7 @@
     dict(type='ImageToTensor', keys=['img', 'ori_img']),
     dict(
         type='Collect',
-        keys=['img', 'sample_idSx'],
+        keys=['img', 'sample_idx'],
         meta_keys=[
             'image_path', 'center', 'scale', 'rotation', 'ori_img',
             'crop_transform'

demo/estimate_smplerx.py

@@ -23,14 +23,23 @@
 )
 from mmhuman3d.utils.ffmpeg_utils import array_to_images
 
-try:
-    from mmdet.apis import inference_detector, init_detector
-    has_mmdet = True
-except (ImportError, ModuleNotFoundError):
-    has_mmdet = False
+# try:
+from mmdet.apis import inference_detector, init_detector
+has_mmdet = True
+# except (ImportError, ModuleNotFoundError):
+#     has_mmdet = False
 
 
 
+
+
+
+import pdb
+
+logger = mmcv.utils.get_logger('mmdet')
+logger.setLevel('WARNING')
+
+
 def get_detection_result(args, frames_iter):
     person_det_model = init_detector(
         args.det_config, args.det_checkpoint, device=args.device.lower())
@@ -81,6 +90,7 @@ def single_person_with_mmdet(args, frames_iter):
         device=args.device.lower())
 
     ckpt = torch.load(args.mesh_reg_checkpoint)
+
     new_state_dict = OrderedDict()
     for k, v in ckpt['network'].items():
         k = k.replace('module.', '').replace('backbone', 'encoder').replace('body_rotation_net', 'body_regressor').replace(
@@ -90,6 +100,8 @@ def single_person_with_mmdet(args, frames_iter):
 
     for i, result in enumerate(mmcv.track_iter_progress(result_list)):
         frame_id = frame_id_list[i]
+        # pdb.set_trace()
+
         mesh_results = inference_image_based_model(
             mesh_model,
             frames_iter[frame_id],
@@ -103,6 +115,7 @@ def single_person_with_mmdet(args, frames_iter):
         smplx_results['right_hand_pose'].append(mesh_results[0]['smplx_rhand_pose'].reshape(-1,3).cpu().numpy())
         smplx_results['jaw_pose'].append(mesh_results[0]['smplx_jaw_pose'].reshape(-1,3).cpu().numpy())
         smplx_results['expression'].append(mesh_results[0]['smplx_expr'].reshape(-1,10).cpu().numpy())
+        smplx_results['transl'].append(mesh_results[0]['cam_trans'].reshape(-1,3).cpu().numpy())
 
         pred_cams.append(mesh_results[0]['cam_trans'].reshape(-1,3).cpu().numpy())
         bboxes_xyxy.append(mesh_results[0]['bbox'])
@@ -136,35 +149,45 @@ def single_person_with_mmdet(args, frames_iter):
         human_data['smplx'] = smplx
         human_data['pred_cams'] = pred_cams
         human_data.dump(osp.join(args.output, 'inference_result.npz'))
+    pdb.set_trace()
+    # if args.show_path is not None:
+    #     frames_folder = osp.join(args.show_path, 'images')
+    #     os.makedirs(frames_folder, exist_ok=True)
+    #     array_to_images(
+    #         np.array(frames_iter)[frame_id_list], output_folder=frames_folder)
+    #     # create body model
+    #     body_model_config = dict(
+    #         type='smplx',
+    #         num_betas=10,
+    #         use_face_contour=True,
+    #         use_pca=False,
+    #         flat_hand_mean=True,
+    #         model_path=args.body_model_dir,
+    #         keypoint_src='smplx',
+    #         keypoint_dst='smplx',
+    #     )
+    #     visualize_smpl_hmr(
+    #         poses=fullpose.reshape(-1, 1, 165),
+    #         betas=smplx_results['betas'],
+    #         cam_transl=pred_cams,
+    #         bbox=bboxes_xyxy,
+    #         output_path=os.path.join(args.show_path, 'smplx.mp4'),
+    #         render_choice=args.render_choice,
+    #         resolution=frames_iter[0].shape[:2],
+    #         origin_frames=frames_folder,
+    #         body_model_config=body_model_config,
+    #         overwrite=True)
+    #     shutil.rmtree(frames_folder)
+
+    # visualize
+
+
+
+
+
+
+
 
-    if args.show_path is not None:
-        frames_folder = osp.join(args.show_path, 'images')
-        os.makedirs(frames_folder, exist_ok=True)
-        array_to_images(
-            np.array(frames_iter)[frame_id_list], output_folder=frames_folder)
-        # create body model
-        body_model_config = dict(
-            type='smplx',
-            num_betas=10,
-            use_face_contour=True,
-            use_pca=False,
-            flat_hand_mean=True,
-            model_path=args.body_model_dir,
-            keypoint_src='smplx',
-            keypoint_dst='smplx',
-        )
-        visualize_smpl_hmr(
-            poses=fullpose.reshape(-1, 1, 165),
-            betas=smplx_results['betas'],
-            cam_transl=pred_cams,
-            bbox=bboxes_xyxy,
-            output_path=os.path.join(args.show_path, 'smplx.mp4'),
-            render_choice=args.render_choice,
-            resolution=frames_iter[0].shape[:2],
-            origin_frames=frames_folder,
-            body_model_config=body_model_config,
-            overwrite=True)
-        shutil.rmtree(frames_folder)
 
 def main(args):
 
@@ -181,12 +204,12 @@ def main(args):
     parser.add_argument(
         '--mesh_reg_config',
         type=str,
-        default='configs/expose/expose.py',
+        default='configs/smpler_x/config_ft_3dpw.py',
         help='Config file for mesh regression')
     parser.add_argument(
         '--mesh_reg_checkpoint',
         type=str,
-        default='data/pretrained_models/expose-d9d5dbf7_20220708.pth',
+        default='../smplerx/checkpoints/smpler_x_h32.pth.tar',
         help='Checkpoint file for mesh regression')
     parser.add_argument(
         '--single_person_demo',
@@ -198,8 +221,7 @@ def main(args):
         help='Config file for detection')
     parser.add_argument(
         '--det_checkpoint',
-        default='https://download.openmmlab.com/mmdetection/v2.0/faster_rcnn/'
-        'faster_rcnn_r50_fpn_1x_coco/'
+        default='data/pretrained_models/'
         'faster_rcnn_r50_fpn_1x_coco_20200130-047c8118.pth',
         help='Checkpoint file for detection')
     parser.add_argument(
@@ -216,17 +238,17 @@ def main(args):
     parser.add_argument(
         '--input_path',
         type=str,
-        default='demo/resources/single_person_demo.mp4',
+        default='../vid_input/ma.mp4',
         help='Input path')
     parser.add_argument(
         '--output',
         type=str,
-        default='demo_result',
+        default='../vid_output',
         help='directory to save output result file')
     parser.add_argument(
         '--show_path',
         type=str,
-        default='demo_result',
+        default='../vid_output',
         help='directory to save rendered images or video')
     parser.add_argument(
         '--render_choice',

mmhuman3d/models/architectures/SMPLer_X.py

@@ -1,13 +1,14 @@
 from abc import ABCMeta
 from typing import Optional, Tuple, Union
+import torchgeometry as tgm
 
 import torch
 import torch.nn as nn
 from torch.nn import functional as F
 from ..heads.smplerx_head import PositionNet, HandRotationNet, FaceRegressor, BoxNet, HandRoI, BodyRotationNet
 
 from mmhuman3d.models.utils.human_models import smpl_x
-from mmhuman3d.utils.transforms import rot6d_to_aa
+from mmhuman3d.utils.transforms import rot6d_to_aa, rotmat_to_aa
 import math
 import mmcv
 import copy
@@ -208,6 +209,8 @@ def forward_test(self,  img: torch.Tensor, img_metas: dict, **kwargs):
         body_img = F.interpolate(img, self.input_body_shape)
 
         # 1. Encoder
+        # import pdb; pdb.set_trace()
+        body_img = body_img.float()
         img_feat, task_tokens = self.encoder(body_img)  # task_token:[bs, N, c]
         shape_token, cam_token, expr_token, jaw_pose_token, hand_token, body_pose_token = \
             task_tokens[:, 0], task_tokens[:, 1], task_tokens[:, 2], task_tokens[:, 3], task_tokens[:, 4:6], task_tokens[:, 6:]
@@ -339,6 +342,8 @@ def rot6d_to_axis_angle(x):
 
     rot_mat = torch.cat([rot_mat, torch.zeros((batch_size, 3, 1)).cuda().float()], 2)  # 3x4 rotation matrix
     axis_angle = tgm.rotation_matrix_to_angle_axis(rot_mat).reshape(-1, 3)  # axis-angle
+    # import pdb; pdb.set_trace()
+    # axis_angle = rotmat_to_aa(rot_mat[...,:3, :3]).reshape(-1, 3)  # axis-angle
     axis_angle[torch.isnan(axis_angle)] = 0.0
     return axis_angle
 

mmhuman3d/models/architectures/builder.py

@@ -8,6 +8,7 @@
 from .hybrik import HybrIK_trainer
 from .mesh_estimator import ImageBodyModelEstimator, VideoBodyModelEstimator
 from .pymafx import PyMAFX
+from .SMPLer_X import SMPLer_X
 
 
 def build_from_cfg(cfg, registry, default_args=None):
@@ -29,6 +30,7 @@ def build_from_cfg(cfg, registry, default_args=None):
 ARCHITECTURES.register_module(
     name='CliffImageBodyModelEstimator', module=CliffImageBodyModelEstimator)
 ARCHITECTURES.register_module(name='PyMAFX', module=PyMAFX)
+ARCHITECTURES.register_module(name='SMPLer_X', module=SMPLer_X)
 
 
 def build_architecture(cfg):

mmhuman3d/models/utils/human_models.py

@@ -7,7 +7,7 @@
 
 class SMPLX(object):
     def __init__(self):
-        self.human_model_path = '/mnt/cache/weichen1/ruisi/mmhuman3d/data/body_models'
+        self.human_model_path = '/home/weichen/mocap_backend/mmhuman3d/data/body_models'
         self.layer_arg = {'create_global_orient': False, 'create_body_pose': False, 'create_left_hand_pose': False, 'create_right_hand_pose': False, 'create_jaw_pose': False, 'create_leye_pose': False, 'create_reye_pose': False, 'create_betas': False, 'create_expression': False, 'create_transl': False}
         self.layer = {'neutral': smplx.create(self.human_model_path, 'smplx', gender='NEUTRAL', use_pca=False, use_face_contour=True, **self.layer_arg),
                         'male': smplx.create(self.human_model_path, 'smplx', gender='MALE', use_pca=False, use_face_contour=True, **self.layer_arg),
@@ -19,6 +19,8 @@ def __init__(self):
         self.expr_code_dim = 10
         with open(osp.join(self.human_model_path, 'smplx', 'SMPLX_to_J14.pkl'), 'rb') as f:
             self.j14_regressor = pickle.load(f, encoding='latin1')
+
+        # self.j14_regressor = np.load(osp.join(self.human_model_path, 'smplx', 'SMPLX_to_J14.npy'))
         with open(osp.join(self.human_model_path, 'smplx', 'MANO_SMPLX_vertex_ids.pkl'), 'rb') as f:
             self.hand_vertex_idx = pickle.load(f, encoding='latin1')
         self.face_vertex_idx = np.load(osp.join(self.human_model_path, 'smplx', 'SMPL-X__FLAME_vertex_ids.npy'))
@@ -153,7 +155,7 @@ def reduce_joint_set(self, joint):
 
 class SMPL(object):
     def __init__(self):
-        self.human_model_path = '/mnt/cache/weichen1/ruisi/mmhuman3d/data/body_models'
+        self.human_model_path =  '/home/weichen/mocap_backend/mmhuman3d/data/body_models'
         self.layer_arg = {'create_body_pose': False, 'create_betas': False, 'create_global_orient': False, 'create_transl': False}
         self.layer = {'neutral': smplx.create(self.human_model_path, 'smpl', gender='NEUTRAL', **self.layer_arg), 'male': smplx.create(self.human_model_path, 'smpl', gender='MALE', **self.layer_arg), 'female': smplx.create(self.human_model_path, 'smpl', gender='FEMALE', **self.layer_arg)}
         self.vertex_num = 6890