🚀 [vis] add the real-time visualization

2021-06-04 17:12:30 +08:00 · 2021-06-04 17:12:30 +08:00 · 584ba2c1e8
commit 584ba2c1e8
parent fa4bd6ddaa
21 changed files with 1258 additions and 52 deletions
--- a/Readme.md
+++ b/Readme.md
@ -2,7 +2,7 @@
 * @Date: 2021-01-13 20:32:12
 * @Author: Qing Shuai
 * @LastEditors: Qing Shuai
- * @LastEditTime: 2021-04-14 16:00:04
+ * @LastEditTime: 2021-06-04 17:12:01
 * @FilePath: /EasyMocapRelease/Readme.md
 -->

@ -74,9 +74,11 @@ This project is used by many other projects:
 - [Pose guided synchronization](./doc/todo.md) (comming soon)
 - [Annotator](apps/calibration/Readme.md): a simple GUI annotator based on OpenCV
 - [Exporting of multiple data formats(bvh, asf/amc, ...)](./doc/02_output.md)
+- [Real-time visualization](./doc/realtime_visualization.md)

 ## Updates

+- 06/04/2021: The **real-time 3D visualization** part is released!
 - 04/12/2021: Mirrored-Human part is released. We also release the calibration tool and the annotator.

 ## Installation
--- a/apps/vis/vis_client.py
+++ b/apps/vis/vis_client.py
@ -0,0 +1,54 @@
+'''
+  @ Date: 2021-05-24 18:57:48
+  @ Author: Qing Shuai
+  @ LastEditors: Qing Shuai
+  @ LastEditTime: 2021-06-04 16:43:00
+  @ FilePath: /EasyMocapRelease/apps/vis/vis_client.py
+'''
+import socket
+import time
+from easymocap.socket.base_client import BaseSocketClient
+import os
+
+def send_rand(client):
+    import numpy as np
+    for _ in range(1000):
+        k3d = np.random.rand(25, 4)
+        data = [
+            {
+                'id': 0,
+                'keypoints3d': k3d
+            }
+        ]
+        client.send(data)
+        time.sleep(0.005)
+    client.close()
+
+def send_dir(client, path):
+    from os.path import join
+    from glob import glob
+    from tqdm import tqdm
+    from easymocap.mytools.reader import read_keypoints3d
+    results = sorted(glob(join(path, '*.json')))
+    for result in tqdm(results):
+        data = read_keypoints3d(result)
+        client.send(data)
+        time.sleep(0.005)
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--host', type=str, default='auto')
+    parser.add_argument('--port', type=int, default=9999)
+    parser.add_argument('--path', type=str, default=None)
+    parser.add_argument('--debug', action='store_true')
+    args = parser.parse_args()
+
+    if args.host == 'auto':
+        args.host = socket.gethostname()
+    client = BaseSocketClient(args.host, args.port)
+
+    if args.path is not None and os.path.isdir(args.path):
+        send_dir(client, args.path)
+    else:
+        send_rand(client)
--- a/apps/vis/vis_server.py
+++ b/apps/vis/vis_server.py
@ -0,0 +1,19 @@
+'''
+  @ Date: 2021-05-24 18:51:58
+  @ Author: Qing Shuai
+  @ LastEditors: Qing Shuai
+  @ LastEditTime: 2021-06-04 17:00:15
+  @ FilePath: /EasyMocapRelease/apps/vis/vis_server.py
+'''
+# socket server for 3D visualization
+from easymocap.socket.o3d import VisOpen3DSocket
+from easymocap.config.vis_socket import Config
+
+def main(cfg):
+    server = VisOpen3DSocket(cfg.host, cfg.port, cfg)
+    while True:
+        server.update()
+
+if __name__ == "__main__":
+    cfg = Config.load_from_args()
+    main(cfg)
--- a/config/vis/o3d_scene.yml
+++ b/config/vis/o3d_scene.yml
@ -0,0 +1,39 @@
+host: 'auto'
+port: 9999
+
+width: 1920
+height: 1080
+
+max_human: 5
+track: True
+block: True # block visualization or not, True for visualize each frame, False in realtime applications
+debug: False
+write: False
+out: 'none'
+
+body_model:
+  module: "easymocap.visualize.skelmodel.SkelModel"
+  args:
+    body_type: "body25"
+    joint_radius: 0.02
+    gender: "neutral"
+    model_type: "smpl"
+
+scene:
+  "easymocap.visualize.o3dwrapper.create_coord":
+    camera: [0, 0, 0]
+    radius: 1
+  "easymocap.visualize.o3dwrapper.create_bbox":
+    min_bound: [-3, -3, 0]
+    max_bound: [3, 3, 2]
+    flip: False
+  "easymocap.visualize.o3dwrapper.create_ground":
+    center: [0, 0, 0]
+    xdir: [1, 0, 0]
+    ydir: [0, 1, 0]
+    step: 1
+    xrange: 3
+    yrange: 3
+    white: [1., 1., 1.]
+    black: [0.,0.,0.]
+    two_sides: True
--- a/doc/assets/vis_client.png
+++ b/doc/assets/vis_client.png
--- a/doc/assets/vis_server.png
+++ b/doc/assets/vis_server.png
--- a/doc/quickstart.md
+++ b/doc/quickstart.md
@ -2,7 +2,7 @@
 * @Date: 2021-04-02 11:53:16
 * @Author: Qing Shuai
 * @LastEditors: Qing Shuai
- * @LastEditTime: 2021-04-13 16:56:19
+ * @LastEditTime: 2021-05-27 20:15:52
 * @FilePath: /EasyMocapRelease/doc/quickstart.md
 -->
 # Quick Start
@ -15,11 +15,15 @@ We provide an example multiview dataset[[dropbox](https://www.dropbox.com/s/24mb
 data=path/to/data
 out=path/to/output
 # 0. extract the video to images
-python3 scripts/preprocess/extract_video.py ${data}
+python3 scripts/preprocess/extract_video.py ${data} --handface
 # 2.1 example for SMPL reconstruction
-python3 apps/demo/mv1p.py ${data} --out ${out} --vis_det --vis_repro --undis --sub_vis 1 7 13 19 --vis_smpl
+python3 apps/demo/mv1p.py ${data} --out ${out}/smpl --vis_det --vis_repro --undis --sub_vis 1 7 13 19 --vis_smpl
 # 2.2 example for SMPL-X reconstruction
-python3 apps/demo/mv1p.py ${data} --out ${out} --vis_det --vis_repro --undis --sub_vis 1 7 13 19 --body bodyhandface --model smplx --gender male --vis_smpl
+python3 apps/demo/mv1p.py ${data} --out ${out}/smplx --vis_det --vis_repro --undis --sub_vis 1 7 13 19 --body bodyhandface --model smplx --gender male --vis_smpl
+# 2.3 example for MANO reconstruction
+#     MANO model is required for this part
+python3 apps/demo/mv1p.py ${data} --out ${out}/manol --vis_det --vis_repro --undis --sub_vis 1 7 13 19 --body handl --model manol --gender male --vis_smpl
+python3 apps/demo/mv1p.py ${data} --out ${out}/manor --vis_det --vis_repro --undis --sub_vis 1 7 13 19 --body handr --model manor --gender male --vis_smpl
 ```

 # Demo On Your Dataset
@ -70,6 +74,7 @@ The output flags:
 - `--vis_repro`: visualize the reprojection
 - `--sub_vis`: use to specify the views to visualize. If not set, the code will use all views
 - `--vis_smpl`: use to render the SMPL mesh to images.
+- `--write_smpl_full`: use to write the full poses of the SMPL parameters

 ### 3. Output

--- a/doc/realtime_visualization.md
+++ b/doc/realtime_visualization.md
@ -0,0 +1,74 @@
+<!--
+ * @Date: 2021-06-04 15:56:55
+ * @Author: Qing Shuai
+ * @LastEditors: Qing Shuai
+ * @LastEditTime: 2021-06-04 17:11:48
+ * @FilePath: /EasyMocapRelease/doc/realtime_visualization.md
+-->
+# EasyMoCap -> Real-time Visualization
+
+We are the first one to release a real-time visualization tool for both skeletons and SMPL/SMPL+H/SMPL-X/MANO models.
+
+## Install
+
+Please install `EasyMocap` first. This part requires `Open3D==0.9.0`:
+
+```bash
+python3 -m pip install open3d==0.9.0
+```
+
+## Open the server
+Before any visualization, you should run a server:
+
+```bash
+python3 apps/vis/vis_server.py --cfg config/vis/o3d_scene.yml host <your_ip_address> port <set_a_port>
+```
+
+This step will open the visualization window:
+
+![](./assets/vis_server.png)
+
+You can alternate the viewpoints free. The configuration file `config/vis/o3d_scene.yml` defines the scene and other properties. In the default setting, we define the xyz-axis in the origin, the bounding box of the scene and a chessboard in the ground.
+
+## Send the data
+
+If you are success to open the server, you can visualize your 3D data anywhere. We provide an example code:
+
+```bash
+python3 apps/vis/vis_client.py --path <path/to/your/keypoints3d> --host <previous_ip_address> --port <previous_port>
+```
+
+Take the `zju-ls-feng` results as example, you can show the skeleton in the main window:
+
+![](./assets/vis_client.png)
+
+## Embed this feature to your code
+
+To add this visualization to your other code, you can follow these steps:
+
+```bash
+# 1. import the base client
+from easymocap.socket.base_client import BaseSocketClient
+# 2. set the ip address and port
+client = BaseSocketClient(host, port)
+# 3. send the data
+client.send(data)
+```
+
+The format of data is:
+```python
+data = [
+    {
+        'id': 0,
+        'keypoints3d': numpy.ndarray # (nJoints, 4) , (x, y, z, c) for each joint
+    },
+    {
+        'id': 1,
+        'keypoints3d': numpy.ndarray # (nJoints, 4)
+    }
+]
+```
+
+## Define your scene
+
+In the configuration file, we main define the `body_model` and `scene`. You can replace them for your data.
--- a/easymocap/config/init.py
+++ b/easymocap/config/init.py
@ -0,0 +1,9 @@
+'''
+  @ Date: 2021-06-04 13:58:01
+  @ Author: Qing Shuai
+  @ LastEditors: Qing Shuai
+  @ LastEditTime: 2021-06-04 13:58:43
+  @ FilePath: /EasyMocap/easymocap/config/__init__.py
+'''
+from .baseconfig import Config
+from .baseconfig import load_object
--- a/easymocap/config/baseconfig.py
+++ b/easymocap/config/baseconfig.py
@ -0,0 +1,54 @@
+'''
+  @ Date: 2021-05-28 14:18:20
+  @ Author: Qing Shuai
+  @ LastEditors: Qing Shuai
+  @ LastEditTime: 2021-06-04 15:43:31
+  @ FilePath: /EasyMocap/easymocap/config/baseconfig.py
+'''
+from .yacs import CfgNode as CN
+
+class Config:
+    @classmethod
+    def load_from_args(cls):
+        import argparse
+        parser = argparse.ArgumentParser()
+        parser.add_argument('--cfg', type=str, default='config/vis/base.yml')
+        parser.add_argument("opts", default=None, nargs=argparse.REMAINDER)
+        args = parser.parse_args()
+        return cls.load(filename=args.cfg, opts=args.opts)
+
+    @classmethod
+    def load(cls, filename=None, opts=[]) -> CN:
+        cfg = CN()
+        cfg = cls.init(cfg)
+        if filename is not None:
+            cfg.merge_from_file(filename)
+        if len(opts) > 0:
+            cfg.merge_from_list(opts)
+        cls.parse(cfg)
+        cls.print(cfg)
+        return cfg
+    
+    @staticmethod
+    def init(cfg):
+        return cfg
+    
+    @staticmethod
+    def parse(cfg):
+        pass
+
+    @staticmethod
+    def print(cfg):
+        print('[Info] --------------')
+        print('[Info] Configuration:')
+        print('[Info] --------------')
+        print(cfg)
+
+import importlib
+def load_object(module_name, module_args):
+    module_path = '.'.join(module_name.split('.')[:-1])
+    # scene_module = importlib.import_module(cfg.scene_module)
+    module = importlib.import_module(module_path)
+    name = module_name.split('.')[-1]
+    obj = getattr(module, name)(**module_args)
+    return obj
--- a/easymocap/config/vis_socket.py
+++ b/easymocap/config/vis_socket.py
@ -0,0 +1,65 @@
+'''
+  @ Date: 2021-05-30 11:17:18
+  @ Author: Qing Shuai
+  @ LastEditors: Qing Shuai
+  @ LastEditTime: 2021-06-04 15:44:56
+  @ FilePath: /EasyMocap/easymocap/config/vis_socket.py
+'''
+from .baseconfig import CN
+from .baseconfig import Config as BaseConfig
+import socket
+import numpy as np
+
+class Config(BaseConfig):
+    @staticmethod
+    def init(cfg):
+        # input and output
+        cfg.host = 'auto'
+        cfg.port = 9999
+        cfg.width = 1920
+        cfg.height = 1080
+        
+        cfg.body = 'body25'
+        cfg.max_human = 5
+        cfg.track = True
+        cfg.block = True # block visualization or not, True for visualize each frame, False in realtime applications
+        cfg.debug = False
+        cfg.write = False
+        cfg.out = '/'
+        # scene:
+        cfg.scene_module = "easymocap.visualize.o3dwrapper"
+        cfg.scene = CN()
+        cfg.extra = CN()
+        # skel
+        cfg.skel = CN()
+        cfg.skel.joint_radius = 0.02
+        # camera
+        cfg.camera = CN()
+        cfg.camera.phi = 0
+        cfg.camera.theta = -90 + 60
+        cfg.camera.cx = 0.
+        cfg.camera.cy = 0.
+        cfg.camera.cz = 6.
+        cfg.camera.set_camera = False
+        cfg.camera.camera_pose = []
+        return cfg
+    
+    @staticmethod
+    def parse(cfg):
+        if cfg.host == 'auto':
+            cfg.host = socket.gethostname()
+        if cfg.camera.set_camera:
+            pass
+        else:# use default camera
+            # theta, phi = cfg.camera.theta, cfg.camera.phi
+            theta, phi = np.deg2rad(cfg.camera.theta), np.deg2rad(cfg.camera.phi)
+            cx, cy, cz = cfg.camera.cx, cfg.camera.cy, cfg.camera.cz
+            st, ct = np.sin(theta), np.cos(theta)
+            sp, cp = np.sin(phi), np.cos(phi)
+            dist = 6
+            camera_pose = np.array([
+                    [cp, -st*sp, ct*sp, cx],
+                    [sp, st*cp, -ct*cp, cy],
+                    [0., ct, st, cz],
+                    [0.0, 0.0, 0.0, 1.0]])
+            cfg.camera.camera_pose = camera_pose.tolist()
--- a/easymocap/config/yacs.py
+++ b/easymocap/config/yacs.py
@ -0,0 +1,501 @@
+# Copyright (c) 2018-present, Facebook, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+##############################################################################
+
+"""YACS -- Yet Another Configuration System is designed to be a simple
+configuration management system for academic and industrial research
+projects.
+See README.md for usage and examples.
+"""
+
+import copy
+import io
+import logging
+import os
+from ast import literal_eval
+
+import yaml
+
+
+# Flag for py2 and py3 compatibility to use when separate code paths are necessary
+# When _PY2 is False, we assume Python 3 is in use
+_PY2 = False
+
+# Filename extensions for loading configs from files
+_YAML_EXTS = {"", ".yaml", ".yml"}
+_PY_EXTS = {".py"}
+
+# py2 and py3 compatibility for checking file object type
+# We simply use this to infer py2 vs py3
+try:
+    _FILE_TYPES = (file, io.IOBase)
+    _PY2 = True
+except NameError:
+    _FILE_TYPES = (io.IOBase,)
+
+# CfgNodes can only contain a limited set of valid types
+_VALID_TYPES = {tuple, list, str, int, float, bool}
+# py2 allow for str and unicode
+if _PY2:
+    _VALID_TYPES = _VALID_TYPES.union({unicode})  # noqa: F821
+
+# Utilities for importing modules from file paths
+if _PY2:
+    # imp is available in both py2 and py3 for now, but is deprecated in py3
+    import imp
+else:
+    import importlib.util
+
+logger = logging.getLogger(__name__)
+
+
+class CfgNode(dict):
+    """
+    CfgNode represents an internal node in the configuration tree. It's a simple
+    dict-like container that allows for attribute-based access to keys.
+    """
+
+    IMMUTABLE = "__immutable__"
+    DEPRECATED_KEYS = "__deprecated_keys__"
+    RENAMED_KEYS = "__renamed_keys__"
+
+    def __init__(self, init_dict=None, key_list=None):
+        # Recursively convert nested dictionaries in init_dict into CfgNodes
+        init_dict = {} if init_dict is None else init_dict
+        key_list = [] if key_list is None else key_list
+        for k, v in init_dict.items():
+            if type(v) is dict:
+                # Convert dict to CfgNode
+                init_dict[k] = CfgNode(v, key_list=key_list + [k])
+            else:
+                # Check for valid leaf type or nested CfgNode
+                _assert_with_logging(
+                    _valid_type(v, allow_cfg_node=True),
+                    "Key {} with value {} is not a valid type; valid types: {}".format(
+                        ".".join(key_list + [k]), type(v), _VALID_TYPES
+                    ),
+                )
+        super(CfgNode, self).__init__(init_dict)
+        # Manage if the CfgNode is frozen or not
+        self.__dict__[CfgNode.IMMUTABLE] = False
+        # Deprecated options
+        # If an option is removed from the code and you don't want to break existing
+        # yaml configs, you can add the full config key as a string to the set below.
+        self.__dict__[CfgNode.DEPRECATED_KEYS] = set()
+        # Renamed options
+        # If you rename a config option, record the mapping from the old name to the new
+        # name in the dictionary below. Optionally, if the type also changed, you can
+        # make the value a tuple that specifies first the renamed key and then
+        # instructions for how to edit the config file.
+        self.__dict__[CfgNode.RENAMED_KEYS] = {
+            # 'EXAMPLE.OLD.KEY': 'EXAMPLE.NEW.KEY',  # Dummy example to follow
+            # 'EXAMPLE.OLD.KEY': (                   # A more complex example to follow
+            #     'EXAMPLE.NEW.KEY',
+            #     "Also convert to a tuple, e.g., 'foo' -> ('foo',) or "
+            #     + "'foo:bar' -> ('foo', 'bar')"
+            # ),
+        }
+
+    def __getattr__(self, name):
+        if name in self:
+            return self[name]
+        else:
+            raise AttributeError(name)
+
+    def __setattr__(self, name, value):
+        if self.is_frozen():
+            raise AttributeError(
+                "Attempted to set {} to {}, but CfgNode is immutable".format(
+                    name, value
+                )
+            )
+
+        _assert_with_logging(
+            name not in self.__dict__,
+            "Invalid attempt to modify internal CfgNode state: {}".format(name),
+        )
+        _assert_with_logging(
+            _valid_type(value, allow_cfg_node=True),
+            "Invalid type {} for key {}; valid types = {}".format(
+                type(value), name, _VALID_TYPES
+            ),
+        )
+
+        self[name] = value
+
+    def __str__(self):
+        def _indent(s_, num_spaces):
+            s = s_.split("\n")
+            if len(s) == 1:
+                return s_
+            first = s.pop(0)
+            s = [(num_spaces * " ") + line for line in s]
+            s = "\n".join(s)
+            s = first + "\n" + s
+            return s
+
+        r = ""
+        s = []
+        for k, v in self.items():
+            seperator = "\n" if isinstance(v, CfgNode) else " "
+            attr_str = "{}:{}{}".format(str(k), seperator, str(v))
+            attr_str = _indent(attr_str, 4)
+            s.append(attr_str)
+        r += "\n".join(s)
+        return r
+
+    def __repr__(self):
+        return "{}({})".format(self.__class__.__name__, super(CfgNode, self).__repr__())
+
+    def dump(self):
+        """Dump to a string."""
+        self_as_dict = _to_dict(self)
+        return yaml.safe_dump(self_as_dict)
+
+    def merge_from_file(self, cfg_filename):
+        """Load a yaml config file and merge it this CfgNode."""
+        with open(cfg_filename, "r") as f:
+            cfg = load_cfg(f)
+        if 'parent' in cfg.keys():
+            if cfg.parent != 'none':
+                print('[Config] merge from parent file: {}'.format(cfg.parent))
+                self.merge_from_file(cfg.parent)
+        self.merge_from_other_cfg(cfg)
+
+    def merge_from_other_cfg(self, cfg_other):
+        """Merge `cfg_other` into this CfgNode."""
+        _merge_a_into_b(cfg_other, self, self, [])
+
+    def merge_from_list(self, cfg_list):
+        """Merge config (keys, values) in a list (e.g., from command line) into
+        this CfgNode. For example, `cfg_list = ['FOO.BAR', 0.5]`.
+        """
+        _assert_with_logging(
+            len(cfg_list) % 2 == 0,
+            "Override list has odd length: {}; it must be a list of pairs".format(
+                cfg_list
+            ),
+        )
+        root = self
+        for full_key, v in zip(cfg_list[0::2], cfg_list[1::2]):
+            if root.key_is_deprecated(full_key):
+                continue
+            if root.key_is_renamed(full_key):
+                root.raise_key_rename_error(full_key)
+            key_list = full_key.split(".")
+            d = self
+            for subkey in key_list[:-1]:
+                _assert_with_logging(
+                    subkey in d, "Non-existent key: {}".format(full_key)
+                )
+                d = d[subkey]
+            subkey = key_list[-1]
+            _assert_with_logging(subkey in d, "Non-existent key: {}".format(full_key))
+            value = _decode_cfg_value(v)
+            value = _check_and_coerce_cfg_value_type(value, d[subkey], subkey, full_key)
+            d[subkey] = value
+
+    def freeze(self):
+        """Make this CfgNode and all of its children immutable."""
+        self._immutable(True)
+
+    def defrost(self):
+        """Make this CfgNode and all of its children mutable."""
+        self._immutable(False)
+
+    def is_frozen(self):
+        """Return mutability."""
+        return self.__dict__[CfgNode.IMMUTABLE]
+
+    def _immutable(self, is_immutable):
+        """Set immutability to is_immutable and recursively apply the setting
+        to all nested CfgNodes.
+        """
+        self.__dict__[CfgNode.IMMUTABLE] = is_immutable
+        # Recursively set immutable state
+        for v in self.__dict__.values():
+            if isinstance(v, CfgNode):
+                v._immutable(is_immutable)
+        for v in self.values():
+            if isinstance(v, CfgNode):
+                v._immutable(is_immutable)
+
+    def clone(self):
+        """Recursively copy this CfgNode."""
+        return copy.deepcopy(self)
+
+    def register_deprecated_key(self, key):
+        """Register key (e.g. `FOO.BAR`) a deprecated option. When merging deprecated
+        keys a warning is generated and the key is ignored.
+        """
+        _assert_with_logging(
+            key not in self.__dict__[CfgNode.DEPRECATED_KEYS],
+            "key {} is already registered as a deprecated key".format(key),
+        )
+        self.__dict__[CfgNode.DEPRECATED_KEYS].add(key)
+
+    def register_renamed_key(self, old_name, new_name, message=None):
+        """Register a key as having been renamed from `old_name` to `new_name`.
+        When merging a renamed key, an exception is thrown alerting to user to
+        the fact that the key has been renamed.
+        """
+        _assert_with_logging(
+            old_name not in self.__dict__[CfgNode.RENAMED_KEYS],
+            "key {} is already registered as a renamed cfg key".format(old_name),
+        )
+        value = new_name
+        if message:
+            value = (new_name, message)
+        self.__dict__[CfgNode.RENAMED_KEYS][old_name] = value
+
+    def key_is_deprecated(self, full_key):
+        """Test if a key is deprecated."""
+        if full_key in self.__dict__[CfgNode.DEPRECATED_KEYS]:
+            logger.warning("Deprecated config key (ignoring): {}".format(full_key))
+            return True
+        return False
+
+    def key_is_renamed(self, full_key):
+        """Test if a key is renamed."""
+        return full_key in self.__dict__[CfgNode.RENAMED_KEYS]
+
+    def raise_key_rename_error(self, full_key):
+        new_key = self.__dict__[CfgNode.RENAMED_KEYS][full_key]
+        if isinstance(new_key, tuple):
+            msg = " Note: " + new_key[1]
+            new_key = new_key[0]
+        else:
+            msg = ""
+        raise KeyError(
+            "Key {} was renamed to {}; please update your config.{}".format(
+                full_key, new_key, msg
+            )
+        )
+
+
+def load_cfg(cfg_file_obj_or_str):
+    """Load a cfg. Supports loading from:
+        - A file object backed by a YAML file
+        - A file object backed by a Python source file that exports an attribute
+          "cfg" that is either a dict or a CfgNode
+        - A string that can be parsed as valid YAML
+    """
+    _assert_with_logging(
+        isinstance(cfg_file_obj_or_str, _FILE_TYPES + (str,)),
+        "Expected first argument to be of type {} or {}, but it was {}".format(
+            _FILE_TYPES, str, type(cfg_file_obj_or_str)
+        ),
+    )
+    if isinstance(cfg_file_obj_or_str, str):
+        return _load_cfg_from_yaml_str(cfg_file_obj_or_str)
+    elif isinstance(cfg_file_obj_or_str, _FILE_TYPES):
+        return _load_cfg_from_file(cfg_file_obj_or_str)
+    else:
+        raise NotImplementedError("Impossible to reach here (unless there's a bug)")
+
+
+def _load_cfg_from_file(file_obj):
+    """Load a config from a YAML file or a Python source file."""
+    _, file_extension = os.path.splitext(file_obj.name)
+    if file_extension in _YAML_EXTS:
+        return _load_cfg_from_yaml_str(file_obj.read())
+    elif file_extension in _PY_EXTS:
+        return _load_cfg_py_source(file_obj.name)
+    else:
+        raise Exception(
+            "Attempt to load from an unsupported file type {}; "
+            "only {} are supported".format(file_obj, _YAML_EXTS.union(_PY_EXTS))
+        )
+
+
+def _load_cfg_from_yaml_str(str_obj):
+    """Load a config from a YAML string encoding."""
+    cfg_as_dict = yaml.safe_load(str_obj)
+    return CfgNode(cfg_as_dict)
+
+
+def _load_cfg_py_source(filename):
+    """Load a config from a Python source file."""
+    module = _load_module_from_file("yacs.config.override", filename)
+    _assert_with_logging(
+        hasattr(module, "cfg"),
+        "Python module from file {} must have 'cfg' attr".format(filename),
+    )
+    VALID_ATTR_TYPES = {dict, CfgNode}
+    _assert_with_logging(
+        type(module.cfg) in VALID_ATTR_TYPES,
+        "Imported module 'cfg' attr must be in {} but is {} instead".format(
+            VALID_ATTR_TYPES, type(module.cfg)
+        ),
+    )
+    if type(module.cfg) is dict:
+        return CfgNode(module.cfg)
+    else:
+        return module.cfg
+
+
+def _to_dict(cfg_node):
+    """Recursively convert all CfgNode objects to dict objects."""
+
+    def convert_to_dict(cfg_node, key_list):
+        if not isinstance(cfg_node, CfgNode):
+            _assert_with_logging(
+                _valid_type(cfg_node),
+                "Key {} with value {} is not a valid type; valid types: {}".format(
+                    ".".join(key_list), type(cfg_node), _VALID_TYPES
+                ),
+            )
+            return cfg_node
+        else:
+            cfg_dict = dict(cfg_node)
+            for k, v in cfg_dict.items():
+                cfg_dict[k] = convert_to_dict(v, key_list + [k])
+            return cfg_dict
+
+    return convert_to_dict(cfg_node, [])
+
+
+def _valid_type(value, allow_cfg_node=False):
+    return (type(value) in _VALID_TYPES) or (allow_cfg_node and type(value) == CfgNode)
+
+
+def _merge_a_into_b(a, b, root, key_list):
+    """Merge config dictionary a into config dictionary b, clobbering the
+    options in b whenever they are also specified in a.
+    """
+    _assert_with_logging(
+        isinstance(a, CfgNode),
+        "`a` (cur type {}) must be an instance of {}".format(type(a), CfgNode),
+    )
+    _assert_with_logging(
+        isinstance(b, CfgNode),
+        "`b` (cur type {}) must be an instance of {}".format(type(b), CfgNode),
+    )
+
+    for k, v_ in a.items():
+        full_key = ".".join(key_list + [k])
+        # a must specify keys that are in b
+        if k not in b:
+            if root.key_is_deprecated(full_key):
+                continue
+            elif root.key_is_renamed(full_key):
+                root.raise_key_rename_error(full_key)
+            else:
+                v = copy.deepcopy(v_)
+                v = _decode_cfg_value(v)
+                b.update({k: v})
+        else:
+            v = copy.deepcopy(v_)
+            v = _decode_cfg_value(v)
+            v = _check_and_coerce_cfg_value_type(v, b[k], k, full_key)
+
+        # Recursively merge dicts
+        if isinstance(v, CfgNode):
+            try:
+                _merge_a_into_b(v, b[k], root, key_list + [k])
+            except BaseException:
+                raise
+        else:
+            b[k] = v
+
+
+def _decode_cfg_value(v):
+    """Decodes a raw config value (e.g., from a yaml config files or command
+    line argument) into a Python object.
+    """
+    # Configs parsed from raw yaml will contain dictionary keys that need to be
+    # converted to CfgNode objects
+    if isinstance(v, dict):
+        return CfgNode(v)
+    # All remaining processing is only applied to strings
+    if not isinstance(v, str):
+        return v
+    # Try to interpret `v` as a:
+    #   string, number, tuple, list, dict, boolean, or None
+    try:
+        v = literal_eval(v)
+    # The following two excepts allow v to pass through when it represents a
+    # string.
+    #
+    # Longer explanation:
+    # The type of v is always a string (before calling literal_eval), but
+    # sometimes it *represents* a string and other times a data structure, like
+    # a list. In the case that v represents a string, what we got back from the
+    # yaml parser is 'foo' *without quotes* (so, not '"foo"'). literal_eval is
+    # ok with '"foo"', but will raise a ValueError if given 'foo'. In other
+    # cases, like paths (v = 'foo/bar' and not v = '"foo/bar"'), literal_eval
+    # will raise a SyntaxError.
+    except ValueError:
+        pass
+    except SyntaxError:
+        pass
+    return v
+
+
+def _check_and_coerce_cfg_value_type(replacement, original, key, full_key):
+    """Checks that `replacement`, which is intended to replace `original` is of
+    the right type. The type is correct if it matches exactly or is one of a few
+    cases in which the type can be easily coerced.
+    """
+    original_type = type(original)
+    replacement_type = type(replacement)
+
+    # The types must match (with some exceptions)
+    if replacement_type == original_type:
+        return replacement
+
+    # Cast replacement from from_type to to_type if the replacement and original
+    # types match from_type and to_type
+    def conditional_cast(from_type, to_type):
+        if replacement_type == from_type and original_type == to_type:
+            return True, to_type(replacement)
+        else:
+            return False, None
+
+    # Conditionally casts
+    # list <-> tuple
+    casts = [(tuple, list), (list, tuple), (int, float), (float, int)]
+    # For py2: allow converting from str (bytes) to a unicode string
+    try:
+        casts.append((str, unicode))  # noqa: F821
+    except Exception:
+        pass
+
+    for (from_type, to_type) in casts:
+        converted, converted_value = conditional_cast(from_type, to_type)
+        if converted:
+            return converted_value
+
+    raise ValueError(
+        "Type mismatch ({} vs. {}) with values ({} vs. {}) for config "
+        "key: {}".format(
+            original_type, replacement_type, original, replacement, full_key
+        )
+    )
+
+
+def _assert_with_logging(cond, msg):
+    if not cond:
+        logger.debug(msg)
+    assert cond, msg
+
+
+def _load_module_from_file(name, filename):
+    if _PY2:
+        module = imp.load_source(name, filename)
+    else:
+        spec = importlib.util.spec_from_file_location(name, filename)
+        module = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(module)
+    return module
--- a/easymocap/mytools/utils.py
+++ b/easymocap/mytools/utils.py
@ -2,8 +2,8 @@
  @ Date: 2021-01-15 11:12:00
  @ Author: Qing Shuai
  @ LastEditors: Qing Shuai
-  @ LastEditTime: 2021-03-08 21:07:48
-  @ FilePath: /EasyMocap/code/mytools/utils.py
+  @ LastEditTime: 2021-05-27 14:55:40
+  @ FilePath: /EasyMocap/easymocap/mytools/utils.py
 '''
 import time
 import tabulate
@ -30,4 +30,10 @@ class Timer:
        end = time.time()
        Timer.records[self.name].append((end-self.start)*1000)
        if not self.silent:
-            print('-> [{:20s}]: {:5.1f}ms'.format(self.name, (end-self.start)*1000))
+            t = (end - self.start)*1000
+            if t > 10000:
+                print('-> [{:20s}]: {:5.1f}s'.format(self.name, t/1000))
+            elif t > 1e3*60*60:
+                print('-> [{:20s}]: {:5.1f}min'.format(self.name, t/1e3/60))
+            else:
+                print('-> [{:20s}]: {:5.1f}ms'.format(self.name, (end-self.start)*1000))
--- a/easymocap/mytools/vis_base.py
+++ b/easymocap/mytools/vis_base.py
@ -2,7 +2,7 @@
  @ Date: 2020-11-28 17:23:04
  @ Author: Qing Shuai
  @ LastEditors: Qing Shuai
-  @ LastEditTime: 2021-03-28 22:19:34
+  @ LastEditTime: 2021-06-03 22:31:31
  @ FilePath: /EasyMocap/easymocap/mytools/vis_base.py
 '''
 import cv2
@ -57,17 +57,26 @@ def get_rgb(index):
        col = tuple([int(c*255) for c in col[::-1]])
    return col

-def plot_point(img, x, y, r, col, pid=-1):
-    cv2.circle(img, (int(x+0.5), int(y+0.5)), r, col, -1)
+def get_rgb_01(index):
+    col = get_rgb(index)
+    return [i*1./255 for i in col[:3]]
+
+def plot_point(img, x, y, r, col, pid=-1, font_scale=-1, circle_type=-1):
+    cv2.circle(img, (int(x+0.5), int(y+0.5)), r, col, circle_type)
+    if font_scale == -1:
+        font_scale = img.shape[0]/4000
    if pid != -1:
-        cv2.putText(img, '{}'.format(pid), (int(x+0.5), int(y+0.5)), cv2.FONT_HERSHEY_SIMPLEX, 1, col, 2)
+        cv2.putText(img, '{}'.format(pid), (int(x+0.5), int(y+0.5)), cv2.FONT_HERSHEY_SIMPLEX, font_scale, col, 1)


 def plot_line(img, pt1, pt2, lw, col):
    cv2.line(img, (int(pt1[0]+0.5), int(pt1[1]+0.5)), (int(pt2[0]+0.5), int(pt2[1]+0.5)),
        col, lw)

-def plot_cross(img, x, y, col, width=10, lw=2):
+def plot_cross(img, x, y, col, width=-1, lw=-1):
+    if lw == -1:
+        lw = int(round(img.shape[0]/1000))
+        width = lw * 5
    cv2.line(img, (int(x-width), int(y)), (int(x+width), int(y)), col, lw)
    cv2.line(img, (int(x), int(y-width)), (int(x), int(y+width)), col, lw)
    
@ -82,7 +91,8 @@ def plot_bbox(img, bbox, pid, vis_id=True):
    lw = max(img.shape[0]//300, 2)
    cv2.rectangle(img, (x1, y1), (x2, y2), color, lw)
    if vis_id:
-        cv2.putText(img, '{}'.format(pid), (x1, y1+20), cv2.FONT_HERSHEY_SIMPLEX, 1, color, 2)
+        font_scale = img.shape[0]/1000
+        cv2.putText(img, '{}'.format(pid), (x1, y1+int(25*font_scale)), cv2.FONT_HERSHEY_SIMPLEX, font_scale, color, 2)

 def plot_keypoints(img, points, pid, config, vis_conf=False, use_limb_color=True, lw=2):
    for ii, (i, j) in enumerate(config['kintree']):
@ -108,33 +118,44 @@ def plot_keypoints(img, points, pid, config, vis_conf=False, use_limb_color=True

 def plot_points2d(img, points2d, lines, lw=4, col=(0, 255, 0), putText=True):
    # 将2d点画上去
+    if points2d.shape[1] == 2:
+        points2d = np.hstack([points2d, np.ones((points2d.shape[0], 1))])
    for i, (x, y, v) in enumerate(points2d):
        if v < 0.01:
            continue
        c = col
        plot_cross(img, x, y, width=10, col=c, lw=lw)
        if putText:
-            cv2.putText(img, '{}'.format(i), (int(x), int(y)), cv2.FONT_HERSHEY_SIMPLEX, 1, c, 2)
+            font_scale = img.shape[0]/2000
+            cv2.putText(img, '{}'.format(i), (int(x), int(y)), cv2.FONT_HERSHEY_SIMPLEX, font_scale, c, 2)
    for i, j in lines:
        if points2d[i][2] < 0.01 or points2d[j][2] < 0.01:
            continue
-        plot_line(img, points2d[i], points2d[j], 2, (255, 255, 255))
+        plot_line(img, points2d[i], points2d[j], 2, col)

-def merge(images, row=-1, col=-1, resize=False, ret_range=False):
-    if row == -1 and col == -1:
+row_col_ = {
+    2: (2, 1),
+    7: (2, 4),
+    8: (2, 4),
+    9: (3, 3),
+    26: (4, 7)
+}
+def get_row_col(l):
+    if l in row_col_.keys():
+        return row_col_[l]
+    else:
        from math import sqrt
-        row = int(sqrt(len(images)) + 0.5)
-        col = int(len(images)/ row + 0.5)
+        row = int(sqrt(l) + 0.5)
+        col = int(l/ row + 0.5)
+        if row*col<l:
+            col = col + 1
        if row > col:
            row, col = col, row
-        if len(images) == 8:
-            # basketball 场景
-            row, col = 2, 4
-            images = [images[i] for i in [0, 1, 2, 3, 7, 6, 5, 4]]
-        if len(images) == 7:
-            row, col = 3, 3
-        elif len(images) == 2:
-            row, col = 2, 1
+        return row, col
+
+def merge(images, row=-1, col=-1, resize=False, ret_range=False, **kwargs):
+    if row == -1 and col == -1:
+        row, col = get_row_col(len(images))
    height = images[0].shape[0]
    width = images[0].shape[1]
    ret_img = np.zeros((height * row, width * col, images[0].shape[2]), dtype=np.uint8) + 255
@ -149,8 +170,9 @@ def merge(images, row=-1, col=-1, resize=False, ret_range=False):
            ret_img[height * i: height * (i+1), width * j: width * (j+1)] = img
            ranges.append((width*j, height*i, width*(j+1), height*(i+1)))
    if resize:
-        scale = min(1000/ret_img.shape[0], 1800/ret_img.shape[1])
-        while ret_img.shape[0] > 2000:
+        min_height = 3000
+        if ret_img.shape[0] > min_height:
+            scale = min_height/ret_img.shape[0]
            ret_img = cv2.resize(ret_img, None, fx=scale, fy=scale)
    if ret_range:
        return ret_img, ranges
--- a/easymocap/socket/base.py
+++ b/easymocap/socket/base.py
@ -0,0 +1,80 @@
+'''
+  @ Date: 2021-05-25 11:14:48
+  @ Author: Qing Shuai
+  @ LastEditors: Qing Shuai
+  @ LastEditTime: 2021-06-02 13:00:35
+  @ FilePath: /EasyMocap/easymocap/socket/base.py
+'''
+import socket
+import time
+from threading import Thread
+from queue import Queue
+
+def log(x):
+    from datetime import datetime
+    time_now = datetime.now().strftime("%m-%d-%H:%M:%S.%f ")
+    print(time_now + x)
+
+class BaseSocket:
+    def __init__(self, host, port, debug=False) -> None:
+        # 创建 socket 对象
+        print('[Info] server start')
+        serversocket = socket.socket(
+                    socket.AF_INET, socket.SOCK_STREAM)
+        serversocket.bind((host, port))
+        serversocket.listen(1)
+        self.serversocket = serversocket
+        self.t = Thread(target=self.run)
+        self.t.start()
+        self.queue = Queue()
+        self.debug = debug
+        self.disconnect = False
+    
+    @staticmethod
+    def recvLine(sock):
+        flag = True
+        result = b''
+        while not result.endswith(b'\n'):
+            res = sock.recv(1)
+            if not res:
+                flag = False
+                break
+            result += res
+        return flag, result.strip().decode('ascii')
+
+    @staticmethod
+    def recvAll(sock, l):
+        l = int(l)
+        result = b''
+        while (len(result) < l):
+            t = sock.recv(l - len(result))
+            result += t
+        return result.decode('ascii')
+
+    def run(self):
+        while True:
+            clientsocket, addr = self.serversocket.accept()
+            print("[Info] Connect: %s" % str(addr))
+            while True:
+                flag, l = self.recvLine(clientsocket)
+                if not flag:
+                    print("[Info] Disonnect: %s" % str(addr))
+                    break
+                data = self.recvAll(clientsocket, l)
+                if self.debug:log('[Info] Recv data')
+                self.queue.put(data)
+            clientsocket.close()
+    
+    def update(self):
+        time.sleep(1)
+        while not self.queue.empty():
+            log('update')
+            data = self.queue.get()
+            self.main(data)
+    
+    def main(self, datas):
+        print(datas)
+
+    def __del__(self):
+        self.serversocket.close()
+        self.t.join()
--- a/easymocap/socket/base_client.py
+++ b/easymocap/socket/base_client.py
@ -0,0 +1,25 @@
+'''
+  @ Date: 2021-05-25 13:39:07
+  @ Author: Qing Shuai
+  @ LastEditors: Qing Shuai
+  @ LastEditTime: 2021-06-04 16:43:39
+  @ FilePath: /EasyMocapRelease/easymocap/socket/base_client.py
+'''
+import socket
+from .utils import encode_detect
+
+class BaseSocketClient:
+    def __init__(self, host, port) -> None:
+        if host == 'auto':
+            host = socket.gethostname()
+        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        s.connect((host, port))
+        self.s = s
+    
+    def send(self, data):
+        val = encode_detect(data)
+        self.s.send(bytes('{}\n'.format(len(val)), 'ascii'))
+        self.s.sendall(val)
+    
+    def close(self):
+        self.s.close()
--- a/easymocap/socket/o3d.py
+++ b/easymocap/socket/o3d.py
@ -0,0 +1,114 @@
+'''
+  @ Date: 2021-05-25 11:15:53
+  @ Author: Qing Shuai
+  @ LastEditors: Qing Shuai
+  @ LastEditTime: 2021-06-04 17:06:17
+  @ FilePath: /EasyMocapRelease/easymocap/socket/o3d.py
+'''
+import open3d as o3d
+from ..config import load_object
+from ..visualize.o3dwrapper import Vector3dVector, create_mesh, load_mesh
+from ..mytools import Timer
+from ..mytools.vis_base import get_rgb_01
+from .base import BaseSocket, log
+import json
+import numpy as np
+from os.path import join
+import os
+
+class VisOpen3DSocket(BaseSocket):
+    def __init__(self, host, port, cfg) -> None:
+        # output
+        self.write = cfg.write
+        self.out = cfg.out
+        if self.write:
+            print('[Info] capture the screen to {}'.format(self.out))
+            os.makedirs(self.out, exist_ok=True)
+        # scene
+        vis = o3d.visualization.Visualizer()
+        vis.create_window(window_name='Visualizer', width=cfg.width, height=cfg.height)
+        self.vis = vis
+        # load the scene
+        for key, mesh_args in cfg.scene.items():
+            mesh = load_object(key, mesh_args)
+            self.vis.add_geometry(mesh)
+        for key, val in cfg.extra.items():
+            mesh = load_mesh(val["path"])
+            trans = np.array(val['transform']).reshape(4, 4)
+            mesh.transform(trans)
+            self.vis.add_geometry(mesh)
+        # create vis => update => super() init
+        super().__init__(host, port, debug=cfg.debug)
+        self.block = cfg.block
+        self.body_model = load_object(cfg.body_model.module, cfg.body_model.args)
+        zero_params = self.body_model.init_params(1)
+        self.max_human = cfg.max_human
+        self.track = cfg.track
+        self.camera_pose = cfg.camera.camera_pose
+        self.init_camera(cfg.camera.camera_pose)
+        self.zero_vertices = Vector3dVector(np.zeros((self.body_model.nVertices, 3)))
+
+        self.vertices, self.meshes = [], []
+        for i in range(self.max_human):
+            self.add_human(zero_params)
+
+        self.count = 0
+    
+    def add_human(self, zero_params):
+        vertices = self.body_model(zero_params)[0]
+        self.vertices.append(vertices)
+        mesh = create_mesh(vertices=vertices, faces=self.body_model.faces)
+        self.meshes.append(mesh)
+        self.vis.add_geometry(mesh)
+        self.init_camera(self.camera_pose)
+
+    def init_camera(self, camera_pose):
+        ctr = self.vis.get_view_control()
+        init_param = ctr.convert_to_pinhole_camera_parameters()
+        # init_param.intrinsic.set_intrinsics(init_param.intrinsic.width, init_param.intrinsic.height, fx, fy, cx, cy)
+        init_param.extrinsic = np.array(camera_pose)
+        ctr.convert_from_pinhole_camera_parameters(init_param) 
+
+    def main(self, datas):
+        if self.debug:log('[Info] Load data {}'.format(self.count))
+        datas = json.loads(datas)
+        with Timer('forward'):
+            for i, data in enumerate(datas):
+                if i >= len(self.meshes):
+                    print('[Error] the number of human exceeds!')
+                    self.add_human(np.array(data['keypoints3d']))
+                vertices = self.body_model(np.array(data['keypoints3d']))
+                self.vertices[i] = Vector3dVector(vertices[0])
+            for i in range(len(datas), len(self.meshes)):
+                self.vertices[i] = self.zero_vertices
+        # Open3D will lock the thread here
+        with Timer('set vertices'):
+            for i in range(len(self.vertices)):
+                self.meshes[i].vertices = self.vertices[i]
+                if i < len(datas) and self.track:
+                    col = get_rgb_01(datas[i]['id'])
+                    self.meshes[i].paint_uniform_color(col)
+
+    def update(self):
+        if not self.queue.empty():
+            if self.debug:log('Update' + str(self.queue.qsize()))
+            datas = self.queue.get()
+            if not self.block:
+                while self.queue.qsize() > 0:
+                    datas = self.queue.get()
+            self.main(datas)
+            with Timer('update geometry'):
+                for mesh in self.meshes:
+                    mesh.compute_triangle_normals()
+                    self.vis.update_geometry(mesh)
+                self.vis.poll_events()
+                self.vis.update_renderer()
+            if self.write:
+                outname = join(self.out, '{:06d}.jpg'.format(self.count))
+                with Timer('capture'):
+                    self.vis.capture_screen_image(outname)
+            self.count += 1
+        else:
+            with Timer('update renderer', True):
+                self.vis.poll_events()
+                self.vis.update_renderer()
--- a/easymocap/socket/utils.py
+++ b/easymocap/socket/utils.py
@ -0,0 +1,23 @@
+'''
+  @ Date: 2021-05-24 20:07:34
+  @ Author: Qing Shuai
+  @ LastEditors: Qing Shuai
+  @ LastEditTime: 2021-06-04 16:29:35
+  @ FilePath: /EasyMocapRelease/media/qing/Project/mirror/EasyMocap/easymocap/socket/utils.py
+'''
+import cv2
+import numpy as np
+from ..mytools.file_utils import write_common_results
+
+def encode_detect(data):
+    res = write_common_results(None, data, ['keypoints3d'])
+    res = res.replace('\r', '').replace('\n', '').replace(' ', '')
+    return res.encode('ascii')
+
+def encode_image(image):
+    fourcc = [int(cv2.IMWRITE_JPEG_QUALITY), 90]
+    #frame을 binary 형태로 변환 jpg로 decoding
+    result, img_encode = cv2.imencode('.jpg', image, fourcc)
+    data = np.array(img_encode) # numpy array로 안바꿔주면 ERROR
+    stringData = data.tostring()
+    return stringData
--- a/easymocap/visualize/geometry.py
+++ b/easymocap/visualize/geometry.py
@ -2,12 +2,13 @@
  @ Date: 2021-01-17 22:44:34
  @ Author: Qing Shuai
  @ LastEditors: Qing Shuai
-  @ LastEditTime: 2021-04-03 20:00:02
-  @ FilePath: /EasyMocap/code/visualize/geometry.py
+  @ LastEditTime: 2021-05-25 14:01:24
+  @ FilePath: /EasyMocap/easymocap/visualize/geometry.py
 '''
 import numpy as np
 import cv2
 import numpy as np
+from tqdm import tqdm

 def create_ground(
    center=[0, 0, 0], xdir=[1, 0, 0], ydir=[0, 1, 0], # 位置
@ -19,15 +20,15 @@ def create_ground(
        center = np.array(center)
        xdir = np.array(xdir)
        ydir = np.array(ydir)
-    print(center, xdir, ydir)
+    print('[Vis Info] {}, x: {}, y: {}'.format(center, xdir, ydir))
    xdir = xdir * step
    ydir = ydir * step
    vertls, trils, colls = [],[],[]
    cnt = 0
    min_x = -xrange if two_sides else 0
    min_y = -yrange if two_sides else 0
-    for i in range(min_x, xrange+1):
-        for j in range(min_y, yrange+1):
+    for i in range(min_x, xrange):
+        for j in range(min_y, yrange):
            point0 = center + i*xdir + j*ydir
            point1 = center + (i+1)*xdir + j*ydir
            point2 = center + (i+1)*xdir + (j+1)*ydir
@ -107,3 +108,46 @@ def create_cameras(cameras):
            'vertices': vertices, 'faces': triangles, 'name': 'camera_{}'.format(nv), 'vid': nv
        })
    return meshes
+
+import os
+current_dir = os.path.dirname(os.path.realpath(__file__))
+
+def create_cameras_texture(cameras, imgnames, scale=5e-3):
+    import trimesh
+    import pyrender
+    from PIL import Image
+    from os.path import join
+    cam_path = join(current_dir, 'objs', 'background.obj')
+    meshes = []
+    for nv, (key, camera) in enumerate(tqdm(cameras.items(), desc='loading images')):
+        cam_trimesh = trimesh.load(cam_path, process=False)
+        vert = np.asarray(cam_trimesh.vertices)
+        K, R, T = camera['K'], camera['R'], camera['T']
+        img = Image.open(imgnames[nv])
+        height, width = img.height, img.width
+        vert[:, 0] *= width
+        vert[:, 1] *= height
+        vert[:, 2] *= 0
+        vert[:, 0] -= vert[:, 0]*0.5
+        vert[:, 1] -= vert[:, 1]*0.5
+        vert[:, 1] = - vert[:, 1]
+        vert[:, :2] *= scale
+        # vert[:, 2] = 1
+        cam_trimesh.vertices = (vert - T.T) @ R
+        cam_trimesh.visual.material.image = img
+        cam_mesh = pyrender.Mesh.from_trimesh(cam_trimesh, smooth=True)
+        meshes.append(cam_mesh)
+    return meshes
+
+def create_mesh_pyrender(vert, faces, col):
+    import trimesh
+    import pyrender
+    mesh = trimesh.Trimesh(vert, faces, process=False)
+    material = pyrender.MetallicRoughnessMaterial(
+        metallicFactor=0.0,
+        alphaMode='OPAQUE',
+        baseColorFactor=col)
+    mesh = pyrender.Mesh.from_trimesh(
+        mesh,
+        material=material)
+    return mesh
--- a/easymocap/visualize/o3dwrapper.py
+++ b/easymocap/visualize/o3dwrapper.py
@ -0,0 +1,45 @@
+'''
+  @ Date: 2021-04-25 15:52:01
+  @ Author: Qing Shuai
+  @ LastEditors: Qing Shuai
+  @ LastEditTime: 2021-05-25 11:48:49
+  @ FilePath: /EasyMocap/easymocap/visualize/o3dwrapper.py
+'''
+import open3d as o3d
+import numpy as np
+
+Vector3dVector = o3d.utility.Vector3dVector
+Vector3iVector = o3d.utility.Vector3iVector
+Vector2iVector = o3d.utility.Vector2iVector
+TriangleMesh = o3d.geometry.TriangleMesh
+load_mesh = o3d.io.read_triangle_mesh
+
+def create_mesh(vertices, faces, colors=None, **kwargs):
+    mesh = TriangleMesh()
+    mesh.vertices = Vector3dVector(vertices)
+    mesh.triangles = Vector3iVector(faces)
+    if colors is not None:
+        mesh.vertex_colors = Vector3dVector(colors)
+    else:
+        mesh.paint_uniform_color([1., 0.8, 0.8])
+    mesh.compute_vertex_normals()
+    return mesh
+
+def create_ground(**kwargs):
+    from .geometry import create_ground as create_ground_
+    ground = create_ground_(**kwargs)
+    return create_mesh(**ground)
+
+def create_coord(camera = [0,0,0], radius=1):
+    camera_frame = TriangleMesh.create_coordinate_frame(
+            size=radius, origin=camera)
+    return camera_frame
+
+def create_bbox(min_bound=(-3., -3., 0), max_bound=(3., 3., 2), flip=False):
+    if flip:
+        min_bound_ = min_bound.copy()
+        max_bound_ = max_bound.copy()
+        min_bound = [min_bound_[0], -max_bound_[1], -max_bound_[2]]
+        max_bound = [max_bound_[0], -min_bound_[1], -min_bound_[2]]
+    bbox = o3d.geometry.AxisAlignedBoundingBox(min_bound, max_bound)
+    return bbox
--- a/easymocap/visualize/skelmodel.py
+++ b/easymocap/visualize/skelmodel.py
@ -2,13 +2,14 @@
  @ Date: 2021-01-17 21:38:19
  @ Author: Qing Shuai
  @ LastEditors: Qing Shuai
-  @ LastEditTime: 2021-01-22 23:08:18
-  @ FilePath: /EasyMocap/code/visualize/skelmodel.py
+  @ LastEditTime: 2021-06-04 15:52:49
+  @ FilePath: /EasyMocap/easymocap/visualize/skelmodel.py
 '''
 import numpy as np
 import cv2
 from os.path import join
 import os
+from ..dataset.config import CONFIG

 def calTransformation(v_i, v_j, r, adaptr=False, ratio=10):
    """ from to vertices to T
@ -27,7 +28,7 @@ def calTransformation(v_i, v_j, r, adaptr=False, ratio=10):
    rotdir = rotdir * np.arccos(np.dot(direc, xaxis))
    rotmat, _ = cv2.Rodrigues(rotdir)
    # set the minimal radius for the finger and face
-    shrink = max(length/ratio, 0.005)
+    shrink = min(max(length/ratio, 0.005), 0.05)
    eigval = np.array([[length/2/r, 0, 0], [0, shrink, 0], [0, 0, shrink]])
    T = np.eye(4)
    T[:3,:3] = rotmat @ eigval @ rotmat.T
@ -35,33 +36,36 @@ def calTransformation(v_i, v_j, r, adaptr=False, ratio=10):
    return T, r, length

 class SkelModel:
-    def __init__(self, nJoints, kintree) -> None:
-        self.nJoints = nJoints
-        self.kintree = kintree
+    def __init__(self, nJoints=None, kintree=None, body_type=None, joint_radius=0.02, **kwargs) -> None:
+        if nJoints is not None:
+            self.nJoints = nJoints
+            self.kintree = kintree
+        else:
+            config = CONFIG[body_type]
+            self.nJoints = config['nJoints']
+            self.kintree = config['kintree']
        cur_dir = os.path.dirname(__file__)
        faces = np.loadtxt(join(cur_dir, 'sphere_faces_20.txt'), dtype=np.int)
        self.vertices = np.loadtxt(join(cur_dir, 'sphere_vertices_20.txt'))
        # compose faces
        faces_all = []
-        for nj in range(nJoints):
+        for nj in range(self.nJoints):
            faces_all.append(faces + nj*self.vertices.shape[0])
-        for nk in range(len(kintree)):
-            faces_all.append(faces + nJoints*self.vertices.shape[0] + nk*self.vertices.shape[0])
+        for nk in range(len(self.kintree)):
+            faces_all.append(faces + self.nJoints*self.vertices.shape[0] + nk*self.vertices.shape[0])
        self.faces = np.vstack(faces_all)
+        self.nVertices = self.vertices.shape[0] * self.nJoints + self.vertices.shape[0] * len(self.kintree)
+        self.joint_radius = joint_radius

-    def __call__(self, keypoints3d, id=0, return_verts=True, return_tensor=False):
+    def __call__(self, keypoints3d, id=0, return_verts=True, return_tensor=False, **kwargs):
        vertices_all = []
-        r = 0.02
+        r = self.joint_radius
        # joints 
        for nj in range(self.nJoints):
-            if nj > 25:
-                r_ = r * 0.4
-            else:
-                r_ = r
            if keypoints3d[nj, -1] < 0.01:
                vertices_all.append(self.vertices*0.001)
                continue
-            vertices_all.append(self.vertices*r_ + keypoints3d[nj:nj+1, :3])
+            vertices_all.append(self.vertices*r + keypoints3d[nj:nj+1, :3])
        # limb
        for nk, (i, j) in enumerate(self.kintree):
            if keypoints3d[i][-1] < 0.1 or keypoints3d[j][-1] < 0.1:
@ -74,4 +78,25 @@ class SkelModel:
            vertices = self.vertices @ T[:3, :3].T + T[:3, 3:].T
            vertices_all.append(vertices)
        vertices = np.vstack(vertices_all)
-        return vertices[None, :, :]
+        return vertices[None, :, :]
+    
+    def init_params(self, nFrames):
+        return np.zeros((self.nJoints, 4))
+
+class SMPLSKEL:
+    def __init__(self, model_type, gender, body_type) -> None:
+        from ..smplmodel import load_model
+        config = CONFIG[body_type]
+        self.smpl_model = load_model(gender, model_type=model_type, skel_type=body_type)
+        self.body_model = SkelModel(config['nJoints'], config['kintree'])
+
+    def __call__(self, return_verts=True, **kwargs):
+        keypoints3d = self.smpl_model(return_verts=False, return_tensor=False, **kwargs)
+        if not return_verts:
+            return keypoints3d
+        else:
+            verts = self.body_model(return_verts=True, return_tensor=False, keypoints3d=keypoints3d[0])
+            return verts
+    
+    def init_params(self, nFrames):
+        return np.zeros((self.body_model.nJoints, 4))