|
|
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="pkl", |
|
|
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 |
|
|
|
