SHOW/render_pkl_release.py

import os
import sys
import os.path as osp
from pathlib import Path
import matplotlib
import platform
import warnings

if platform.system() == 'Windows':
    matplotlib.use('TkAgg')

if platform.system() == "Linux":
    os.environ['PYOPENGL_PLATFORM'] = 'egl'

sys.path.extend([
    osp.join(osp.dirname(__file__)),
])
warnings.filterwarnings("ignore")

import torch
from pathlib import Path
import numpy as np
from tqdm import tqdm
from loguru import logger
from easydict import EasyDict
import mmcv
import os
import mmcv
from loguru import logger
from pathlib import Path
import trimesh
import numpy as np
import pyrender
import os
import PIL.Image as pil_img
import cv2
import smplx
from loguru import logger

# from SHOW.utils.video import images_to_video
from mmhuman3d.utils.ffmpeg_utils import images_to_video


DEFAULT_SMPLX_CONFIG = dict(
    create_global_orient=True,
    create_body_pose=True,
    create_betas=True,
    create_left_hand_pose=True,
    create_right_hand_pose=True,
    create_expression=True,
    create_jaw_pose=True,
    create_leye_pose=True,
    create_reye_pose=True,
    create_transl=True,
)

colors_dict = {
    'red': np.array([0.5, 0.2, 0.2]),
    'pink': np.array([0.7, 0.5, 0.5]),
    'neutral': np.array([0.7, 0.7, 0.6]),
    # 'purple': np.array([0.5, 0.5, 0.7]),
    'purple': np.array([0.55, 0.4, 0.9]),
    'green': np.array([0.5, 0.55, 0.3]),
    'sky': np.array([0.3, 0.5, 0.55]),
    'white': np.array([1.0, 0.98, 0.94]),
}


def save_one_results(
    vertices,
    faces,
    img_size,  #(height,width)
    center,  #(cx,cy)
    focal_length,  #(focalx,focaly)
    camera_pose,  #K:(4,4)
    meta_data={},
    color_type='sky',
    input_renderer=None,
):

    input_img = meta_data.get('input_img', None)
    output_name = meta_data.get('output_name', None)

    if 1:
        color = colors_dict[color_type]

        out_mesh = trimesh.Trimesh(vertices, faces, process=False)

        out_mesh.apply_transform(
            trimesh.transformations.rotation_matrix(np.radians(180),
                                                    [1, 0, 0]))

        mesh = pyrender.Mesh.from_trimesh(
            out_mesh,
            material=pyrender.MetallicRoughnessMaterial(
                metallicFactor=0.2,
                roughnessFactor=0.6,
                alphaMode='OPAQUE',
                baseColorFactor=(color[0], color[1], color[2], 1.0)))
        # 黑色背景
        bg_color = [0.0, 0.0, 0.0, 0.0]

    # 渲染器参数设置
    if input_renderer is None:
        renderer = pyrender.OffscreenRenderer(viewport_width=img_size[1],
                                              viewport_height=img_size[0],
                                              point_size=1.0)
    else:
        renderer = input_renderer

    scene = pyrender.Scene(bg_color=bg_color, ambient_light=(0.3, 0.3, 0.3))
    scene.add(mesh, 'mesh')

    camera = pyrender.camera.IntrinsicsCamera(fx=focal_length[0],
                                              fy=focal_length[1],
                                              cx=center[0],
                                              cy=center[1])
    scene.add(camera, pose=camera_pose)

    if 0:
        light_node = pyrender.DirectionalLight()
        scene.add(light_node, pose=camera_pose)
    else:
        light = pyrender.PointLight(color=np.array([1.0, 1.0, 1.0]) * 0.2,
                                    intensity=1)
        light_pose = np.eye(4)
        light_pose[:3, 3] = [0, -1, 1]
        scene.add(light, pose=light_pose)

        light_pose[:3, 3] = [0, 1, 1]
        scene.add(light, pose=light_pose)

        light_pose[:3, 3] = [1, 1, 2]
        scene.add(light, pose=light_pose)

        spot_l = pyrender.SpotLight(color=np.ones(3),
                                    intensity=15.0,
                                    innerConeAngle=np.pi / 3,
                                    outerConeAngle=np.pi / 2)

        light_pose[:3, 3] = [1, 2, 2]
        scene.add(spot_l, pose=light_pose)

        light_pose[:3, 3] = [-1, 2, 2]
        scene.add(spot_l, pose=light_pose)

    color, _ = renderer.render(scene, flags=pyrender.RenderFlags.RGBA)
    color = color.astype(np.float32) / 255.0

    if input_img != None:
        if type(input_img) == str:
            input_img = cv2.imread(input_img)
            input_img = cv2.cvtColor(input_img, cv2.COLOR_BGR2RGB)

        assert (type(input_img) == np.ndarray)

        if input_img.max() > 1:
            input_img = input_img.astype(np.float32) / 255.0

        valid_mask = (color[:, :, -1] > 0)[:, :, np.newaxis]
        output_img = (color[:, :, :-1] * valid_mask +
                      (1 - valid_mask) * input_img)
        output_img = pil_img.fromarray((output_img * 255.).astype(np.uint8))

    if output_name != None:
        output_img = (color[:, :, :-1])
        output_img = pil_img.fromarray((output_img * 255.).astype(np.uint8))

        Path(output_name).parent.mkdir(exist_ok=True, parents=True)
        output_img.save(output_name)

    if input_renderer is None:
        renderer.delete()

    return output_img


def render_pkl_api(
        pkl_file_path,  # pkl path
        out_images_path,  # image path folder
        output_video_path=None,  # endwith .mp4
        smplx_model_path='../models/smplx/SMPLX_NEUTRAL_2020_org.npz',  #smplx neutral 2020 npz file
        **kwargs):
    dtype = torch.float32

    smplx_model_path = os.path.abspath(
        os.path.join(os.path.dirname(__file__), smplx_model_path))

    all_var = mmcv.load(pkl_file_path)
    if isinstance(all_var, list):
        all_var = all_var[0]
    all_var = EasyDict(all_var)

    body_model = smplx.create(**DEFAULT_SMPLX_CONFIG,
                              dtype=dtype,
                              model_path=smplx_model_path,
                              use_face_contour=True,
                              use_pca=True,
                              flat_hand_mean=False,
                              use_hands=True,
                              use_face=True,
                              num_pca_comps=12,
                              num_betas=300,
                              num_expression_coeffs=100,
                              batch_size=all_var.batch_size).to(device='cuda')

    def get_smplx_to_pyrender_K(cam_transl) -> np.ndarray:
        if isinstance(cam_transl, torch.Tensor):
            T = cam_transl.detach().cpu().numpy()
        T[1] *= -1
        K = np.eye(4)
        K[:3, 3] = T
        return K

    camera_pose = get_smplx_to_pyrender_K(
        torch.from_numpy(all_var.camera_transl))
    smplx_params = dict(
        body_pose=all_var.body_pose_axis,
        betas=all_var.betas,
        global_orient=all_var.global_orient,
        transl=all_var.transl,
        left_hand_pose=all_var.left_hand_pose,
        right_hand_pose=all_var.right_hand_pose,
        jaw_pose=all_var.jaw_pose,
        leye_pose=all_var.leye_pose,
        reye_pose=all_var.reye_pose,
        expression=all_var.expression,
    )
    for key, val in smplx_params.items():
        if isinstance(val, torch.Tensor):
            smplx_params[key] = smplx_params[key].detach().to('cuda')
        if isinstance(val, np.ndarray):
            smplx_params[key] = torch.from_numpy(smplx_params[key]).to('cuda')
    model_output = body_model(return_verts=True, **smplx_params)

    import pyrender
    input_renderer = pyrender.OffscreenRenderer(viewport_width=all_var.width,
                                                viewport_height=all_var.height,
                                                point_size=1.0)

    vertices_ = model_output.vertices.detach().cpu().numpy()
    Path(out_images_path).mkdir(parents=True, exist_ok=True)
    logger.info(f'saving final images to: {out_images_path}')

    for idx in tqdm(range(vertices_.shape[0])):  # idx=0
        vertices = vertices_[idx]
        meta_data = dict(
            # input_img=os.path.join(img_folder, f'{idx+1:06d}.png'),
            output_name=os.path.join(out_images_path, f'{idx+1:06d}.png'), )

        tmp = save_one_results(
            vertices,
            body_model.faces,
            img_size=[all_var.height, all_var.width],
            center=all_var.center,
            focal_length=[all_var.focal_length, all_var.focal_length],
            camera_pose=camera_pose,
            meta_data=meta_data,
            input_renderer=input_renderer,
        )

    input_renderer.delete()

    def files_num_in_dir(dir_name):
        if not Path(dir_name).exists():
            return -1
        return len(os.listdir(dir_name))

    def is_empty_dir(dir_name):
        if not os.path.exists(dir_name):
            return 1
        return int(files_num_in_dir(dir_name) == 0)

    if output_video_path is not None:
        logger.info(f'saving final video to: {output_video_path}')
        Path(output_video_path).parent.mkdir(parents=True, exist_ok=True)
        if not is_empty_dir(out_images_path):
            images_to_video(
                input_folder=out_images_path,
                output_path=output_video_path,
                img_format=None,
                fps=30,
            )


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--pkl_file_path', type=str, default='all.pkl')
    parser.add_argument('--out_images_path',
                        type=str,
                        default='../results/ours_images')
    parser.add_argument('--output_video_path',
                        type=str,
                        default='../results/ours.mp4')
    parser.add_argument(
        '--smplx_model_path',
        type=str,
        default=
        'path_to_models/smplx/SMPLX_NEUTRAL_2020_org.npz')
    args = parser.parse_args()

    render_pkl_api(pkl_file_path=args.pkl_file_path,
                   out_images_path=args.out_images_path,
                   output_video_path=args.output_video_path,
                   smplx_model_path=args.smplx_model_path)