mGPT/render/pyrender/smpl_render.py

import os

os.environ['PYOPENGL_PLATFORM'] = 'egl'
import torch
import numpy as np
import cv2

import matplotlib.pyplot as plt
import glob
import pickle
import pyrender
import trimesh
from shapely import geometry
from smplx import SMPL as _SMPL
from smplx.utils import SMPLOutput as ModelOutput
from scipy.spatial.transform.rotation import Rotation as RRR


class SMPL(_SMPL):
    """ Extension of the official SMPL implementation to support more joints """

    def __init__(self, *args, **kwargs):
        super(SMPL, self).__init__(*args, **kwargs)
        # joints = [constants.JOINT_MAP[i] for i in constants.JOINT_NAMES]
        # J_regressor_extra = np.load(config.JOINT_REGRESSOR_TRAIN_EXTRA)
        # self.register_buffer('J_regressor_extra', torch.tensor(J_regressor_extra, dtype=torch.float32))
        # self.joint_map = torch.tensor(joints, dtype=torch.long)

    def forward(self, *args, **kwargs):
        kwargs['get_skin'] = True
        smpl_output = super(SMPL, self).forward(*args, **kwargs)
        # extra_joints = vertices2joints(self.J_regressor_extra, smpl_output.vertices)        #Additional 9 joints #Check doc/J_regressor_extra.png
        # joints = torch.cat([smpl_output.joints, extra_joints], dim=1)               #[N, 24 + 21, 3]  + [N, 9, 3]
        # joints = joints[:, self.joint_map, :]
        joints = smpl_output.joints
        output = ModelOutput(vertices=smpl_output.vertices,
                             global_orient=smpl_output.global_orient,
                             body_pose=smpl_output.body_pose,
                             joints=joints,
                             betas=smpl_output.betas,
                             full_pose=smpl_output.full_pose)
        return output


class Renderer:
    """
    Renderer used for visualizing the SMPL model
    Code adapted from https://github.com/vchoutas/smplify-x
    """

    def __init__(self,
                 vertices,
                 focal_length=5000,
                 img_res=(224, 224),
                 faces=None):
        self.renderer = pyrender.OffscreenRenderer(viewport_width=img_res[0],
                                                   viewport_height=img_res[1],
                                                   point_size=1.0)
        self.focal_length = focal_length
        self.camera_center = [img_res[0] // 2, img_res[1] // 2]
        self.faces = faces

        if torch.cuda.is_available():
            self.device = torch.device("cuda")
        else:
            self.device = torch.device("cpu")

        vertices = np.concatenate(vertices)
        # Center the first root to the first frame
        vertices -= vertices[[0], [0], :]
        # Remove the floor
        vertices[..., 2] -= vertices[..., 2].min()
        data = vertices[..., [2, 0, 1]]
        minx, miny, _ = data.min(axis=(0, 1))
        maxx, maxy, _ = data.max(axis=(0, 1))
        minz, maxz = -0.5, 0.5
        minx = minx - 0.5
        maxx = maxx + 0.5
        miny = miny - 0.5
        maxy = maxy + 0.5

        polygon = geometry.Polygon([[minx, minz], [minx, maxz], [maxx, maxz],
                                    [maxx, minz]])
        self.polygon_mesh = trimesh.creation.extrude_polygon(polygon, 1e-5)
        self.polygon_mesh.visual.face_colors = [0, 0, 0, 0.21]
        self.rot = trimesh.transformations.rotation_matrix(
            np.radians(180), [1, 0, 0])
        # self.polygon_mesh.apply_transform(self.rot)

    def __call__(self, vertices, camera_translation):
        scene = pyrender.Scene(bg_color=(1., 1., 1., 0.8),
                               ambient_light=(0.4, 0.4, 0.4))

        material = pyrender.MetallicRoughnessMaterial(
            metallicFactor=0.4,
            alphaMode='OPAQUE',
            baseColorFactor=(0.658, 0.214, 0.0114, 0.2))
        mesh = trimesh.Trimesh(vertices, self.faces)
        mesh.apply_transform(self.rot)
        mesh = pyrender.Mesh.from_trimesh(mesh, material=material)
        scene.add(mesh, 'mesh')

        polygon_render = pyrender.Mesh.from_trimesh(self.polygon_mesh,
                                                    smooth=False)
        c = np.pi / 2
        scene.add(polygon_render)

        camera_pose = np.eye(4)
        camera_translation[0] *= -1.
        camera_pose[:3, 3] = camera_translation
        camera = pyrender.IntrinsicsCamera(fx=self.focal_length,
                                           fy=self.focal_length,
                                           cx=self.camera_center[0],
                                           cy=self.camera_center[1])
        scene.add(camera, pose=camera_pose)

        light = pyrender.DirectionalLight(color=[1, 1, 1], intensity=300)
        light_pose = np.eye(4)

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

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

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

        color, rend_depth = self.renderer.render(
            scene, flags=pyrender.RenderFlags.RGBA)

        return color


class SMPLRender():

    def __init__(self, SMPL_MODEL_DIR):
        if torch.cuda.is_available():
            self.device = torch.device("cuda")
        else:
            self.device = torch.device("cpu")
        self.smpl = SMPL(SMPL_MODEL_DIR, batch_size=1,
                         create_transl=False).to(self.device)
        self.vertices = []
        self.pred_camera_t = []
        self.focal_length = 5000

    def fit(self, smpl_param, is_headroot=False):
        pose = smpl_param['pred_pose']
        if pose.size == 72:
            pose = pose.reshape(-1, 3)
            pose = RRR.from_rotvec(pose).as_matrix()
            pose = pose.reshape(1, 24, 3, 3)
        pred_betas = torch.from_numpy(smpl_param['pred_shape'].reshape(
            1, 10).astype(np.float32)).to(self.device)
        pred_rotmat = torch.from_numpy(pose.astype(np.float32)).to(self.device)
        pred_camera_t = smpl_param['pred_root'].reshape(1,
                                                        3).astype(np.float32)
        smpl_output = self.smpl(betas=pred_betas,
                                body_pose=pred_rotmat[:, 1:],
                                global_orient=pred_rotmat[:, 0].unsqueeze(1),
                                pose2rot=False)
        vertices = smpl_output.vertices[0].detach().cpu().numpy()
        self.vertices.append(vertices[None])

        pred_camera_t = pred_camera_t[0]

        if is_headroot:
            pred_camera_t = pred_camera_t - smpl_output.joints[
                0, 12].detach().cpu().numpy()

        self.pred_camera_t.append(pred_camera_t)

    def init_renderer(self, res):
        self.renderer = Renderer(vertices=self.vertices,
                                 focal_length=self.focal_length,
                                 img_res=(res[1], res[0]),
                                 faces=self.smpl.faces)

    def render(self, index):
        renderImg = self.renderer(self.vertices[index][0],
                                  self.pred_camera_t[index].copy())
        return renderImg