electron-rare/SMPLer-X
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #19 from caizhongang/mesh_overlay

Browse Source 
[Feature] Add mesh overlay visualization
This commit is contained in:
Zhongang Cai
2023-10-23 19:51:31 +08:00
committed by GitHub
parent bf61a727b3 80721f4fb9
commit 74bc864f66
2 changed files with 384 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+30 -14
View File
@@ -12,6 +12,7 @@
</div>
## News
- [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.
@@ -40,6 +41,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 |
@@ -162,10 +174,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
@@ -175,6 +187,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
@@ -200,6 +226,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.`
@@ -207,23 +234,12 @@ 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!
Currently, this repo supports SMPL-X estimation and a simple visualization (overlay of SMPL-X vertices).
## 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.
## References
- [Hand4Whole](https://github.com/mks0601/Hand4Whole_RELEASE)
- [OSX](https://github.com/IDEA-Research/OSX)
main/render.py
+354
View File
@@ -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)