Merge branch 'main' of https://github.com/caizhongang/SMPLer-X

# Conflicts:
#	README.md

README.md

@@ -14,6 +14,7 @@
 
 ## News
 - [2024-02-29] [HuggingFace](https://huggingface.co/spaces/caizhongang/SMPLer-X) demo is online!
+- [2023-10-23] Support visualization through SMPL-X mesh overlay and add inference docker. 
 - [2023-10-02] [arXiv](https://arxiv.org/abs/2309.17448) preprint is online!
 - [2023-09-28] [Homepage](https://caizhongang.github.io/projects/SMPLer-X/) and [Video](https://youtu.be/DepTqbPpVzY) are online!
 - [2023-07-19] Pretrained models are released.
@@ -42,6 +43,17 @@ pip install -v -e .
 cd ../..
 ```
 
+## Docker Support (Early Stage)
+```
+docker pull wcwcw/smplerx_inference:v0.2
+docker run  --gpus all -v <vid_input_folder>:/smplerx_inference/vid_input \
+        -v <vid_output_folder>:/smplerx_inference/vid_output \
+        wcwcw/smplerx_inference:v0.2 --vid <video_name>.mp4
+# Currently any customization need to be applied to /smplerx_inference/smplerx/inference_docker.py
+```
+- We recently developed a docker for inference at docker hub.
+- This docker image uses SMPLer-X-H32 as inference baseline and was tested at RTX3090 & WSL2 (Ubuntu 20.04).
+
 
 ## Pretrained Models
 |    Model     | Backbone | #Datasets | #Inst. | #Params | MPE  | Download |  FPS  |
@@ -164,10 +176,10 @@ SMPLer-X/
     └── preprocessed_datasets/  # HumanData files
 ```
 ## Inference 
-- Place the video to be inferenced under `SMPLer-X/demo/videos`
+- Place the video for inference under `SMPLer-X/demo/videos`
 - Prepare the pretrained models to be used for inference under `SMPLer-X/pretrained_models`
 - Prepare the mmdet pretrained model and config under `SMPLer-X/pretrained_models`
-- Inference out put will be saved in `SMPLer-X/demo/results`
+- Inference output will be saved in `SMPLer-X/demo/results`
 
 ```bash
 cd main
@@ -177,6 +189,20 @@ sh slurm_inference.sh {VIDEO_FILE} {FORMAT} {FPS} {PRETRAINED_CKPT}
 sh slurm_inference.sh test_video mp4 24 smpler_x_h32
 
 ```
+## 2D Smplx Overlay
+We provide a lightweight visualization script for mesh overlay based on pyrender.
+- Use ffmpeg to split video into images
+- The visualization script takes inference results (see above) as the input.
+```bash
+ffmpeg -i {VIDEO_FILE} -f image2 -vf fps=30 \
+        {SMPLERX INFERENCE DIR}/{VIDEO NAME (no extension)}/orig_img/%06d.jpg \
+        -hide_banner  -loglevel error
+
+cd main && python render.py \
+            --data_path {SMPLERX INFERENCE DIR} --seq {VIDEO NAME} \
+            --image_path {SMPLERX INFERENCE DIR}/{VIDEO NAME} \
+            --render_biggest_person False
+```
 
 
 ## Training
@@ -202,6 +228,7 @@ sh slurm_test.sh {JOB_NAME} {NUM_GPU} {TRAIN_OUTPUT_DIR} {CKPT_ID}
 - NUM_GPU = 1 is recommended for testing
 - Logs and results  will be saved to `SMPLer-X/output/test_{JOB_NAME}_ep{CKPT_ID}_{TEST_DATSET}`
 
+
 ## FAQ
 - `RuntimeError: Subtraction, the '-' operator, with a bool tensor is not supported. If you are trying to invert a mask, use the '~' or 'logical_not()' operator instead.`
   
@@ -209,7 +236,7 @@ sh slurm_test.sh {JOB_NAME} {NUM_GPU} {TRAIN_OUTPUT_DIR} {CKPT_ID}
 
 - `KeyError: 'SinePositionalEncoding is already registered in position encoding'` or any other similar KeyErrors due to duplicate module registration.
 
-  Maually add `force=True` to respective module registration under `main/transformer_utils/mmpose/models/utils`, e.g. `@POSITIONAL_ENCODING.register_module(force=True)` in [this file](main/transformer_utils/mmpose/models/utils/positional_encoding.py)
+  Manually add `force=True` to respective module registration under `main/transformer_utils/mmpose/models/utils`, e.g. `@POSITIONAL_ENCODING.register_module(force=True)` in [this file](main/transformer_utils/mmpose/models/utils/positional_encoding.py)
 
 - How do I animate my virtual characters with SMPLer-X output (like that in the demo video)? 
   - We are working on that, please stay tuned!

main/SMPLer_X.py

@@ -463,22 +463,6 @@ def get_model(mode):
 
     encoder = vit.backbone
 
-    # apply adapters
-    # currently, adapters have only been tested for ViTPose
-    adapter_name = getattr(cfg, 'adapter_name', None)
-    if adapter_name == 'lora':
-        from lora_utils import apply_adapter
-        encoder = apply_adapter(encoder)
-        print(f"Apply adapter {adapter_name}.")
-    elif adapter_name == 'vit_adapter':
-        from vit_adapter_utils import apply_adapter
-        encoder = apply_adapter(encoder, cfg.model_type)
-        print(f"Apply adapter {adapter_name}.")
-    elif adapter_name is not None:
-        raise NotImplementedError('Undefined adapter: {}'.format(adapter_name))
-    else:
-        print("No adapter used.")
-
     model = Model(encoder, body_position_net, body_rotation_net, box_net, hand_position_net, hand_roi_net, hand_rotation_net,
                   face_regressor)
     return model

main/render.py

@@ -0,0 +1,354 @@
+import numpy as np
+import glob
+import random
+import cv2
+import os
+import argparse
+import torch
+import pyrender
+import trimesh
+import pandas as pd
+import json
+
+from tqdm import tqdm
+from multiprocessing import Pool
+
+# from mmhuman3d.models.body_models.builder import build_body_model
+import smplx
+import pdb
+
+smpl_shape = {'betas': (-1, 10), 'transl': (-1, 3), 'global_orient': (-1, 3), 'body_pose': (-1, 69)}
+smplx_shape = {'betas': (-1, 10), 'transl': (-1, 3), 'global_orient': (-1, 3), 
+        'body_pose': (-1, 21, 3), 'left_hand_pose': (-1, 15, 3), 'right_hand_pose': (-1, 15, 3), 
+        'leye_pose': (-1, 3), 'reye_pose': (-1, 3), 'jaw_pose': (-1, 3), 'expression': (-1, 10)}
+smplx_shape_except_expression = {'betas': (-1, 10), 'transl': (-1, 3), 'global_orient': (-1, 3), 
+        'body_pose': (-1, 21, 3), 'left_hand_pose': (-1, 15, 3), 'right_hand_pose': (-1, 15, 3), 
+        'leye_pose': (-1, 3), 'reye_pose': (-1, 3), 'jaw_pose': (-1, 3)}
+# smplx_shape = smplx_shape_except_expression
+
+def render_pose(img, body_model_param, body_model, camera, return_mask=False):
+
+    # the inverse is same
+    pyrender2opencv = np.array([[1.0, 0, 0, 0],
+                                [0, -1, 0, 0],
+                                [0, 0, -1, 0],
+                                [0, 0, 0, 1]])
+    
+    output = body_model(**body_model_param, return_verts=True)
+    
+    vertices = output['vertices'].detach().cpu().numpy().squeeze()
+    faces = body_model.faces
+
+    # render material
+    base_color = (1.0, 193/255, 193/255, 1.0)
+    material = pyrender.MetallicRoughnessMaterial(
+            metallicFactor=0,
+            alphaMode='OPAQUE',
+            baseColorFactor=base_color)
+    
+    material_new = pyrender.MetallicRoughnessMaterial(
+            metallicFactor=0.1,
+            roughnessFactor=0.4,
+            alphaMode='OPAQUE',
+            emissiveFactor=(0.2, 0.2, 0.2),
+            baseColorFactor=(0.7, 0.7, 0.7, 1))  
+    material = material_new
+    
+    # get body mesh
+    body_trimesh = trimesh.Trimesh(vertices, faces, process=False)
+    body_mesh = pyrender.Mesh.from_trimesh(body_trimesh, material=material)
+
+    # prepare camera and light
+    light = pyrender.DirectionalLight(color=np.ones(3), intensity=2.0)
+    cam_pose = pyrender2opencv @ np.eye(4)
+    
+    # build scene
+    scene = pyrender.Scene(bg_color=[0.0, 0.0, 0.0, 0.0],
+                                    ambient_light=(0.3, 0.3, 0.3))
+    scene.add(camera, pose=cam_pose)
+    scene.add(light, pose=cam_pose)
+    scene.add(body_mesh, 'mesh')
+
+    # render scene
+    os.environ["PYOPENGL_PLATFORM"] = "osmesa" # include this line if use in vscode
+    r = pyrender.OffscreenRenderer(viewport_width=img.shape[1],
+                                    viewport_height=img.shape[0],
+                                    point_size=1.0)
+    
+    color, _ = r.render(scene, flags=pyrender.RenderFlags.RGBA)
+    color = color.astype(np.float32) / 255.0
+    # alpha = 1.0  # set transparency in [0.0, 1.0]
+    # color[:, :, -1] = color[:, :, -1] * alpha
+    valid_mask = (color[:, :, -1] > 0)[:, :, np.newaxis]
+    img = img / 255
+    # output_img = (color[:, :, :-1] * valid_mask + (1 - valid_mask) * img)
+    color = cv2.cvtColor(color, cv2.COLOR_BGR2RGB)
+    output_img = (color[:, :, :] * valid_mask + (1 - valid_mask) * img)
+
+    # output_img = color
+
+    img = (output_img * 255).astype(np.uint8)
+    # img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+
+    if return_mask:
+        return img, valid_mask, (color * 255).astype(np.uint8)
+
+    return img
+
+
+def render_multi_pose(img,
+                      body_model_params,
+                      body_model,
+                      cameras):
+
+    masks, colors = [], []
+
+    # calculate distance based on transl
+    dists, valid_idx = [], []
+    for i, body_model_param in enumerate(body_model_params):
+        dist = np.linalg.norm(body_model_param['transl'].detach().cpu()) * 2/ (cameras[i].fx + cameras[i].fy)
+        if dist not in dists:
+            valid_idx.append(i)
+            dists.append(dist)
+
+    # pdb.set_trace()
+
+    # select by valid idx
+    body_model_params = [body_model_params[i] for i in valid_idx]
+    cameras = [cameras[i] for i in valid_idx]
+
+    # sort by dist
+
+    body_model_params = [x for _, x in sorted(zip(dists, body_model_params), reverse=True)]
+    cameras = [x for _, x in sorted(zip(dists, cameras), reverse=True)]
+
+
+    # render separate masks
+    for i, body_model_param in enumerate(body_model_params):
+
+        _, mask, color = render_pose(
+            img=img,
+            body_model_param=body_model_param,
+            body_model=body_model,
+            camera=cameras[i],
+            return_mask=True,
+        )
+        masks.append(mask)
+        colors.append(color)
+    # sum masks
+    mask_sum = np.sum(masks, axis=0)
+    mask_all = (mask_sum > 0)
+
+    # pp_occ = 1 - np.sum(mask_all) / np.sum(mask_sum)
+
+    # overlay colors to img
+    for i, color in enumerate(colors):
+        mask = masks[i]
+        img = img * (1 - mask) + color * mask
+
+    img = img.astype(np.uint8)
+    # img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+
+    return img
+
+def render_frame(framestamp, anno_ps, image_base_path, seq, smplx_model, args):
+    annos = [p for p in anno_ps if framestamp in os.path.basename(p)]
+    annos = [p for p in annos if 'person' not in os.path.basename(p)]
+
+    body_model_params = []
+    cameras = []
+    bbox_sizes = []
+    try:
+        # image_path = os.path.join(seq, f'0{framestamp}.jpg').replace(args.data_path, args.image_path)
+        image_path = os.path.join(image_base_path, f'0{framestamp}.jpg')
+        # pdb.set_trace()
+        image = cv2.imread(image_path)
+    except:
+
+        pass
+    # pdb.set_trace()
+    for anno_p in annos:
+
+        anno = dict(np.load(anno_p, allow_pickle=True))
+
+        meta = json.load(open(os.path.join(seq, 'meta', 
+                                        os.path.basename(anno_p).replace('.npz', '.json')
+                                        )))
+
+        bbox_size = meta['bbox'][2] * meta['bbox'][3]
+        focal_length = meta['focal']
+        principal_point = meta['princpt']
+        camera = pyrender.camera.IntrinsicsCamera(
+                fx=focal_length[0], fy=focal_length[1],
+                cx=principal_point[0], cy=principal_point[1],)
+
+        # prepare body model params
+        intersect_key = list(set(anno.keys()) & set(smplx_shape.keys()))
+        body_model_param_tensor = {key: torch.tensor(
+                np.array(anno[key]).reshape(smplx_shape[key]), device=args.device, dtype=torch.float32)
+                        for key in intersect_key if len(anno[key]) > 0}
+        
+        cameras.append(camera)
+        body_model_params.append(body_model_param_tensor)
+        bbox_sizes.append(bbox_size)
+
+    # render pose
+    if args.render_biggest_person == 'True':
+        bid = bbox_sizes.index(max(bbox_sizes))
+        rendered_image = render_pose(img=image,
+                        body_model_param=body_model_params[bid],
+                        body_model=smplx_model,
+                        camera=cameras[bid])
+    else:
+        rendered_image = render_multi_pose(img=image, 
+                        body_model_params=body_model_params, 
+                        body_model=smplx_model,
+                        cameras=cameras)
+
+    sp = seq.replace(f'{args.data_path}{os.path.sep}', '')
+    save_path = os.path.join(args.data_path, 'output', sp)
+    os.makedirs(save_path, exist_ok=True)
+
+    save_name = os.path.join(save_path, framestamp+'.jpg')
+    cv2.imwrite(save_name, rendered_image)
+
+
+def call_frame_render(args):
+    return render_frame(*args)
+
+def visualize_seqs(args):
+
+    if args.seq == 'default':
+        seqs = glob.glob(os.path.join(args.data_path, '**/smplx'), recursive=True)
+        seqs = [os.path.dirname(p) for p in seqs]
+    else:
+        seqs = glob.glob(os.path.join(args.data_path, args.seq), recursive=True)    
+    
+    kwargs = dict(gender='neutral',
+        num_betas=10,
+        use_face_contour=True,
+        flat_hand_mean=args.flat_hand_mean,
+        use_pca=False,
+        batch_size=1)
+
+    smplx_model = smplx.create(
+        '../common/utils/human_model_files', 'smplx', 
+        **kwargs).to(args.device)
+
+    # seqs = [p for p in seqs if 'dance' not in os.path.basename(p)]
+
+    # pdb.set_trace()
+    
+    for i, seq in enumerate(seqs):
+        
+        # prepare image path
+        if args.load_mode == 'smplerx':
+            image_base_path = os.path.join(seq, 'orig_img')
+        else:
+            image_base_path = os.path.join(seq, 'frames')
+
+        assert os.path.exists(image_base_path)
+
+        smplx_path = os.path.join(seq, 'smplx')
+        anno_ps = sorted(glob.glob(os.path.join(smplx_path, '*.npz')))
+
+        # group by framestamps
+        framestamps = sorted(list(set([os.path.basename(p)[:5] for p in anno_ps 
+                                       if 'person' not in os.path.basename(p)]
+                                       )))
+
+        for framestamp in tqdm(framestamps, leave=False, desc=f'Seqs {os.path.basename(seq)}'
+                               f' : {i}/{len(seqs)}'):
+                
+            annos = [p for p in anno_ps if framestamp in os.path.basename(p)]
+            annos = [p for p in annos if 'person' not in os.path.basename(p)]
+
+            body_model_params = []
+            cameras = []
+            bbox_sizes = []
+            # pdb.set_trace()
+            try:
+                if args.load_mode == 'smplerx':
+                    image_path = os.path.join(image_base_path, f'{int(framestamp):06d}.jpg')
+                    image = cv2.imread(image_path)
+                else:
+                    image_path = os.path.join(image_base_path, f'{int(framestamp):04d}.png')
+                    image = cv2.imread(image_path)
+                    image = cv2.resize(image, (1280, 720), interpolation = cv2.INTER_AREA)
+                # image = cv2.imread(image_path)
+                
+            except:
+                pass
+
+            for anno_p in annos:
+
+                anno = dict(np.load(anno_p, allow_pickle=True))
+
+                meta = json.load(open(os.path.join(seq, 'meta', 
+                                                os.path.basename(anno_p).replace('.npz', '.json')
+                                                )))
+
+                bbox_size = meta['bbox'][2] * meta['bbox'][3]
+                focal_length = meta['focal']
+                principal_point = meta['princpt']
+                camera = pyrender.camera.IntrinsicsCamera(
+                        fx=focal_length[0], fy=focal_length[1],
+                        cx=principal_point[0], cy=principal_point[1],)
+
+                # prepare body model params
+                intersect_key = list(set(anno.keys()) & set(smplx_shape.keys()))
+                body_model_param_tensor = {key: torch.tensor(
+                        np.array(anno[key]).reshape(smplx_shape[key]), device=args.device, dtype=torch.float32)
+                                for key in intersect_key if len(anno[key]) > 0}
+                
+                cameras.append(camera)
+                body_model_params.append(body_model_param_tensor)
+                bbox_sizes.append(bbox_size)
+
+            # render pose
+            if args.render_biggest_person == 'True':
+                bid = bbox_sizes.index(max(bbox_sizes))
+                rendered_image = render_pose(img=image,
+                                body_model_param=body_model_params[bid],
+                                body_model=smplx_model,
+                                camera=cameras[bid])
+            else:
+                rendered_image = render_multi_pose(img=image, 
+                                body_model_params=body_model_params, 
+                                body_model=smplx_model,
+                                cameras=cameras)
+            # save image
+            sp = seq.replace(f'{args.data_path}{os.path.sep}', '')
+            save_path = os.path.join(args.data_path, sp, f'{args.load_mode}_overlay_img')
+            os.makedirs(save_path, exist_ok=True)
+
+            save_name = os.path.join(save_path, f'{int(framestamp):06d}.jpg')
+            cv2.imwrite(save_name, rendered_image)
+            
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument('--data_path', type=str, required=False,
+                        help='path to the data folder')
+    parser.add_argument('--load_mode', type=str, required=False,
+                        default='smplerx',
+                        help='load mode: smplerx or other test mode, please select smplerx')
+    parser.add_argument('--seq', type=str, required=False,
+                        help='seq name or seq pattern',
+                        default='default')
+    parser.add_argument('--image_path', type=str, required=False,
+                        help='path to the image folder')
+
+    # optional args
+    parser.add_argument('--flat_hand_mean', type=bool, required=False,
+                        help='use flat hand mean for smplx',
+                        default=False)
+    parser.add_argument('--render_biggest_person', type=str, required=False,
+                        help='render biggest person in the frame',
+                        default='True')
+    args = parser.parse_args()
+
+    args.device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+
+    visualize_seqs(args)