From 78f63e628d956f2de25c8cb13205d4db8dc39870 Mon Sep 17 00:00:00 2001 From: Wei-Chen-hub <1259566226@qq.com> Date: Mon, 23 Oct 2023 16:55:02 +0800 Subject: [PATCH] add smplx 2d overlay --- README.md | 26 ++-- main/render.py | 354 +++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 369 insertions(+), 11 deletions(-) create mode 100644 main/render.py diff --git a/README.md b/README.md index ba345ea..7bff3a6 100644 --- a/README.md +++ b/README.md @@ -12,6 +12,7 @@ ## News +- [2023-10-23] More tools: smplx mesh overlay script and inference docker are added ! - [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. @@ -175,6 +176,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 light pyrender script for mesh overlay projection. +- Overlay script uses result from above inference +- Use ffmpeg to split video to images to support overlay +```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 @@ -213,17 +228,6 @@ sh slurm_test.sh {JOB_NAME} {NUM_GPU} {TRAIN_OUTPUT_DIR} {CKPT_ID} - 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 :/smplerx_inference/vid_input \ - -v :/smplerx_inference/vid_output \ - wcwcw/smplerx_inference:v0.2 --vid .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) diff --git a/main/render.py b/main/render.py new file mode 100644 index 0000000..335d287 --- /dev/null +++ b/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) \ No newline at end of file