OpenMotionLab/MotionGPT

1import os2import torch3import numpy as np4import cv25 6import matplotlib.pyplot as plt7import glob8import pickle9import pyrender10import trimesh11import smplx12from pathlib import Path13from shapely import geometry14from smplx import SMPL as _SMPL15from smplx.utils import SMPLOutput as ModelOutput16from scipy.spatial.transform.rotation import Rotation as RRR17 18class Renderer:19    """20    Renderer used for visualizing the SMPL model21    Code adapted from https://github.com/vchoutas/smplify-x22    """23    def __init__(self, vertices, focal_length=5000, img_res=(224,224), faces=None):24        self.renderer = pyrender.OffscreenRenderer(viewport_width=img_res[0],25                                       viewport_height=img_res[1],26                                       point_size=2.0)27        28        self.focal_length = focal_length29        self.camera_center = [img_res[0] // 2, img_res[1] // 2]30        self.faces = faces31        32        if torch.cuda.is_available():33            self.device = torch.device("cuda")34        else:35            self.device = torch.device("cpu")36 37        self.rot = trimesh.transformations.rotation_matrix(np.radians(180), [1, 0, 0])38        39        minx, miny, minz = vertices.min(axis=(0, 1))40        maxx, maxy, maxz = vertices.max(axis=(0, 1))41        minx = minx - 0.542        maxx = maxx + 0.543        minz = minz - 0.544        maxz = maxz + 0.545        46        floor = geometry.Polygon([[minx, minz], [minx, maxz], [maxx, maxz], [maxx, minz]])47        self.floor = trimesh.creation.extrude_polygon(floor, 1e-5)48        self.floor.visual.face_colors = [0, 0, 0, 0.2]49        self.floor.apply_transform(self.rot)50        self.floor_pose =np.array([[ 1, 0, 0, 0],51        [ 0, np.cos(np.pi / 2), -np.sin(np.pi / 2), miny],52        [ 0, np.sin(np.pi / 2), np.cos(np.pi / 2), 0],53        [ 0, 0, 0, 1]])54        55        c = -np.pi / 656        self.camera_pose = [[ 1, 0, 0, (minx+maxx)/2],57                        [ 0, np.cos(c), -np.sin(c), 1.5],58                        [ 0, np.sin(c), np.cos(c), max(4, minz+(1.5-miny)*2, (maxx-minx))],59                        [ 0, 0, 0, 1]60                        ]61        62    def __call__(self, vertices, camera_translation):63 64        floor_render = pyrender.Mesh.from_trimesh(self.floor, smooth=False)65        66        material = pyrender.MetallicRoughnessMaterial(67            metallicFactor=0.1,68            alphaMode='OPAQUE',69            baseColorFactor=(0.658, 0.214, 0.0114, 0.2))70        mesh = trimesh.Trimesh(vertices, self.faces)71        mesh.apply_transform(self.rot)72        mesh = pyrender.Mesh.from_trimesh(mesh, material=material)73        74        camera = pyrender.PerspectiveCamera(yfov=(np.pi / 3.0))75        76        light = pyrender.DirectionalLight(color=[1,1,1], intensity=350)77        spot_l = pyrender.SpotLight(color=np.ones(3), intensity=300.0,78                        innerConeAngle=np.pi/16, outerConeAngle=np.pi/6)79        point_l = pyrender.PointLight(color=np.ones(3), intensity=300.0)80        81        scene = pyrender.Scene(bg_color=(1.,1.,1.,0.8),ambient_light=(0.4, 0.4, 0.4))82        scene.add(floor_render, pose=self.floor_pose)83        scene.add(mesh, 'mesh')84        85        light_pose = np.eye(4)86        light_pose[:3, 3] = np.array([0, -1, 1])87        scene.add(light, pose=light_pose)88 89        light_pose[:3, 3] = np.array([0, 1, 1])90        scene.add(light, pose=light_pose)91 92        light_pose[:3, 3] = np.array([1, 1, 2])93        scene.add(light, pose=light_pose)94        95        scene.add(camera, pose=self.camera_pose)96        97        flags = pyrender.RenderFlags.RGBA | pyrender.RenderFlags.SHADOWS_DIRECTIONAL98        color, rend_depth = self.renderer.render(scene, flags=flags)99        100        return color101 102class SMPLRender():103    def __init__(self, SMPL_MODEL_DIR):104        if torch.cuda.is_available():105            self.device = torch.device("cuda")106        else:107            self.device = torch.device("cpu")108        # self.smpl = SMPL(SMPL_MODEL_DIR, batch_size=1, create_transl=False).to(self.device)109        self.smpl = smplx.create(Path(SMPL_MODEL_DIR).parent, model_type="smpl", gender="neutral", ext="npz", batch_size=1).to(self.device)110 111        self.pred_camera_t = []112        self.focal_length = 110113        114    def init_renderer(self, res, smpl_param, is_headroot=False):115        poses = smpl_param['pred_pose']116        pred_rotmats = []117        for pose in poses:118            if pose.size==72:119                pose = pose.reshape(-1,3)120                pose = RRR.from_rotvec(pose).as_matrix()121                pose = pose.reshape(1,24,3,3)122            pred_rotmats.append(torch.from_numpy(pose.astype(np.float32)[None]).to(self.device))123        pred_rotmat = torch.cat(pred_rotmats, dim=0)124 125        pred_betas = torch.from_numpy(smpl_param['pred_shape'].reshape(1, 10).astype(np.float32)).to(self.device)126        pred_root = torch.tensor(smpl_param['pred_root'].reshape(-1, 3).astype(np.float32),device=self.device)127        smpl_output = self.smpl(betas=pred_betas, body_pose=pred_rotmat[:, 1:],transl=pred_root, global_orient=pred_rotmat[:, :1], pose2rot=False)128        129        self.vertices = smpl_output.vertices.detach().cpu().numpy()130 131        pred_root = pred_root[0]132 133        if is_headroot:134            pred_root = pred_root - smpl_output.joints[0,12].detach().cpu().numpy()135 136        self.pred_camera_t.append(pred_root)137        138        self.renderer = Renderer(vertices=self.vertices, focal_length=self.focal_length,139                                img_res=(res[1], res[0]), faces=self.smpl.faces)140        141 142    def render(self, index):143        renderImg = self.renderer(self.vertices[index, ...], self.pred_camera_t)144        return renderImg145