From 88c828e69ce9fc4c7a0d36ac7339d9e1ce200164 Mon Sep 17 00:00:00 2001 From: shuaiqing Date: Mon, 16 Oct 2023 19:49:00 +0800 Subject: [PATCH] add vis_smpl_by_open3dgui --- apps/vis3d/vis_smpl.py | 208 +++++++++++++++++++++++++++++++++ easymocap/vis3d/basegui.py | 227 +++++++++++++++++++++++++++++++++++++ 2 files changed, 435 insertions(+) create mode 100644 apps/vis3d/vis_smpl.py create mode 100644 easymocap/vis3d/basegui.py diff --git a/apps/vis3d/vis_smpl.py b/apps/vis3d/vis_smpl.py new file mode 100644 index 0000000..9b5853a --- /dev/null +++ b/apps/vis3d/vis_smpl.py @@ -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() \ No newline at end of file diff --git a/easymocap/vis3d/basegui.py b/easymocap/vis3d/basegui.py new file mode 100644 index 0000000..98af496 --- /dev/null +++ b/easymocap/vis3d/basegui.py @@ -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 ') + 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) \ No newline at end of file