update

app.py

@@ -1,136 +0,0 @@
-import os
-import sys
-import os.path as osp
-from pathlib import Path
-import cv2
-import gradio as gr
-import torch
-import math
-import spaces
-from huggingface_hub import hf_hub_download
-try:
-    import mmpose
-except:
-    os.system('pip install /home/user/app/main/transformer_utils')
-hf_hub_download(repo_id="caizhongang/SMPLer-X", filename="smpler_x_h32.pth.tar", local_dir="/home/user/app/pretrained_models")
-os.system('cp -rf /home/user/app/assets/conversions.py /usr/local/lib/python3.10/site-packages/torchgeometry/core/conversions.py')
-DEFAULT_MODEL='smpler_x_h32'
-OUT_FOLDER = '/home/user/app/demo_out'
-os.makedirs(OUT_FOLDER, exist_ok=True)
-num_gpus = 1 if torch.cuda.is_available() else -1
-print("!!!", torch.cuda.is_available())      
-print(torch.cuda.device_count())   
-print(torch.version.cuda)  
-index = torch.cuda.current_device()
-print(index)  
-print(torch.cuda.get_device_name(index))
-# from main.inference import Inferer
-# inferer = Inferer(DEFAULT_MODEL, num_gpus, OUT_FOLDER)
-
-@spaces.GPU(enable_queue=True, duration=300)
-def infer(video_input, in_threshold=0.5, num_people="Single person", render_mesh=False):
-    from main.inference import Inferer
-    inferer = Inferer(DEFAULT_MODEL, num_gpus, OUT_FOLDER)
-    os.system(f'rm -rf {OUT_FOLDER}/*')
-    multi_person = False if (num_people == "Single person") else True
-    cap = cv2.VideoCapture(video_input)
-    fps = math.ceil(cap.get(5))
-    width = int(cap.get(3))
-    height = int(cap.get(4))
-    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-    video_path = osp.join(OUT_FOLDER, f'out.m4v')
-    final_video_path = osp.join(OUT_FOLDER, f'out.mp4')
-    video_output = cv2.VideoWriter(video_path, fourcc, fps, (width, height))
-    success = 1
-    frame = 0
-    while success:
-        success, original_img = cap.read()
-        if not success:
-            break
-        frame += 1
-        img, mesh_paths, smplx_paths = inferer.infer(original_img, in_threshold, frame, multi_person, not(render_mesh))
-        video_output.write(img)
-        yield img, None, None, None
-    cap.release()
-    video_output.release()
-    cv2.destroyAllWindows()
-    os.system(f'ffmpeg -i {video_path} -c copy {final_video_path}')
-
-    #Compress mesh and smplx files
-    save_path_mesh = os.path.join(OUT_FOLDER, 'mesh')
-    save_mesh_file = os.path.join(OUT_FOLDER, 'mesh.zip')
-    os.makedirs(save_path_mesh, exist_ok= True)
-    save_path_smplx = os.path.join(OUT_FOLDER, 'smplx')
-    save_smplx_file = os.path.join(OUT_FOLDER, 'smplx.zip')
-    os.makedirs(save_path_smplx, exist_ok= True)
-    os.system(f'zip -r {save_mesh_file} {save_path_mesh}')
-    os.system(f'zip -r {save_smplx_file} {save_path_smplx}')
-    yield img, video_path, save_mesh_file, save_smplx_file
-
-TITLE = '''<h1 align="center">SMPLer-X: Scaling Up Expressive Human Pose and Shape Estimation</h1>'''
-VIDEO = '''
-<center><iframe width="960" height="540" 
-src="https://www.youtube.com/embed/DepTqbPpVzY?si=qSeQuX-bgm_rON7E"title="SMPLer-X: Scaling Up Expressive Human Pose and Shape Estimation" frameborder="0" allow="accelerometer; autoplay; clipboard-write; encrypted-media; gyroscope; picture-in-picture; web-share" allowfullscreen>
-</iframe>
-</center><br>'''
-DESCRIPTION = '''
-<b>Official Gradio demo</b> for <a href="https://caizhongang.com/projects/SMPLer-X/"><b>SMPLer-X: Scaling Up Expressive Human Pose and Shape Estimation</b></a>.<br>
-<p>
-Note: You can drop a video at the panel (or select one of the examples) 
-to obtain the 3D parametric reconstructions of the detected humans.
-</p>
-'''
-
-with gr.Blocks(title="SMPLer-X", css=".gradio-container") as demo:
-
-    gr.Markdown(TITLE)
-    gr.HTML(VIDEO)
-    gr.Markdown(DESCRIPTION)
-
-    with gr.Row():
-        with gr.Column():
-            video_input = gr.Video(label="Input video", elem_classes="video")
-            threshold = gr.Slider(0, 1.0, value=0.5, label='BBox detection threshold')
-        with gr.Column(scale=2):
-            num_people = gr.Radio(
-                choices=["Single person", "Multiple people"],
-                value="Single person",
-                label="Number of people",
-                info="Choose how many people are there in the video. Choose 'single person' for faster inference.",
-                interactive=True,
-                scale=1,)
-            gr.HTML("""<br/>""")
-            mesh_as_vertices = gr.Checkbox(
-                label="Render as mesh", 
-                info="By default, the estimated SMPL-X parameters are rendered as vertices for faster visualization. Check this option if you want to visualize meshes instead.",
-                interactive=True,
-                scale=1,)
-
-            send_button = gr.Button("Infer")
-    gr.HTML("""<br/>""")
-
-    with gr.Row():
-        with gr.Column():
-            processed_frames = gr.Image(label="Last processed frame")
-            video_output = gr.Video(elem_classes="video")
-        with gr.Column():
-            meshes_output = gr.File(label="3D meshes")
-            smplx_output = gr.File(label= "SMPL-X models")
-    # example_images = gr.Examples([])
-    send_button.click(fn=infer, inputs=[video_input, threshold, num_people, mesh_as_vertices], outputs=[processed_frames, video_output, meshes_output, smplx_output])
-    # with gr.Row():
-    example_videos = gr.Examples([
-        ['/home/user/app/assets/01.mp4'], 
-        ['/home/user/app/assets/02.mp4'], 
-        ['/home/user/app/assets/03.mp4'],
-        ['/home/user/app/assets/04.mp4'],
-        ['/home/user/app/assets/05.mp4'], 
-        ['/home/user/app/assets/06.mp4'], 
-        ['/home/user/app/assets/07.mp4'],
-        ['/home/user/app/assets/08.mp4'],
-        ['/home/user/app/assets/09.mp4'],
-        ], 
-        inputs=[video_input, 0.5])
-
-#demo.queue()
-demo.queue().launch(debug=True)

common/utils/smplx/smplx/body_models.py

@@ -138,8 +138,8 @@ class SMPL(nn.Module):
         self.batch_size = batch_size
         shapedirs = data_struct.shapedirs
         if (shapedirs.shape[-1] < self.SHAPE_SPACE_DIM):
-            print(f'WARNING: You are using a {self.name()} model, with only'
-                  ' 10 shape coefficients.')
+            # print(f'WARNING: You are using a {self.name()} model, with only'
+            #       ' 10 shape coefficients.')
             num_betas = min(num_betas, 10)
         else:
             num_betas = min(num_betas, self.SHAPE_SPACE_DIM)
@@ -979,8 +979,8 @@ class SMPLX(SMPLH):
             shapedirs = shapedirs[:, :, None]
         if (shapedirs.shape[-1] < self.SHAPE_SPACE_DIM +
                 self.EXPRESSION_SPACE_DIM):
-            print(f'WARNING: You are using a {self.name()} model, with only'
-                  ' 10 shape and 10 expression coefficients.')
+            # print(f'WARNING: You are using a {self.name()} model, with only'
+            #       ' 10 shape and 10 expression coefficients.')
             expr_start_idx = 10
             expr_end_idx = 20
             num_expression_coeffs = min(num_expression_coeffs, 10)
@@ -1823,8 +1823,8 @@ class FLAME(SMPL):
             shapedirs = shapedirs[:, :, None]
         if (shapedirs.shape[-1] < self.SHAPE_SPACE_DIM +
                 self.EXPRESSION_SPACE_DIM):
-            print(f'WARNING: You are using a {self.name()} model, with only'
-                  ' 10 shape and 10 expression coefficients.')
+            # print(f'WARNING: You are using a {self.name()} model, with only'
+            #       ' 10 shape and 10 expression coefficients.')
             expr_start_idx = 10
             expr_end_idx = 20
             num_expression_coeffs = min(num_expression_coeffs, 10)

common/utils/transforms.py

@@ -3,7 +3,6 @@ import numpy as np
 import scipy
 from config import cfg
 from torch.nn import functional as F
-import torchgeometry as tgm
 
 
 def cam2pixel(cam_coord, f, c):
@@ -69,6 +68,105 @@ def transform_joint_to_other_db(src_joint, src_name, dst_name):
     return new_joint
 
 
+def rotation_matrix_to_angle_axis(rotation_matrix):
+    # todo add check that matrix is a valid rotation matrix
+    quaternion = rotation_matrix_to_quaternion(rotation_matrix)
+    return quaternion_to_angle_axis(quaternion)
+
+
+def rotation_matrix_to_quaternion(rotation_matrix, eps=1e-6):
+    if not torch.is_tensor(rotation_matrix):
+        raise TypeError("Input type is not a torch.Tensor. Got {}".format(
+            type(rotation_matrix)))
+
+    if len(rotation_matrix.shape) > 3:
+        raise ValueError(
+            "Input size must be a three dimensional tensor. Got {}".format(
+                rotation_matrix.shape))
+    if not rotation_matrix.shape[-2:] == (3, 4):
+        raise ValueError(
+            "Input size must be a N x 3 x 4  tensor. Got {}".format(
+                rotation_matrix.shape))
+
+    rmat_t = torch.transpose(rotation_matrix, 1, 2)
+
+    mask_d2 = rmat_t[:, 2, 2] < eps
+
+    mask_d0_d1 = rmat_t[:, 0, 0] > rmat_t[:, 1, 1]
+    mask_d0_nd1 = rmat_t[:, 0, 0] < -rmat_t[:, 1, 1]
+
+    t0 = 1 + rmat_t[:, 0, 0] - rmat_t[:, 1, 1] - rmat_t[:, 2, 2]
+    q0 = torch.stack([rmat_t[:, 1, 2] - rmat_t[:, 2, 1],
+                      t0, rmat_t[:, 0, 1] + rmat_t[:, 1, 0],
+                      rmat_t[:, 2, 0] + rmat_t[:, 0, 2]], -1)
+    t0_rep = t0.repeat(4, 1).t()
+
+    t1 = 1 - rmat_t[:, 0, 0] + rmat_t[:, 1, 1] - rmat_t[:, 2, 2]
+    q1 = torch.stack([rmat_t[:, 2, 0] - rmat_t[:, 0, 2],
+                      rmat_t[:, 0, 1] + rmat_t[:, 1, 0],
+                      t1, rmat_t[:, 1, 2] + rmat_t[:, 2, 1]], -1)
+    t1_rep = t1.repeat(4, 1).t()
+
+    t2 = 1 - rmat_t[:, 0, 0] - rmat_t[:, 1, 1] + rmat_t[:, 2, 2]
+    q2 = torch.stack([rmat_t[:, 0, 1] - rmat_t[:, 1, 0],
+                      rmat_t[:, 2, 0] + rmat_t[:, 0, 2],
+                      rmat_t[:, 1, 2] + rmat_t[:, 2, 1], t2], -1)
+    t2_rep = t2.repeat(4, 1).t()
+
+    t3 = 1 + rmat_t[:, 0, 0] + rmat_t[:, 1, 1] + rmat_t[:, 2, 2]
+    q3 = torch.stack([t3, rmat_t[:, 1, 2] - rmat_t[:, 2, 1],
+                      rmat_t[:, 2, 0] - rmat_t[:, 0, 2],
+                      rmat_t[:, 0, 1] - rmat_t[:, 1, 0]], -1)
+    t3_rep = t3.repeat(4, 1).t()
+
+    mask_c0 = mask_d2 * mask_d0_d1
+    mask_c1 = mask_d2 * (~mask_d0_d1)
+    mask_c2 = (~mask_d2) * mask_d0_nd1
+    mask_c3 = (~mask_d2) * (~mask_d0_nd1)
+    mask_c0 = mask_c0.view(-1, 1).type_as(q0)
+    mask_c1 = mask_c1.view(-1, 1).type_as(q1)
+    mask_c2 = mask_c2.view(-1, 1).type_as(q2)
+    mask_c3 = mask_c3.view(-1, 1).type_as(q3)
+
+    q = q0 * mask_c0 + q1 * mask_c1 + q2 * mask_c2 + q3 * mask_c3
+    q /= torch.sqrt(t0_rep * mask_c0 + t1_rep * mask_c1 +  # noqa
+                    t2_rep * mask_c2 + t3_rep * mask_c3)  # noqa
+    q *= 0.5
+    return q
+
+
+def quaternion_to_angle_axis(quaternion: torch.Tensor) -> torch.Tensor:
+    if not torch.is_tensor(quaternion):
+        raise TypeError("Input type is not a torch.Tensor. Got {}".format(
+            type(quaternion)))
+
+    if not quaternion.shape[-1] == 4:
+        raise ValueError("Input must be a tensor of shape Nx4 or 4. Got {}"
+                         .format(quaternion.shape))
+    # unpack input and compute conversion
+    q1: torch.Tensor = quaternion[..., 1]
+    q2: torch.Tensor = quaternion[..., 2]
+    q3: torch.Tensor = quaternion[..., 3]
+    sin_squared_theta: torch.Tensor = q1 * q1 + q2 * q2 + q3 * q3
+
+    sin_theta: torch.Tensor = torch.sqrt(sin_squared_theta)
+    cos_theta: torch.Tensor = quaternion[..., 0]
+    two_theta: torch.Tensor = 2.0 * torch.where(
+        cos_theta < 0.0,
+        torch.atan2(-sin_theta, -cos_theta),
+        torch.atan2(sin_theta, cos_theta))
+
+    k_pos: torch.Tensor = two_theta / sin_theta
+    k_neg: torch.Tensor = 2.0 * torch.ones_like(sin_theta)
+    k: torch.Tensor = torch.where(sin_squared_theta > 0.0, k_pos, k_neg)
+
+    angle_axis: torch.Tensor = torch.zeros_like(quaternion)[..., :3]
+    angle_axis[..., 0] += q1 * k
+    angle_axis[..., 1] += q2 * k
+    angle_axis[..., 2] += q3 * k
+    return angle_axis
+
+
 def rot6d_to_axis_angle(x):
     batch_size = x.shape[0]
 
@@ -81,7 +179,7 @@ def rot6d_to_axis_angle(x):
     rot_mat = torch.stack((b1, b2, b3), dim=-1)  # 3x3 rotation matrix
 
     rot_mat = torch.cat([rot_mat, torch.zeros((batch_size, 3, 1)).to(cfg.device).float()], 2)  # 3x4 rotation matrix
-    axis_angle = tgm.rotation_matrix_to_angle_axis(rot_mat).reshape(-1, 3)  # axis-angle
+    axis_angle = rotation_matrix_to_angle_axis(rot_mat).reshape(-1, 3)  # axis-angle
     axis_angle[torch.isnan(axis_angle)] = 0.0
     return axis_angle
 

demo.py

@@ -1,74 +0,0 @@
-import os
-import sys
-import os.path as osp
-import argparse
-from pathlib import Path
-import cv2
-import torch
-import math
-import mmpose
-import shutil
-import time
-from OpenGL import GL
-from OpenGL.GL import *
-os.environ["PYOPENGL_PLATFORM"] = "osmesa"
-import pyrender
-try:
-    import mmpose
-except:
-    os.system('pip install main/transformer_utils')
-
-def parse_args():
-    parser = argparse.ArgumentParser()
-    parser.add_argument('--show_verts', action="store_true")
-    parser.add_argument('--multi_person', action="store_true")
-    parser.add_argument('--in_threshold', type=float, default=0.5)
-    parser.add_argument('--output_folder', type=str, default='demo_out')
-    parser.add_argument('--pretrained_model', type=str, default='smpler_x_h32')
-    parser.add_argument('--input_video', type=str, default='')
-    args = parser.parse_args()
-    return args
-
-def infer():
-    args = parse_args()
-    os.makedirs(args.output_folder, exist_ok=True)
-    num_gpus = 1 if torch.cuda.is_available() else -1
-
-    from main.inference import Inferer
-    inferer = Inferer(args.pretrained_model, num_gpus, args.output_folder)
-
-    cap = cv2.VideoCapture(args.input_video)
-    fps = math.ceil(cap.get(5))
-    width = int(cap.get(3))
-    height = int(cap.get(4))
-    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-    video_path = osp.join(args.output_folder, f'out.m4v')
-    final_video_path = osp.join(args.output_folder, f'out.mp4')
-    video_output = cv2.VideoWriter(video_path, fourcc, fps, (width, height))
-    success = 1
-    frame = 0
-    while success:
-        success, original_img = cap.read()
-        if not success:
-            break
-        frame += 1
-        img, mesh_paths, smplx_paths = inferer.infer(original_img, args.in_threshold, frame, args.multi_person, args.show_verts)
-        video_output.write(img)
-    cap.release()
-    video_output.release()
-    cv2.destroyAllWindows()
-    os.system(f'ffmpeg -i {video_path} -c copy {final_video_path}')
-
-    #Compress mesh and smplx files
-    save_path_mesh = os.path.join(args.output_folder, 'mesh')
-    save_mesh_file = os.path.join(args.output_folder, 'mesh.zip')
-    os.makedirs(save_path_mesh, exist_ok= True)
-    save_path_smplx = os.path.join(args.output_folder, 'smplx')
-    save_smplx_file = os.path.join(args.output_folder, 'smplx.zip')
-    os.makedirs(save_path_smplx, exist_ok= True)
-    os.system(f'zip -r {save_mesh_file} {save_path_mesh}')
-    os.system(f'zip -r {save_smplx_file} {save_path_smplx}')
-
-if __name__ == "__main__":
-    main()
-

demo_out/test_logs.txt

@@ -0,0 +1,16 @@
+05-28 18:40:45 Load checkpoint from /home/weichen/SMPLer-X/main/../pretrained_models/smpler_x_h32.pth.tar
+05-28 18:40:45 Creating graph...
+05-28 18:41:36 Load checkpoint from /home/weichen/SMPLer-X/main/../pretrained_models/smpler_x_h32.pth.tar
+05-28 18:41:36 Creating graph...
+05-28 18:44:29 Load checkpoint from /home/weichen/SMPLer-X/main/../pretrained_models/smpler_x_h32.pth.tar
+05-28 18:44:29 Creating graph...
+05-28 18:44:41 Load checkpoint from /home/weichen/SMPLer-X/main/../pretrained_models/smpler_x_h32.pth.tar
+05-28 18:44:41 Creating graph...
+05-29 13:30:00 Load checkpoint from /home/weichen/SMPLer-X/main/../pretrained_models/smpler_x_h32.pth.tar
+05-29 13:30:00 Creating graph...
+05-29 13:32:26 Load checkpoint from /home/weichen/SMPLer-X/main/../pretrained_models/smpler_x_h32.pth.tar
+05-29 13:32:26 Creating graph...
+05-29 13:37:15 Load checkpoint from /home/weichen/SMPLer-X/main/../pretrained_models/smpler_x_h32.pth.tar
+05-29 13:37:15 Creating graph...
+05-29 13:58:51 Load checkpoint from /home/weichen/SMPLer-X/main/../pretrained_models/smpler_x_h32.pth.tar
+05-29 13:58:51 Creating graph...

inf_demo.py

@@ -0,0 +1,67 @@
+import os
+import argparse
+import cv2
+import readline
+import tqdm
+import time
+
+import torch
+from torch.utils.data import DataLoader
+from PIL import Image
+
+# from smplerx.main.inference import Inferer
+from main.inference import SmplerxData, Inferer
+
+
+import pdb
+
+
+def inference(args):
+
+    # load model
+    num_gpus = 1 if torch.cuda.is_available() else -1
+    inferer = Inferer(args.pretrained_model, num_gpus)
+
+    # test annotations
+    annotations = [
+        {'image_path': '/home/weichen/wc_workspace/laoyouji/frames/Season_1/S01E01/frame025000.jpg', 'bbox': [0, 0, 1000, 1000]},
+        {'image_path': '/home/weichen/wc_workspace/laoyouji/frames/Season_1/S01E01/frame026000.jpg', 'bbox': [0, 0, 1000, 1000]}
+    ]
+
+    annotations = annotations*500
+    anno_len = len(annotations)
+    
+    start_time = time.time()
+    batch_size = 1
+    
+    dataset = SmplerxData(annotations=annotations)
+    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=4)
+
+    preproces_time = time.time()
+
+    for batch in tqdm.tqdm(dataloader):
+
+        smplx_pred, meta, mesh = inferer.batch_infer_given_bbox(batch['image'], batch['bbox'])
+
+    end_time = time.time()
+
+    # print report, time in seconds
+    print(f'Instance number: {anno_len}, Batch size: {batch_size}')
+    print(f'Preprocess time: {preproces_time-start_time:02f}, FPS: {anno_len/(preproces_time-start_time):02f}')
+    print(f'Inference time: {end_time-preproces_time:02f}, FPS: {anno_len/(end_time-preproces_time):02f}')
+
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument('--in_threshold', type=float, default=0.5)
+    parser.add_argument('--pretrained_model', type=str, default='smpler_x_h32')
+
+    args = parser.parse_args()
+
+    inference(args)
+
+
+
+

main/SMPLer_X.py

@@ -8,7 +8,7 @@ from utils.transforms import rot6d_to_axis_angle, restore_bbox
 from config import cfg
 import math
 import copy
-from mmpose.models import build_posenet
+from mmpose_smplerx.models import build_posenet
 from mmengine.config import Config
 
 class Model(nn.Module):

main/config.py

@@ -4,21 +4,44 @@ import sys
 import datetime
 from mmengine.config import Config as MMConfig
 
+import os
+import os.path as osp
+
+# write your path here
+
+encoder_config_file = '/home/weichen/sst/smplerx/main/transformer_utils/configs/smpler_x/encoder/body_encoder_huge.py'
+human_model_path = '/home/weichen/wc_workspace/models/human_model'
+
+# <configure your smplerx model path here>
+model_path_dict = {
+    'smpler_x_h32': '/home/weichen/wc_workspace/models/smplerx/smpler_x_h32_correct.pth.tar',
+    'smpler_x_l32': '',
+    'smpler_x_b32': '',
+    'smpler_x_s32': '',
+}
+
 
 class Config:
     def get_config_fromfile(self, config_path):
+
         self.config_path = config_path
         cfg = MMConfig.fromfile(self.config_path)
+
+        # update config
+        cfg.encoder_config_file = encoder_config_file
+        cfg.human_model_path = human_model_path
+
         self.__dict__.update(dict(cfg))
 
         # update dir
         self.cur_dir = osp.dirname(os.path.abspath(__file__))
         self.root_dir = osp.join(self.cur_dir, '..')
         self.data_dir = osp.join(self.root_dir, 'dataset')
-        self.human_model_path = osp.join(self.root_dir, 'common', 'utils', 'human_model_files')
+        self.human_model_path = human_model_path
 
         ## add some paths to the system root dir
         sys.path.insert(0, osp.join(self.root_dir, 'common'))
+
                 
     def prepare_dirs(self, exp_name):
         time_str = datetime.datetime.now().strftime('%Y%m%d_%H%M%S')
@@ -63,4 +86,7 @@ class Config:
         # Save
         cfg_save = MMConfig(self.__dict__)
 
-cfg = Config()
+cfg = Config()
+cfg.human_model_path = human_model_path
+
+

main/inference.py

@@ -2,133 +2,127 @@ import os
 import sys
 import os.path as osp
 import argparse
+import json
+from tqdm import tqdm
+
+import cv2
 import numpy as np
-import torchvision.transforms as transforms
-import torch.backends.cudnn as cudnn
+
 import torch
+import torch.nn as nn
+import torch.backends.cudnn as cudnn
+import torchvision.transforms as transforms
+from torch.utils.data import Dataset
+
 CUR_DIR = osp.dirname(os.path.abspath(__file__))
 sys.path.insert(0, osp.join(CUR_DIR, '..', 'main'))
 sys.path.insert(0, osp.join(CUR_DIR , '..', 'common'))
-from config import cfg
-import cv2
-from tqdm import tqdm
-import json
-from typing import Literal, Union
-from mmdet.apis import init_detector, inference_detector
-from utils.inference_utils import process_mmdet_results, non_max_suppression
+from config import cfg, model_path_dict
+from base import Demoer
+from utils.preprocessing import process_bbox, generate_patch_image
+from utils.human_models import smpl_x
+
+class SmplerxData(Dataset):
+    def __init__(self, annotations):
+        self.annotations = annotations
+
+    def __len__(self):
+        return len(self.annotations)
+
+    def __getitem__(self, idx):
+        img_path = self.annotations[idx]['image_path']
+        image = cv2.imread(img_path)
+        bbox = self.annotations[idx]['bbox']
+        if bbox[2] < 50 or bbox[3] < 150:
+            return None
+        img_shape = image.shape  # (width, height)
+
+        # prepare input image
+        transform = transforms.ToTensor()
+        original_img_height, original_img_width = image.shape[:2]
+        bbox = process_bbox(bbox, original_img_width, original_img_height)
+        img, img2bb_trans, bb2img_trans = generate_patch_image(image, bbox, 1.0, 0.0, False, (512, 384))
+        img = transform(img.astype(np.float32))/255
+        
+        sample = {'image': img, 'bbox': bbox, 'shape': img_shape, 'path': img_path}
+        
+        return sample
+
 
 class Inferer:
 
-    def __init__(self, pretrained_model, num_gpus, output_folder):
+    def __init__(self, pretrained_model, num_gpus, data_parallel=True,
+                 output_folder=osp.join(CUR_DIR, '..', 'demo_out')):
+
         self.output_folder = output_folder
+        self.data_parallel = data_parallel
         self.device = torch.device('cuda') if (num_gpus > 0) else torch.device('cpu')
-        config_path = osp.join(CUR_DIR, './config', f'config_{pretrained_model}.py')
-        ckpt_path = osp.join(CUR_DIR, '../pretrained_models', f'{pretrained_model}.pth.tar')
+
+        # load config and model path
+        ckpt_path = model_path_dict[pretrained_model]
+        config_path = osp.join(CUR_DIR, 'config', f'config_{pretrained_model}.py')
+
         cfg.get_config_fromfile(config_path)
         cfg.update_config(num_gpus, ckpt_path, output_folder, self.device)
         self.cfg = cfg
         cudnn.benchmark = True
-        
         # load model
-        from base import Demoer
         demoer = Demoer()
+        # if num_gpus > 1:
         demoer._make_model()
+        if self.data_parallel:
+            demoer.model = nn.DataParallel(demoer.model)
         demoer.model.eval()
         self.demoer = demoer
-        checkpoint_file = osp.join(CUR_DIR, '../pretrained_models/mmdet/faster_rcnn_r50_fpn_1x_coco_20200130-047c8118.pth')
-        config_file= osp.join(CUR_DIR, '../pretrained_models/mmdet/mmdet_faster_rcnn_r50_fpn_coco.py')
-        model = init_detector(config_file, checkpoint_file, device=self.device)  # or device='cuda:0'
-        self.model = model
 
-    def infer(self, original_img, iou_thr, frame, multi_person=False, mesh_as_vertices=False):
-        from utils.preprocessing import process_bbox, generate_patch_image
-        from utils.vis import render_mesh, save_obj
-        from utils.human_models import smpl_x
-        mesh_paths = []
-        smplx_paths = []
-        # prepare input image
-        transform = transforms.ToTensor()
-        vis_img = original_img.copy()
-        original_img_height, original_img_width = original_img.shape[:2]
 
-        ## mmdet inference
-        mmdet_results = inference_detector(self.model, original_img)
-        
-        pred_instance = mmdet_results.pred_instances.cpu().numpy()
-        bboxes = np.concatenate(
-            (pred_instance.bboxes, pred_instance.scores[:, None]), axis=1)
-        bboxes = bboxes[pred_instance.labels == 0]
-        bboxes = np.expand_dims(bboxes, axis=0)
-        mmdet_box = process_mmdet_results(bboxes, cat_id=0, multi_person=True)
+    def _get_focal(self, bbox):
+        bbox = bbox.cpu().numpy()
+        focal = [self.cfg.focal[0] / self.cfg.input_body_shape[1] * bbox[2],
+                 self.cfg.focal[0] / self.cfg.input_body_shape[0] * bbox[3]]
+        return focal
+    
 
-        # save original image if no bbox
-        if len(mmdet_box[0])<1:
-            return original_img, [], []
-        
-        if not multi_person:
-            # only select the largest bbox
-            num_bbox = 1
-            mmdet_box = mmdet_box[0]
-        else:
-            # keep bbox by NMS with iou_thr
-            mmdet_box = non_max_suppression(mmdet_box[0], iou_thr)
-            num_bbox = len(mmdet_box)
-        
-        ## loop all detected bboxes
-        for bbox_id in range(num_bbox):
-            mmdet_box_xywh = np.zeros((4))
-            mmdet_box_xywh[0] = mmdet_box[bbox_id][0]
-            mmdet_box_xywh[1] = mmdet_box[bbox_id][1]
-            mmdet_box_xywh[2] =  abs(mmdet_box[bbox_id][2]-mmdet_box[bbox_id][0])
-            mmdet_box_xywh[3] =  abs(mmdet_box[bbox_id][3]-mmdet_box[bbox_id][1]) 
+    def _get_princpt(self, bbox):
+        bbox = bbox.cpu().numpy()
+        princpt = [self.cfg.princpt[0] / self.cfg.input_body_shape[1] * bbox[2] + bbox[0],
+                   self.cfg.princpt[1] / self.cfg.input_body_shape[0] * bbox[3] + bbox[1]]
+        return princpt
 
-            # skip small bboxes by bbox_thr in pixel
-            if mmdet_box_xywh[2] < 50 or mmdet_box_xywh[3] < 150:
-                continue
 
-            bbox = process_bbox(mmdet_box_xywh, original_img_width, original_img_height)
-            img, img2bb_trans, bb2img_trans = generate_patch_image(original_img, bbox, 1.0, 0.0, False, self.cfg.input_img_shape)
-            img = transform(img.astype(np.float32))/255
-            img = img.to(cfg.device)[None,:,:,:]
-            inputs = {'img': img}
-            targets = {}
-            meta_info = {}
+    def batch_infer_given_bbox(self, img, bbox, return_mesh=False):
 
-            # mesh recovery
-            with torch.no_grad():
-                out = self.demoer.model(inputs, targets, meta_info, 'test')
+        batch_size = img.shape[0]
+        inputs = {'img': img}
+        targets = {}
+        meta_info = {}
+
+        # mesh recovery
+        with torch.no_grad():
+            out = self.demoer.model(inputs, targets, meta_info, 'test')
+
+        ## save mesh
+        if return_mesh:
             mesh = out['smplx_mesh_cam'].detach().cpu().numpy()[0]
+        else:
+            mesh = None
 
-            ## save mesh
-            save_path_mesh = os.path.join(self.output_folder, 'mesh')
-            os.makedirs(save_path_mesh, exist_ok= True)
-            obj_path = os.path.join(save_path_mesh, f'{frame:05}_{bbox_id}.obj')
-            save_obj(mesh, smpl_x.face, obj_path)
-            mesh_paths.append(obj_path)
-            ## save single person param
-            smplx_pred = {}
-            smplx_pred['global_orient'] = out['smplx_root_pose'].reshape(-1,3).cpu().numpy()
-            smplx_pred['body_pose'] = out['smplx_body_pose'].reshape(-1,3).cpu().numpy()
-            smplx_pred['left_hand_pose'] = out['smplx_lhand_pose'].reshape(-1,3).cpu().numpy()
-            smplx_pred['right_hand_pose'] = out['smplx_rhand_pose'].reshape(-1,3).cpu().numpy()
-            smplx_pred['jaw_pose'] = out['smplx_jaw_pose'].reshape(-1,3).cpu().numpy()
-            smplx_pred['leye_pose'] = np.zeros((1, 3))
-            smplx_pred['reye_pose'] = np.zeros((1, 3))
-            smplx_pred['betas'] = out['smplx_shape'].reshape(-1,10).cpu().numpy()
-            smplx_pred['expression'] = out['smplx_expr'].reshape(-1,10).cpu().numpy()
-            smplx_pred['transl'] =  out['cam_trans'].reshape(-1,3).cpu().numpy()
-            save_path_smplx = os.path.join(self.output_folder, 'smplx')
-            os.makedirs(save_path_smplx, exist_ok= True)
+        ## save single person param
+        smplx_pred = {}
+        smplx_pred['global_orient'] = out['smplx_root_pose'].reshape(batch_size,-1,3).cpu().numpy()
+        smplx_pred['body_pose'] = out['smplx_body_pose'].reshape(batch_size,-1,3).cpu().numpy()
+        smplx_pred['left_hand_pose'] = out['smplx_lhand_pose'].reshape(batch_size,-1,3).cpu().numpy()
+        smplx_pred['right_hand_pose'] = out['smplx_rhand_pose'].reshape(batch_size,-1,3).cpu().numpy()
+        smplx_pred['jaw_pose'] = out['smplx_jaw_pose'].reshape(batch_size,-1,3).cpu().numpy()
+        smplx_pred['leye_pose'] = np.zeros((batch_size,1,3))
+        smplx_pred['reye_pose'] = np.zeros((batch_size,1,3))
+        smplx_pred['betas'] = out['smplx_shape'].reshape(batch_size,-1,10).cpu().numpy()
+        smplx_pred['expression'] = out['smplx_expr'].reshape(batch_size,-1,10).cpu().numpy()
+        smplx_pred['transl'] =  out['cam_trans'].reshape(batch_size,-1,3).cpu().numpy()
 
-            npz_path = os.path.join(save_path_smplx, f'{frame:05}_{bbox_id}.npz')
-            np.savez(npz_path, **smplx_pred)
-            smplx_paths.append(npz_path)
-
-            ## render single person mesh
-            focal = [self.cfg.focal[0] / self.cfg.input_body_shape[1] * bbox[2], self.cfg.focal[1] / self.cfg.input_body_shape[0] * bbox[3]]
-            princpt = [self.cfg.princpt[0] / self.cfg.input_body_shape[1] * bbox[2] + bbox[0], self.cfg.princpt[1] / self.cfg.input_body_shape[0] * bbox[3] + bbox[1]]
-            vis_img = render_mesh(vis_img, mesh, smpl_x.face, {'focal': focal, 'princpt': princpt}, 
-                                  mesh_as_vertices=mesh_as_vertices)
-            vis_img = vis_img.astype('uint8') 
-        return vis_img, mesh_paths, smplx_paths
+        ## save meta
+        meta = {}
+        meta['focal_length'] = [(self._get_focal(box)) for box in bbox]
+        meta['principal_point'] = [(self._get_princpt(box)) for box in bbox]
 
+        return smplx_pred, meta, mesh

main/transformer_utils/mmpose_smplerx/.mim/configs

@@ -0,0 +1 @@
+../../configs

main/transformer_utils/mmpose_smplerx/.mim/model-index.yml

@@ -0,0 +1 @@
+../../model-index.yml

main/transformer_utils/mmpose_smplerx/__init__.py

@@ -1,6 +1,6 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 import mmcv
-import mmpose.ops
+import mmpose_smplerx.ops
 from .version import __version__, short_version
 
 

main/transformer_utils/mmpose_smplerx/core/evaluation/bottom_up_eval.py

@@ -2,7 +2,7 @@
 import numpy as np
 import torch
 
-from mmpose.core.post_processing import (get_warp_matrix, transform_preds,
+from mmpose_smplerx.core.post_processing import (get_warp_matrix, transform_preds,
                                          warp_affine_joints)
 
 

main/transformer_utils/mmpose_smplerx/core/evaluation/top_down_eval.py

@@ -4,7 +4,7 @@ import warnings
 import cv2
 import numpy as np
 
-from mmpose.core.post_processing import transform_preds
+from mmpose_smplerx.core.post_processing import transform_preds
 
 
 def _calc_distances(preds, targets, mask, normalize):

main/transformer_utils/mmpose_smplerx/core/optimizers/layer_decay_optimizer_constructor.py

@@ -4,7 +4,7 @@ import warnings
 from mmengine.dist import get_dist_info
 from mmcv.runner import DefaultOptimizerConstructor
 
-from mmpose.utils import get_root_logger
+from mmpose_smplerx.utils import get_root_logger
 from .builder import OPTIMIZER_BUILDERS
 
 

main/transformer_utils/mmpose_smplerx/core/post_processing/group.py

@@ -7,7 +7,7 @@ import numpy as np
 import torch
 from munkres import Munkres
 
-from mmpose.core.evaluation import post_dark_udp
+from mmpose_smplerx.core.evaluation import post_dark_udp
 
 
 def _py_max_match(scores):

main/transformer_utils/mmpose_smplerx/core/post_processing/smoother.py

@@ -6,7 +6,7 @@ from typing import Dict, Union
 import numpy as np
 from mmengine.config import Config
 from mmengine.utils import is_seq_of
-from mmpose.core.post_processing.temporal_filters import build_filter
+from mmpose_smplerx.core.post_processing.temporal_filters import build_filter
 
 
 class Smoother():

main/transformer_utils/mmpose_smplerx/models/builder.py

@@ -3,7 +3,7 @@ from mmengine.registry import MODELS as MMCV_MODELS
 from mmengine import Registry
 from mmengine.registry import build_from_cfg, build_model_from_cfg
 
-MODELS = Registry('models', parent=MMCV_MODELS, locations=['mmpose.models'])
+MODELS = Registry('models', parent=MMCV_MODELS, locations=['mmpose_smplerx.models'])
 
 BACKBONES = MODELS
 NECKS = MODELS

main/transformer_utils/mmpose_smplerx/models/detectors/poseur.py

@@ -5,14 +5,14 @@ import numpy as np
 from mmcv.image import imwrite
 from mmcv.visualization.image import imshow
 
-from mmpose.core import imshow_keypoints
+from mmpose_smplerx.core import imshow_keypoints
 from .. import builder
 from ..builder import POSENETS
 from .base import BasePose
 import torch
 from config import cfg
 
-from mmpose.core import auto_fp16
+from mmpose_smplerx.core import auto_fp16
 
 from .top_down import TopDown
 

main/transformer_utils/mmpose_smplerx/models/detectors/top_down.py

@@ -7,12 +7,12 @@ from mmcv.image import imwrite
 from mmengine.utils import deprecated_api_warning
 from mmcv.visualization.image import imshow
 
-from mmpose.core import imshow_bboxes, imshow_keypoints
+from mmpose_smplerx.core import imshow_bboxes, imshow_keypoints
 from .. import builder
 from ..builder import POSENETS
 from .base import BasePose
 
-from mmpose.core import auto_fp16
+from mmpose_smplerx.core import auto_fp16
 
 
 @POSENETS.register_module()

main/transformer_utils/mmpose_smplerx/models/heads/poseur_head.py

@@ -8,20 +8,20 @@ from mmengine.model import constant_init, normal_init, bias_init_with_prob
 from mmcv.cnn import build_upsample_layer, Linear
 import torch.nn.functional as F
 
-from mmpose.core.evaluation import (keypoint_pck_accuracy,
+from mmpose_smplerx.core.evaluation import (keypoint_pck_accuracy,
                                     keypoints_from_regression)
-from mmpose.core.post_processing import fliplr_regression
-from mmpose.models.builder import build_loss, HEADS, build_transformer
-from mmpose.core.evaluation import pose_pck_accuracy
-from mmpose.models.utils.transformer import inverse_sigmoid
+from mmpose_smplerx.core.post_processing import fliplr_regression
+from mmpose_smplerx.models.builder import build_loss, HEADS, build_transformer
+from mmpose_smplerx.core.evaluation import pose_pck_accuracy
+from mmpose_smplerx.models.utils.transformer import inverse_sigmoid
 from mmcv.cnn import Conv2d, build_activation_layer
 from mmcv.cnn.bricks.transformer import Linear, FFN, build_positional_encoding
 from mmcv.cnn import ConvModule
 import torch.distributions as distributions
 from .rle_regression_head import nets, nett, RealNVP, nets3d, nett3d
 from easydict import EasyDict
-from mmpose.models.losses.regression_loss import L1Loss
-from mmpose.models.losses.rle_loss import RLELoss_poseur, RLEOHKMLoss
+from mmpose_smplerx.models.losses.regression_loss import L1Loss
+from mmpose_smplerx.models.losses.rle_loss import RLELoss_poseur, RLEOHKMLoss
 from config import cfg
 from utils.human_models import smpl_x
 from torch.distributions.utils import lazy_property

main/transformer_utils/mmpose_smplerx/models/heads/rle_regression_head.py

@@ -1,10 +1,10 @@
 import numpy as np
 import torch.nn as nn
 from mmengine.model import normal_init
-from mmpose.core.evaluation import (keypoint_pck_accuracy,
+from mmpose_smplerx.core.evaluation import (keypoint_pck_accuracy,
                                     keypoints_from_regression)
-from mmpose.core.post_processing import fliplr_regression
-from mmpose.models.builder import HEADS, build_loss
+from mmpose_smplerx.core.post_processing import fliplr_regression
+from mmpose_smplerx.models.builder import HEADS, build_loss
 
 import torch
 import torch.nn as nn

main/transformer_utils/mmpose_smplerx/models/heads/topdown_heatmap_base_head.py

@@ -4,7 +4,7 @@ from abc import ABCMeta, abstractmethod
 import numpy as np
 import torch.nn as nn
 
-from mmpose.core.evaluation.top_down_eval import keypoints_from_heatmaps
+from mmpose_smplerx.core.evaluation.top_down_eval import keypoints_from_heatmaps
 
 
 class TopdownHeatmapBaseHead(nn.Module):

main/transformer_utils/mmpose_smplerx/models/heads/topdown_heatmap_multi_stage_head.py

@@ -7,9 +7,9 @@ from mmcv.cnn import (ConvModule, DepthwiseSeparableConvModule, Linear,
                       build_activation_layer, build_conv_layer,
                       build_norm_layer, build_upsample_layer)
 
-from mmpose.core.evaluation import pose_pck_accuracy
-from mmpose.core.post_processing import flip_back
-from mmpose.models.builder import build_loss
+from mmpose_smplerx.core.evaluation import pose_pck_accuracy
+from mmpose_smplerx.core.post_processing import flip_back
+from mmpose_smplerx.models.builder import build_loss
 from ..builder import HEADS
 from .topdown_heatmap_base_head import TopdownHeatmapBaseHead
 

main/transformer_utils/mmpose_smplerx/models/heads/topdown_heatmap_simple_head.py

@@ -4,10 +4,10 @@ import torch.nn as nn
 from mmengine.model import constant_init, normal_init
 from mmcv.cnn import (build_conv_layer, build_norm_layer, build_upsample_layer)
 
-from mmpose.core.evaluation import pose_pck_accuracy
-from mmpose.core.post_processing import flip_back
-from mmpose.models.builder import build_loss
-from mmpose.models.utils.ops import resize
+from mmpose_smplerx.core.evaluation import pose_pck_accuracy
+from mmpose_smplerx.core.post_processing import flip_back
+from mmpose_smplerx.models.builder import build_loss
+from mmpose_smplerx.models.utils.ops import resize
 from ..builder import HEADS
 from .topdown_heatmap_base_head import TopdownHeatmapBaseHead
 

main/transformer_utils/mmpose_smplerx/models/misc/discriminator.py

@@ -9,7 +9,7 @@ import torch
 import torch.nn as nn
 from mmcv.cnn import normal_init, xavier_init
 
-from mmpose.models.utils.geometry import batch_rodrigues
+from mmpose_smplerx.models.utils.geometry import batch_rodrigues
 
 
 class BaseDiscriminator(nn.Module):