add vis_smpl_by_open3dgui
This commit is contained in:
shuaiqing
parent
2d8584f643
commit
88c828e69c
208
apps/vis3d/vis_smpl.py
Normal file
208
apps/vis3d/vis_smpl.py
Normal file
@ -0,0 +1,208 @@
|
||||
import open3d as o3d
|
||||
from easymocap.config.baseconfig import load_object_from_cmd
|
||||
from easymocap.mytools.reader import read_smpl
|
||||
from easymocap.visualize.o3dwrapper import Vector3dVector, Vector3iVector
|
||||
from easymocap.vis3d.basegui import BaseWindow
|
||||
from easymocap.mytools.vis_base import generate_colorbar
|
||||
from easymocap.mytools.file_utils import myarray2string, write_smpl
|
||||
import open3d.visualization.gui as gui
|
||||
import os
|
||||
import numpy as np
|
||||
class SMPLControl(BaseWindow):
|
||||
def __init__(self, cfg_model, opts_model, out) -> None:
|
||||
super().__init__(out, panel_rate=0.4)
|
||||
self.body_model = load_object_from_cmd(cfg_model, opts_model)
|
||||
if args.param_path is not None and os.path.exists(args.param_path):
|
||||
from easymocap.mytools.reader import read_smpl
|
||||
params = read_smpl(args.param_path)[0]
|
||||
else:
|
||||
params = self.body_model.init_params(nFrames=1)
|
||||
self._params = params
|
||||
vertices = self.body_model(return_verts=True, return_tensor=False, **params)
|
||||
joints = self.body_model(return_verts=False, return_tensor=False, return_smpl_joints=True, **self._params)[0]
|
||||
mesh = o3d.geometry.TriangleMesh()
|
||||
mesh.vertices = Vector3dVector(vertices[0])
|
||||
mesh.triangles = Vector3iVector(self.body_model.faces)
|
||||
# 使用blendweight计算颜色
|
||||
if True:
|
||||
weights = self.body_model.weights.cpu().numpy()
|
||||
nJoints = weights.shape[-1]
|
||||
colorbar = np.array(generate_colorbar(nJoints))[:, ::-1]/255.
|
||||
colors = weights @ colorbar
|
||||
mesh.vertex_colors = Vector3dVector(colors)
|
||||
mat = self.get_default_mat(color=[1,1,1,0.5])
|
||||
else:
|
||||
mat = self.get_default_mat()
|
||||
mesh.compute_vertex_normals()
|
||||
self.add_geometry('smpl', mesh, mat)
|
||||
colorbar = np.array(generate_colorbar(joints.shape[0]))[:, ::-1]/255.
|
||||
for nj in range(joints.shape[0]):
|
||||
sphere = o3d.geometry.TriangleMesh.create_sphere(
|
||||
radius=0.01, resolution=20)
|
||||
sphere.translate(joints[nj])
|
||||
sphere.compute_vertex_normals()
|
||||
sphere.paint_uniform_color(colorbar[nj])
|
||||
# self.add_geometry('joint{}'.format(nj), sphere, mat)
|
||||
# l = self.scene.add_3d_label(joints[nj], "{}".format(nj))
|
||||
self.add_control_smpl(self.panel, self.em)
|
||||
self.cnt = 0
|
||||
|
||||
def update_smpl(self):
|
||||
self._params['id'] = 0
|
||||
write_smpl('/tmp/smpl.json', [self._params])
|
||||
self._params.pop('id')
|
||||
vertices = self.body_model(return_verts=True, return_tensor=False, **self._params)
|
||||
joints = self.body_model(return_verts=False, return_tensor=False, return_smpl_joints=True, **self._params)[0]
|
||||
for nj in range(joints.shape[0]):
|
||||
break
|
||||
sphere = o3d.geometry.TriangleMesh.create_sphere(
|
||||
radius=0.03, resolution=20)
|
||||
sphere.translate(joints[nj])
|
||||
self.factory['joint{}'.format(nj)]['geom'].vertices = sphere.vertices
|
||||
self.update_geometry('joint{}'.format(nj))
|
||||
self.factory['smpl']['geom'].vertices = Vector3dVector(vertices[0])
|
||||
self.factory['smpl']['geom'].compute_vertex_normals()
|
||||
self.update_geometry('smpl')
|
||||
|
||||
def reset_params(self, ):
|
||||
for key, val in self._params.items():
|
||||
val[:] = 0.
|
||||
for slider in self.slider_dict[key]:
|
||||
slider.double_value = 0.
|
||||
print('Reset {}'.format(key))
|
||||
self.update_smpl()
|
||||
|
||||
def export_params(self, ):
|
||||
for key, val in self._params.items():
|
||||
print('{}: {}'.format(key, myarray2string(self._params[key])))
|
||||
self._params['id'] = 0
|
||||
write_smpl('/tmp/smpl.json', [self._params])
|
||||
|
||||
def read_params(self, filename):
|
||||
datas = read_smpl(filename)[0]
|
||||
for key in self._params.keys():
|
||||
if key in datas.keys():
|
||||
self._params[key] = datas[key]
|
||||
self.update_smpl()
|
||||
|
||||
def import_params(self, ):
|
||||
dlg = gui.FileDialog(gui.FileDialog.OPEN, "Choose file to load",
|
||||
self.window.theme)
|
||||
dlg.add_filter(
|
||||
".json",
|
||||
"SMPL parameters")
|
||||
# A file dialog MUST define on_cancel and on_done functions
|
||||
dlg.set_on_cancel(self._on_file_dialog_cancel)
|
||||
dlg.set_on_done(self._on_load_dialog_done)
|
||||
self.window.show_dialog(dlg)
|
||||
|
||||
def _on_file_dialog_cancel(self):
|
||||
self.window.close_dialog()
|
||||
|
||||
def _on_load_dialog_done(self, filename):
|
||||
self.window.close_dialog()
|
||||
self.read_params(filename)
|
||||
|
||||
def create_callback(self, key, index):
|
||||
def change(value):
|
||||
self._params[key][0, index] = value
|
||||
self.update_smpl()
|
||||
if self.out is not None:
|
||||
self.scene.scene.scene.render_to_image(self.capture_callback)
|
||||
return change
|
||||
|
||||
def add_control_smpl(self, panel, em):
|
||||
# Rh, Th
|
||||
print(self._params.keys())
|
||||
# add reset
|
||||
reset_button = gui.Button("reset")
|
||||
reset_button.set_on_clicked(self.reset_params)
|
||||
panel.add_child(reset_button)
|
||||
import_button = gui.Button("import")
|
||||
import_button.set_on_clicked(self.import_params)
|
||||
panel.add_child(import_button)
|
||||
export_button = gui.Button("export")
|
||||
export_button.set_on_clicked(self.export_params)
|
||||
panel.add_child(export_button)
|
||||
self.slider_dict = {}
|
||||
for key in ['Rh', 'Th']:
|
||||
col = gui.CollapsableVert(key, 0.33 * em,
|
||||
gui.Margins(em, 0, 0, 0))
|
||||
for i in range(self._params[key].shape[1]//3):
|
||||
grid = gui.VGrid(5, 10)
|
||||
_label = gui.Label(key + '_{}'.format(i))
|
||||
grid.add_child(_label)
|
||||
self.slider_dict[key] = []
|
||||
for nj in range(3):
|
||||
slider = gui.Slider(gui.Slider.DOUBLE)
|
||||
slider.set_limits(-4., 4.)
|
||||
slider.double_value = self._params[key][0, nj+3*i]
|
||||
slider.set_on_value_changed(self.create_callback(key, nj+3*i))
|
||||
self.slider_dict[key].append(slider)
|
||||
grid.add_child(slider)
|
||||
col.add_child(grid)
|
||||
panel.add_child(col)
|
||||
for key in ['shapes', 'expression']:
|
||||
if key not in self._params.keys():
|
||||
continue
|
||||
nShape = self._params[key].shape[-1]
|
||||
col = gui.CollapsableVert(key, 0.33 * em,
|
||||
gui.Margins(em, 0, 0, 0))
|
||||
grid = gui.VGrid(2, 10)
|
||||
self.slider_dict[key] = []
|
||||
for nj in range(nShape):
|
||||
_label = gui.Label(key + '_{}'.format(nj))
|
||||
grid.add_child(_label)
|
||||
slider = gui.Slider(gui.Slider.DOUBLE)
|
||||
slider.set_limits(-4., 4.)
|
||||
slider.double_value = self._params[key][0, nj]
|
||||
slider.set_on_value_changed(self.create_callback(key, nj))
|
||||
self.slider_dict[key].append(slider)
|
||||
grid.add_child(slider)
|
||||
col.add_child(grid)
|
||||
panel.add_child(col)
|
||||
for key in ['poses', 'handl', 'handr']:
|
||||
if key not in self._params.keys():
|
||||
continue
|
||||
nShape = self._params[key].shape[-1]//3
|
||||
col = gui.CollapsableVert(key, 0.33 * em,
|
||||
gui.Margins(em, 0, 0, 0))
|
||||
grid = gui.VGrid(4, 10)
|
||||
self.slider_dict[key] = []
|
||||
for nj in range(nShape):
|
||||
_label = gui.Label(key + '_{}'.format(nj))
|
||||
grid.add_child(_label)
|
||||
for i in range(3):
|
||||
slider = gui.Slider(gui.Slider.DOUBLE)
|
||||
slider.set_limits(-3., 3.)
|
||||
slider.double_value = self._params[key][0, 3*nj+i]
|
||||
slider.set_on_value_changed(self.create_callback(key, 3*nj+i))
|
||||
self.slider_dict[key].append(slider)
|
||||
grid.add_child(slider)
|
||||
col.add_child(grid)
|
||||
panel.add_child(col)
|
||||
|
||||
if __name__ == "__main__":
|
||||
import argparse
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--cfg', type=str,
|
||||
default='config/model/smpl_neutral.yml')
|
||||
parser.add_argument('--opts', type=str,
|
||||
default=[], nargs="+")
|
||||
|
||||
parser.add_argument('--key', type=str,
|
||||
default='poses')
|
||||
parser.add_argument('--max', type=float, default=1.57)
|
||||
parser.add_argument('--param_path', type=str, default=None)
|
||||
parser.add_argument('--out', type=str, default=None)
|
||||
parser.add_argument('--num', type=int, default=50)
|
||||
parser.add_argument('--start', type=int, default=0)
|
||||
parser.add_argument('--one', action='store_true')
|
||||
parser.add_argument('--debug', action='store_true')
|
||||
args = parser.parse_args()
|
||||
|
||||
gui.Application.instance.initialize()
|
||||
|
||||
w = SMPLControl(args.cfg, args.opts, out=args.out)
|
||||
# Run the event loop. This will not return until the last window is closed.
|
||||
gui.Application.instance.run()
|
||||
227
easymocap/vis3d/basegui.py
Normal file
227
easymocap/vis3d/basegui.py
Normal file
@ -0,0 +1,227 @@
|
||||
import open3d as o3d
|
||||
import os
|
||||
from os.path import join
|
||||
import open3d.visualization.gui as gui
|
||||
import open3d.visualization.rendering as rendering
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
class BaseWindow: # this window is the basic of Open3D new render style
|
||||
colormap = {
|
||||
0: (94/255, 124/255, 226/255, 1.), # 青色
|
||||
1: (255/255, 200/255, 87/255, 1.), # yellow
|
||||
2: (74/255., 189/255., 172/255., 1.), # green
|
||||
3: (8/255, 76/255, 97/255, 1.), # blue
|
||||
4: (219/255, 58/255, 52/255, 1.), # red
|
||||
5: (77/255, 40/255, 49/255, 1.), # brown
|
||||
1000: [1., 0., 0., 1.],
|
||||
1001: [0., 1., 0., 1.],
|
||||
1002: [0., 0., 1., 1.],
|
||||
}
|
||||
def __init__(self, out, panel_rate=0.2) -> None:
|
||||
self.out = out
|
||||
if out is not None:
|
||||
os.makedirs(out, exist_ok=True)
|
||||
self.window = gui.Application.instance.create_window(
|
||||
"EasyMocap Visualization", 1920, 1080)
|
||||
w = self.window # for more concise code
|
||||
em = w.theme.font_size
|
||||
self.em = em
|
||||
# set the frame region
|
||||
self.panel_rate = panel_rate
|
||||
self.window.set_on_layout(self._on_layout)
|
||||
self.window.set_on_close(self._on_close)
|
||||
# create scene
|
||||
self.scene = self.add_scene(w)
|
||||
self.window.add_child(self.scene)
|
||||
self.panel = self.add_panel(w, em)
|
||||
self.factory = {}
|
||||
self._render_cnt = 0
|
||||
self.is_done = False
|
||||
|
||||
def _on_layout(self, layout_context):
|
||||
contentRect = self.window.content_rect
|
||||
panel_width = contentRect.width * self.panel_rate
|
||||
self.scene.frame = gui.Rect(contentRect.x, contentRect.y,
|
||||
contentRect.width - panel_width,
|
||||
contentRect.height)
|
||||
self.panel.frame = gui.Rect(self.scene.frame.get_right(),
|
||||
contentRect.y, panel_width,
|
||||
contentRect.height)
|
||||
|
||||
def _on_close(self):
|
||||
self.is_done = True
|
||||
return True # False would cancel the close
|
||||
|
||||
def add_scene(self, window):
|
||||
scene = gui.SceneWidget()
|
||||
scene.scene = rendering.Open3DScene(window.renderer)
|
||||
scene.scene.set_background([1, 1, 1, 1])
|
||||
scene.scene.scene.set_sun_light(
|
||||
[-1, -1, -1], # direction
|
||||
[1, 1, 1], # color
|
||||
100000) # intensity
|
||||
scene.scene.scene.enable_sun_light(True)
|
||||
scene.scene.show_skybox(True)
|
||||
# scene.scene.show_ground_plane(True)
|
||||
scene.scene.show_axes(True)
|
||||
bbox = o3d.geometry.AxisAlignedBoundingBox([-5, -5, -5],
|
||||
[5, 5, 5])
|
||||
|
||||
if False:
|
||||
scene.setup_camera(60, bbox, [0, 0, 0])
|
||||
# fov, bounds, center_of_rotation
|
||||
else:
|
||||
K = np.array([
|
||||
1000., 0., 500.,
|
||||
0., 1000., 500.,
|
||||
0., 0., 1.
|
||||
]).reshape(3, 3)
|
||||
T = np.array([0., 0., 10.]).reshape(3, 1)
|
||||
RT = np.eye(4)
|
||||
R = cv2.Rodrigues(np.array([1., 0., 0.])*(np.pi/6 +np.pi/2))[0]
|
||||
RT[:3, 3:] = T
|
||||
RT[:3, :3] = R
|
||||
scene.setup_camera(K, RT, 1000, 1000, bbox)
|
||||
return scene
|
||||
|
||||
def add_color(self, name, init, callback):
|
||||
_col = gui.ColorEdit()
|
||||
_col.set_on_value_changed(callback)
|
||||
_col.color_value.set_color(init[0], init[1], init[2])
|
||||
label = gui.Label(name)
|
||||
return label, _col
|
||||
|
||||
def add_vec3d(self, name, init, callback):
|
||||
label = gui.Label(name)
|
||||
# Create a widget for showing/editing a 3D vector
|
||||
vedit = gui.VectorEdit()
|
||||
vedit.vector_value = init
|
||||
vedit.set_on_value_changed(callback)
|
||||
return label, vedit
|
||||
|
||||
def add_panel(self, window, em):
|
||||
panel = gui.Vert(20,
|
||||
gui.Margins(0.5 * em, 0.5 * em, 0.5 * em, 0.5 * em))
|
||||
window.add_child(panel)
|
||||
return panel
|
||||
|
||||
def get_default_mat(self, pid=-1, color=[0., 1., 1., 1.], shader='defaultLit'):
|
||||
mat = rendering.MaterialRecord()
|
||||
if pid != -1 and pid in list(self.colormap.keys()):
|
||||
color = self.colormap[pid]
|
||||
mat.base_color = color
|
||||
mat.shader = shader
|
||||
return mat
|
||||
|
||||
def remove_geometry(self, name):
|
||||
if name in self.factory.keys():
|
||||
self.scene.scene.remove_geometry(name)
|
||||
|
||||
def add_geometry(self, name, geom, mat=None):
|
||||
if name in self.factory.keys():
|
||||
self.scene.scene.remove_geometry(name)
|
||||
self.factory[name] = {
|
||||
'geom': geom,
|
||||
'mat': mat
|
||||
}
|
||||
self.scene.scene.add_geometry(name, geom, mat)
|
||||
|
||||
@staticmethod
|
||||
def _chessboard_params():
|
||||
params = {
|
||||
'center': [0, 0, -0.1],
|
||||
'xdir': [1, 0, 0],
|
||||
'ydir': [0, 1, 0],
|
||||
'step': 1,
|
||||
'xrange': 5,
|
||||
'yrange': 5,
|
||||
'white': [1., 1., 1.],
|
||||
'black': [0.5, 0.5, 0.5],
|
||||
'two_sides': True,
|
||||
}
|
||||
return params
|
||||
|
||||
def add_camera(self, path=None, cameras=None):
|
||||
from ..visualize.o3dwrapper import create_camera
|
||||
mesh = create_camera(path=path, cameras=cameras)
|
||||
self.add_geometry('camera', mesh, self.get_default_mat(color=[1., 1., 1., 1.]))
|
||||
|
||||
def add_chessboard(self, params):
|
||||
from ..visualize.o3dwrapper import create_ground
|
||||
mesh = create_ground(**params)
|
||||
self.add_geometry('chessboard', mesh, self.get_default_mat(color=[1., 1., 1., 1.]))
|
||||
|
||||
def update_geometry(self, name):
|
||||
self.scene.scene.remove_geometry(name)
|
||||
self.scene.scene.add_geometry(name, self.factory[name]['geom'], self.factory[name]['mat'])
|
||||
|
||||
def capture_callback(self, image):
|
||||
if self.out is None:
|
||||
print('[Vis] Please set output folder by --out <path>')
|
||||
return 0
|
||||
outname = join(self.out, '{:06d}.jpg'.format(self._render_cnt))
|
||||
img = image
|
||||
quality = 9 # png
|
||||
if outname.endswith(".jpg"):
|
||||
quality = 100
|
||||
o3d.io.write_image(outname, img, quality)
|
||||
print('[Vis] render image to {}'.format(outname))
|
||||
self._render_cnt += 1
|
||||
|
||||
def add_chessboard_widget(self, param):
|
||||
chessboard = gui.CollapsableVert("Chessboard setting", 10,
|
||||
gui.Margins(self.em, 0, 0, 0))
|
||||
|
||||
grid = gui.VGrid(2, 10)
|
||||
label0, widget0 = self.add_color(
|
||||
'black', param['black'], self._on_chess_col_0)
|
||||
label1, widget1 = self.add_color(
|
||||
'white', param['white'], self._on_chess_col_1)
|
||||
grid.add_child(label0)
|
||||
grid.add_child(widget0)
|
||||
grid.add_child(label1)
|
||||
grid.add_child(widget1)
|
||||
# 增加棋盘格的范围选项
|
||||
label_range, widget_range = self.add_vec3d(
|
||||
'ranges',
|
||||
[param['xrange'], param['yrange'], param['step']],
|
||||
self._on_chess_range)
|
||||
grid.add_child(label_range)
|
||||
grid.add_child(widget_range)
|
||||
# 增加棋盘格的中心选项
|
||||
label_center, widget_center = self.add_vec3d(
|
||||
'center',
|
||||
param['center'],
|
||||
self._on_chess_center)
|
||||
grid.add_child(label_center)
|
||||
grid.add_child(widget_center)
|
||||
chessboard.add_child(grid)
|
||||
self.panel.add_child(chessboard)
|
||||
|
||||
def _on_chess_center(self, ranges):
|
||||
for i in range(3):
|
||||
self.param_chessboard['center'][i] = float(ranges[i])
|
||||
self.add_chessboard(self.param_chessboard)
|
||||
|
||||
def _on_chess_range(self, ranges):
|
||||
self.param_chessboard['xrange'] = int(ranges[0])
|
||||
self.param_chessboard['yrange'] = int(ranges[1])
|
||||
self.param_chessboard['step'] = float(ranges[2])
|
||||
self.add_chessboard(self.param_chessboard)
|
||||
|
||||
def _on_chess_col_0(self, new_color):
|
||||
color = [
|
||||
new_color.red, new_color.green,
|
||||
new_color.blue
|
||||
]
|
||||
self.param_chessboard['black'] = color
|
||||
self.add_chessboard(self.param_chessboard)
|
||||
|
||||
def _on_chess_col_1(self, new_color):
|
||||
color = [
|
||||
new_color.red, new_color.green,
|
||||
new_color.blue
|
||||
]
|
||||
self.param_chessboard['white'] = color
|
||||
self.add_chessboard(self.param_chessboard)
|
||||
Loading…
Reference in New Issue
Block a user