-
Notifications
You must be signed in to change notification settings - Fork 227
/
vedo_visualizer.py
177 lines (160 loc) · 7.99 KB
/
vedo_visualizer.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
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
import os,sys
import vedo
from vedo import *
import vtk
from vtk.util.numpy_support import vtk_to_numpy
import numpy as np
import pickle
import cv2
import config
import constants
from config import args
from utils.temporal_optimization import OneEuroFilter
from utils.cam_utils import convert_cam_to_3d_trans
def convert_cam_to_stand_on_image_trans(cam, enlarge_scale):
trans_3d = convert_cam_to_3d_trans(cam)
stand_on_image_trans = np.zeros(3)
# The x-axis is supposed to be adapted to the proper scale
stand_on_image_trans[0] = trans_3d[0] * 0.3
stand_on_image_trans[1] = 0.42 #0.5 - trans_3d[1] * 0.2
#stand_on_image_trans[1] = 0.56
#stand_on_image_trans[2] = trans_3d[1] - trans_3d[2]/3 + 2.6
stand_on_image_trans[2] = trans_3d[1] * 0.4 #- trans_3d[2]/3
stand_on_image_trans *= enlarge_scale
return stand_on_image_trans
def parse_nvxia_uvmap(uvs, face):
verts_num = np.max(face) + 1
uvs_verts = np.zeros((verts_num, 2))
for uv, f in zip(uvs, face):
uvs_verts[f] = uv[:,:2]
#uvs_verts[:,1] = 1-uvs_verts[:,1]
return uvs_verts
class Vedo_visualizer(object):
def __init__(self):
if args().character == 'smpl':
self.faces = pickle.load(open(os.path.join(args().smpl_model_path,'SMPL_NEUTRAL.pkl'),'rb'), encoding='latin1')['f']
elif args().character == 'nvxia':
params_dict = np.load(os.path.join(args().nvxia_model_path, 'nvxia.npz'), allow_pickle=True)
self.faces = np.array([np.array(face) for face in params_dict['polygons']])
self.texture_file = cv2.imread(os.path.join(args().nvxia_model_path, 'Kachujin_diffuse.png'))[:,:,::-1]
self.uvs = parse_nvxia_uvmap(params_dict['uvmap'],self.faces)
self.scene_bg_color = [240,255,255]
self.default_camera={'pos':{'far':(0,800,1000), 'close':(0,200,800)}[args().soi_camera]}
self.light = Light([0,800,1000], c='white')
vedo.settings.screeshotLargeImage = True
vedo.settings.screeshotScale = 2
def plot_multi_meshes_batch(self, vertices, cam_params, meta_data, reorganize_idx, save_img=True, interactive_show=False):
result_imgs = []
for inds, img_id in enumerate(np.unique(reorganize_idx)):
single_img_verts_inds = np.array(np.where(reorganize_idx==img_id)[0])
save_name = os.path.join(args().output_dir, '3D_meshes-'+os.path.basename(meta_data['imgpath'][single_img_verts_inds[0]]+'.jpg'))
result_img = self.plot_multi_meshes(vertices[single_img_verts_inds].detach().cpu().numpy(), \
cam_params[single_img_verts_inds].detach().cpu().numpy(), meta_data['image'][single_img_verts_inds[0]].cpu().numpy().astype(np.uint8),\
save_name=save_name, interactive_show=interactive_show)
result_imgs.append(save_name)
if save_img:
cv2.imwrite(save_name, result_img[:,:,::-1])
return result_imgs
def plot_multi_meshes(self, vertices, cam_params, img, mesh_colors=None, save_name=None, interactive_show=False):
plt = Plotter(bg=[240,255,255], axes=0, offscreen=not interactive_show)
h,w = img.shape[:2]
pic = Picture(img)
pic.rotateX(-90).z(h//2).x(-w//2)
plt += pic
vertices_vis = []
enlarge_scale = max(h,w)/2.6
for inds, (vert, cam) in enumerate(zip(vertices, cam_params)):
trans_3d = convert_cam_to_stand_on_image_trans(cam, enlarge_scale)
vert[:,1:] *= -1
vert = vert * enlarge_scale / 2.
vert += trans_3d[None]
vertices_vis.append(vert)
vertices_vis = np.stack(vertices_vis, 0)
visulize_list = []
for inds, vert in enumerate(vertices_vis):
mesh = Mesh([vert, self.faces])
if args().character == 'smpl':
mesh = mesh.c([255,255,255]).smooth(niter=20)
if mesh_colors is not None:
mesh.c(mesh_colors[inds].astype(np.uint8))
elif args().character == 'nvxia':
mesh.texture(self.texture_file,tcoords=self.uvs).smooth(niter=20)#.lighting('glossy')
visulize_list.append(mesh)
plt += visulize_list
plt += self.light
image_result = plt.show(camera=self.default_camera) #elevation=10,azimuth=0,,bg=self.bg_path
result_img = plt.topicture(scale=2)
rows, cols, _ = result_img._data.GetDimensions()
vtkimage = result_img._data.GetPointData().GetScalars()
image_result = vtk_to_numpy(vtkimage).reshape((rows, cols, 3))
image_result = image_result[::-1]
#result_img.write(save_name)
#screenshot(save_name) #returnNumpy=True
plt.close()
return image_result
# class vedo_visualizer(object):
# def __init__(self):
# smpl_param_dict = pickle.load(open(os.path.join(args().smpl_model_path,'smpl','SMPL_NEUTRAL.pkl'),'rb'), encoding='latin1')
# self.faces = smpl_param_dict['f']
# #self.verts_mean = smpl_param_dict['v_template']
# #self.load_smpl_tex()
# #self.load_smpl_vtk()
# if args().webcam_mesh_color == 'female_tex':
# self.uv_map = np.load(args().smpl_uvmap)
# self.texture_file = args().smpl_female_texture
# elif args().webcam_mesh_color == 'male_tex':
# self.uv_map = np.load(args().smpl_uvmap)
# self.texture_file = args().smpl_male_texture
# else:
# self.mesh_color = np.array(constants.mesh_color_dict[args().webcam_mesh_color])/255.
# #self.mesh = self.create_smpl_mesh(self.verts_mean)
# self.mesh_smoother = OneEuroFilter(4.0, 0.0)
# self.vp = Plotter(title='Predicted 3D mesh',interactive=0)#
# self.vp_2d = Plotter(title='Input frame',interactive=0)
# #show(self.mesh, axes=1, viewup="y", interactive=0)
# def load_smpl_tex(self):
# import scipy.io as sio
# UV_info = sio.loadmat(os.path.join(args().smpl_model_path,'smpl','UV_Processed.mat'))
# self.vertex_reorder = UV_info['All_vertices'][0]-1
# self.faces = UV_info['All_Faces']-1
# self.uv_map = np.concatenate([UV_info['All_U_norm'], UV_info['All_V_norm']],1)
# def run(self, verts,frame):
# verts[:,1:] = verts[:,1:]*-1
# verts = self.mesh_smoother.process(verts)
# #verts = verts[self.vertex_reorder]
# #self.mesh.points(verts)
# mesh = self.create_smpl_mesh(verts)
# self.vp.show(mesh,viewup=np.array([0,-1,0]))
# self.vp_2d.show(Picture(frame))
# return False
# def create_smpl_mesh(self, verts):
# mesh = Mesh([verts, self.faces])
# mesh.texture(self.texture_file,tcoords=self.uv_map)
# mesh = self.collapse_triangles_with_large_gradient(mesh)
# mesh.computeNormals()
# return mesh
# def collapse_triangles_with_large_gradient(self, mesh, threshold=4.0):
# points = mesh.points()
# new_points = np.array(points)
# mesh_vtk = Mesh(os.path.join(args().smpl_model_path,'smpl_male.vtk'), c='w').texture(self.texture_file).lw(0.1)
# grad = mesh_vtk.gradient("tcoords")
# ugrad, vgrad = np.split(grad, 2, axis=1)
# ugradm, vgradm = mag(ugrad), mag(vgrad)
# gradm = np.log(ugradm*ugradm + vgradm*vgradm)
# largegrad_ids = np.arange(mesh.N())[gradm>threshold]
# for f in mesh.faces():
# if np.isin(f, largegrad_ids).all():
# id1, id2, id3 = f
# uv1, uv2, uv3 = self.uv_map[f]
# d12 = mag(uv1-uv2)
# d23 = mag(uv2-uv3)
# d31 = mag(uv3-uv1)
# idm = np.argmin([d12, d23, d31])
# if idm == 0: # d12, collapse segment to pt3
# new_points[id1] = new_points[id3]
# new_points[id2] = new_points[id3]
# elif idm == 1: # d23
# new_points[id2] = new_points[id1]
# new_points[id3] = new_points[id1]
# mesh.points(new_points)
# return mesh