| | import os |
| | import torch |
| | import numpy as np |
| | import cv2 |
| |
|
| | import matplotlib.pyplot as plt |
| | import glob |
| | import pickle |
| | import pyrender |
| | import trimesh |
| | import smplx |
| | from pathlib import Path |
| | 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 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=2.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") |
| |
|
| | self.rot = trimesh.transformations.rotation_matrix(np.radians(180), [1, 0, 0]) |
| | |
| | minx, miny, minz = vertices.min(axis=(0, 1)) |
| | maxx, maxy, maxz = vertices.max(axis=(0, 1)) |
| | minx = minx - 0.5 |
| | maxx = maxx + 0.5 |
| | minz = minz - 0.5 |
| | maxz = maxz + 0.5 |
| | |
| | floor = geometry.Polygon([[minx, minz], [minx, maxz], [maxx, maxz], [maxx, minz]]) |
| | self.floor = trimesh.creation.extrude_polygon(floor, 1e-5) |
| | self.floor.visual.face_colors = [0, 0, 0, 0.2] |
| | self.floor.apply_transform(self.rot) |
| | self.floor_pose =np.array([[ 1, 0, 0, 0], |
| | [ 0, np.cos(np.pi / 2), -np.sin(np.pi / 2), miny], |
| | [ 0, np.sin(np.pi / 2), np.cos(np.pi / 2), 0], |
| | [ 0, 0, 0, 1]]) |
| | |
| | c = -np.pi / 6 |
| | self.camera_pose = [[ 1, 0, 0, (minx+maxx)/2], |
| | [ 0, np.cos(c), -np.sin(c), 1.5], |
| | [ 0, np.sin(c), np.cos(c), max(4, minz+(1.5-miny)*2, (maxx-minx))], |
| | [ 0, 0, 0, 1] |
| | ] |
| | |
| | def __call__(self, vertices, camera_translation): |
| |
|
| | floor_render = pyrender.Mesh.from_trimesh(self.floor, smooth=False) |
| | |
| | material = pyrender.MetallicRoughnessMaterial( |
| | metallicFactor=0.1, |
| | 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) |
| | |
| | camera = pyrender.PerspectiveCamera(yfov=(np.pi / 3.0)) |
| | |
| | light = pyrender.DirectionalLight(color=[1,1,1], intensity=350) |
| | spot_l = pyrender.SpotLight(color=np.ones(3), intensity=300.0, |
| | innerConeAngle=np.pi/16, outerConeAngle=np.pi/6) |
| | point_l = pyrender.PointLight(color=np.ones(3), intensity=300.0) |
| | |
| | scene = pyrender.Scene(bg_color=(1.,1.,1.,0.8),ambient_light=(0.4, 0.4, 0.4)) |
| | scene.add(floor_render, pose=self.floor_pose) |
| | scene.add(mesh, 'mesh') |
| | |
| | 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) |
| | |
| | scene.add(camera, pose=self.camera_pose) |
| | |
| | flags = pyrender.RenderFlags.RGBA | pyrender.RenderFlags.SHADOWS_DIRECTIONAL |
| | color, rend_depth = self.renderer.render(scene, flags=flags) |
| | |
| | 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 = smplx.create(Path(SMPL_MODEL_DIR).parent, model_type="smpl", gender="neutral", ext="npz", batch_size=1).to(self.device) |
| |
|
| | self.pred_camera_t = [] |
| | self.focal_length = 110 |
| | |
| | def init_renderer(self, res, smpl_param, is_headroot=False): |
| | poses = smpl_param['pred_pose'] |
| | pred_rotmats = [] |
| | for pose in poses: |
| | if pose.size==72: |
| | pose = pose.reshape(-1,3) |
| | pose = RRR.from_rotvec(pose).as_matrix() |
| | pose = pose.reshape(1,24,3,3) |
| | pred_rotmats.append(torch.from_numpy(pose.astype(np.float32)[None]).to(self.device)) |
| | pred_rotmat = torch.cat(pred_rotmats, dim=0) |
| |
|
| | pred_betas = torch.from_numpy(smpl_param['pred_shape'].reshape(1, 10).astype(np.float32)).to(self.device) |
| | pred_root = torch.tensor(smpl_param['pred_root'].reshape(-1, 3).astype(np.float32),device=self.device) |
| | smpl_output = self.smpl(betas=pred_betas, body_pose=pred_rotmat[:, 1:],transl=pred_root, global_orient=pred_rotmat[:, :1], pose2rot=False) |
| | |
| | self.vertices = smpl_output.vertices.detach().cpu().numpy() |
| |
|
| | pred_root = pred_root[0] |
| |
|
| | if is_headroot: |
| | pred_root = pred_root - smpl_output.joints[0,12].detach().cpu().numpy() |
| |
|
| | self.pred_camera_t.append(pred_root) |
| | |
| | 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, ...], self.pred_camera_t) |
| | return renderImg |
| |
|
