add vis_smpl_by_open3dgui

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()

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 <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)