-
Notifications
You must be signed in to change notification settings - Fork 14
/
vis_npy.py
138 lines (112 loc) · 5.56 KB
/
vis_npy.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
import os
import cv2
import torch
import numpy as np
import smplx
from smplx import SMPL
import trimesh
import pyrender
import colorsys
def estimate_focal_length(img_h, img_w):
return (img_w * img_w + img_h * img_h) ** 0.5 # fov: 55 degree
class Renderer(object):
def __init__(self, focal_length=600, img_w=512, img_h=512, faces=None,
same_mesh_color=False):
os.environ['PYOPENGL_PLATFORM'] = 'egl'
self.renderer = pyrender.OffscreenRenderer(viewport_width=img_w,
viewport_height=img_h,
point_size=1.0)
self.camera_center = [img_w // 2, img_h // 2]
self.focal_length = focal_length
self.faces = faces
self.same_mesh_color = same_mesh_color
def render_front_view(self, verts, bg_img_rgb=None, bg_color=(0, 0, 0, 0)):
# Create a scene for each image and render all meshes
scene = pyrender.Scene(bg_color=bg_color, ambient_light=np.ones(3) * 0)
# Create camera. Camera will always be at [0,0,0]
camera = pyrender.camera.IntrinsicsCamera(fx=self.focal_length, fy=self.focal_length,
cx=self.camera_center[0], cy=self.camera_center[1])
scene.add(camera, pose=np.eye(4))
# Create light source
light = pyrender.DirectionalLight(color=[1.0, 1.0, 1.0], intensity=3.0)
# for DirectionalLight, only rotation matters
light_pose = trimesh.transformations.rotation_matrix(np.radians(-45), [1, 0, 0])
scene.add(light, pose=light_pose)
light_pose = trimesh.transformations.rotation_matrix(np.radians(45), [0, 1, 0])
scene.add(light, pose=light_pose)
# Need to flip x-axis
rot = trimesh.transformations.rotation_matrix(np.radians(180), [1, 0, 0])
# multiple person
num_people = len(verts)
# for every person in the scene
for n in range(num_people):
mesh = trimesh.Trimesh(verts[n], self.faces)
mesh.apply_transform(rot)
if self.same_mesh_color:
mesh_color = colorsys.hsv_to_rgb(0.6, 0.5, 1.0)
else:
mesh_color = colorsys.hsv_to_rgb(float(n) / num_people, 0.5, 1.0)
material = pyrender.MetallicRoughnessMaterial(
metallicFactor=0.2,
alphaMode='OPAQUE',
baseColorFactor=mesh_color)
mesh = pyrender.Mesh.from_trimesh(mesh, material=material, wireframe=False)
scene.add(mesh, 'mesh')
# Alpha channel was not working previously, need to check again
# Until this is fixed use hack with depth image to get the opacity
color_rgba, depth_map = self.renderer.render(scene, flags=pyrender.RenderFlags.RGBA)
color_rgb = color_rgba[:, :, :3]
if bg_img_rgb is None:
return color_rgb
else:
mask = depth_map > 0
bg_img_rgb[mask] = color_rgb[mask]
return bg_img_rgb
def render_side_view(self, verts):
centroid = verts.mean(axis=(0, 1)) # n*6890*3 -> 3
# make the centroid at the image center (the X and Y coordinates are zeros)
centroid[:2] = 0
aroundy = cv2.Rodrigues(np.array([0, np.radians(90.), 0]))[0][np.newaxis, ...] # 1*3*3
pred_vert_arr_side = np.matmul((verts - centroid), aroundy) + centroid
side_view = self.render_front_view(pred_vert_arr_side)
return side_view
def delete(self):
"""
Need to delete before creating the renderer next time
"""
self.renderer.delete()
if __name__ == '__main__':
device = torch.device('cuda:{}'.format(0)) if torch.cuda.is_available() else torch.device('cpu')
# motion data dir
file_dir = '/share/wanghaofan/research/OpenDance-10K/data/2nd_data/clips_results/luGqEHIO_WrZoL2LvlN3z1uoLN1V_compress_L2/luGqEHIO_WrZoL2LvlN3z1uoLN1V_compress_L2_cliff_hr48_infill.npz'
# load motion
results = np.load(file_dir)
pred_betas = torch.from_numpy(results['shape']).float().to(device)
pred_rotmat = torch.from_numpy(results['pose']).float().to(device)
pred_cam_full = torch.from_numpy(results['global_t']).float().to(device)
# load smpl model
# smpl_model = smplx.create(constants.SMPL_MODEL_DIR, "smpl").to(device)
smpl_model = SMPL(model_path='./data/smpl', gender='MALE', batch_size=1).eval().to(device)
# shape params: betas=pred_betas or beta=None
pred_output = smpl_model(body_pose=pred_rotmat[:, 1:],
global_orient=pred_rotmat[:, 0:1],
pose2rot=True,
transl=pred_cam_full)
# get vertices
pred_vertices = pred_output.vertices.detach().cpu().numpy()[:150]
# setup renderer for visualization
img_h, img_w, _ = 512, 512, 3
background = np.zeros((img_h, img_w, 3))
focal_length = estimate_focal_length(img_h, img_w)
video = cv2.VideoWriter('temp.mp4',
cv2.VideoWriter_fourcc('m', 'p', '4', 'v'),
60, (img_h, img_w))
for i, mesh in enumerate(pred_vertices):
renderer = Renderer(focal_length=focal_length,
img_w=img_w, img_h=img_h,
faces=smpl_model.faces,
same_mesh_color=False)
front_view = renderer.render_front_view(mesh[None,:], bg_img_rgb=background.copy())
video.write((front_view[:, :, ::-1]*255).astype(np.uint8))
renderer.delete()
video.release()