[Service] Websockect server and client for SMPL(X) verts inference (#125

)

* use uuid in server, disconnect works not properly

* Send bytes instead of json encoding list for verts

* add docs for service

* add dockerfile and docker image for server

* timer is not thread safe, lock it

* make sending verts in bytes an option

* [Fix] Update SMPL server for bridge server (#10)

Modified SMPL server so that it is compatible with the bridge server. See commit history for more details.

* ready for windows test

* Fix installation docs

* Fix PR comments

* Rename verts to stream in docs

---------

Co-authored-by: teemo <35056621+WYK96@users.noreply.github.com>

configs/modules/service/smpl_stream_service.py

@@ -0,0 +1,9 @@
+type = 'SMPLStreamService'
+name = 'smpl_stream_service'
+work_dir = f'temp/{name}'
+body_model_dir = 'xrmocap_data/body_models'
+device = 'cuda:0'
+enable_bytes = True
+enable_cors = True
+port = 29091
+max_http_buffer_size = 128 * 1024 * 1024

dockerfiles/service_ubt18/Dockerfile

@@ -0,0 +1,8 @@
+ARG INPUT_TAG
+FROM $INPUT_TAG
+
+# Install test requirements
+RUN . /opt/miniconda/etc/profile.d/conda.sh && \
+    conda activate openxrlab && \
+    pip install -r https://raw.githubusercontent.com/openxrlab/xrmocap/main/requirements/service.txt && \
+    pip cache purge

dockerfiles/service_ubt18/build_runtime_docker.sh

@@ -0,0 +1,17 @@
+#!/bin/bash
+CUDA_VER=11.6
+PY_VER=3.8
+MMCV_VER=1.6.1
+TORCH_VER=1.12.1
+TORCHV_VER=0.13.1
+CUDA_VER_DIGIT=${CUDA_VER//./}
+PY_VER_DIGIT=${PY_VER//./}
+MMCV_VER_DIGIT=${MMCV_VER//./}
+TORCH_VER_DIGIT=${TORCH_VER//./}
+INPUT_TAG="openxrlab/xrmocap_runtime:ubuntu1804_x64_cuda${CUDA_VER_DIGIT}_py${PY_VER_DIGIT}_torch${TORCH_VER_DIGIT}_mmcv${MMCV_VER_DIGIT}"
+FINAL_TAG="${INPUT_TAG}_service"
+echo "tag to build: $FINAL_TAG"
+BUILD_ARGS="--build-arg CUDA_VER=${CUDA_VER} --build-arg PY_VER=${PY_VER} --build-arg MMCV_VER=${MMCV_VER} --build-arg TORCH_VER=${TORCH_VER} --build-arg TORCHV_VER=${TORCHV_VER} --build-arg INPUT_TAG=${INPUT_TAG}"
+# build according to Dockerfile
+docker build -t $FINAL_TAG -f dockerfiles/service_ubt18/Dockerfile $BUILD_ARGS --progress=plain .
+echo "Successfully tagged $FINAL_TAG"

docs/en/installation.md

@@ -5,6 +5,7 @@
 - [Prepare environment](#prepare-environment)
 - [Run with docker image](#run-with-docker-image)
 - [Test environment](#test-environment)
+- [Client only](#client-only)
 - [Frequently Asked Questions](#frequently-asked-questions)
 
 ## Requirements
@@ -21,14 +22,15 @@
 
 Optional:
 
-| Name                                                     | When it is required            | What's important                                             |
-| :------------------------------------------------------- | :----------------------------- | :----------------------------------------------------------- |
-| [MMPose](https://github.com/open-mmlab/mmpose)           | Keypoints 2D estimation.       | Install `mmcv-full`, instead of `mmcv`.                      |
-| [MMDetection](https://github.com/open-mmlab/mmdetection) | Bbox 2D estimation.            | Install `mmcv-full`, instead of `mmcv`.                      |
-| [MMTracking](https://github.com/open-mmlab/mmtracking)   | Multiple object tracking.      | Install `mmcv-full`, instead of `mmcv`.                      |
-| [MMDeploy](https://github.com/open-mmlab/mmdeploy)       | Faster mmdet+mmpose inference. | Install `mmcv-full`, `cudnn` and `TensorRT`.                 |
-| [Aniposelib](https://github.com/google/aistplusplus_api) | Triangulation.                 | Install from [github](https://github.com/liruilong940607/aniposelib), instead of pypi. |
-| [Minimal Pytorch Rasterizer](https://github.com/rmbashirov/minimal_pytorch_rasterizer) | SMPL mesh fast visualization.                 | Tested on torch-1.12.0. |
+| Name                                                         | When it is required                     | What's important                                             |
+| :----------------------------------------------------------- | :-------------------------------------- | :----------------------------------------------------------- |
+| [MMPose](https://github.com/open-mmlab/mmpose)               | Keypoints 2D estimation.                | Install `mmcv-full`, instead of `mmcv`.                      |
+| [MMDetection](https://github.com/open-mmlab/mmdetection)     | Bbox 2D estimation.                     | Install `mmcv-full`, instead of `mmcv`.                      |
+| [MMTracking](https://github.com/open-mmlab/mmtracking)       | Multiple object tracking.               | Install `mmcv-full`, instead of `mmcv`.                      |
+| [MMDeploy](https://github.com/open-mmlab/mmdeploy)           | Faster mmdet+mmpose inference.          | Install `mmcv-full`, `cudnn` and `TensorRT`.                 |
+| [Aniposelib](https://github.com/google/aistplusplus_api)     | Triangulation.                          | Install from [github](https://github.com/liruilong940607/aniposelib), instead of pypi. |
+| [Minimal Pytorch Rasterizer](https://github.com/rmbashirov/minimal_pytorch_rasterizer) | SMPL mesh fast visualization.           | Tested on torch-1.12.0.                                      |
+| [Flask](https://flask.palletsprojects.com/en/2.3.x/)         | Starting an http or a websocket server. |                                                              |
 
 ## A from-scratch setup script
 
@@ -92,6 +94,8 @@ cd xrmocap
 pip install -r requirements/build.txt
 # install requirements for runtime
 pip install -r requirements/runtime.txt
+# install requirements for services
+pip install -r requirements/service.txt
 
 # install xrmocap
 rm -rf .eggs && pip install -e .
@@ -196,7 +200,15 @@ cd /opt && \
 
 **Note3:** We've only tested mmdeploy 0.12.0, other version may not work as expectation.
 
-#### g. Run unittests or demos
+#### g. Install requirements for service
+
+You will only need this when you are going to start a server defined in `xrmocap.service`.
+
+```bash
+pip install -r requirements/service.txt
+```
+
+#### h. Run unittests or demos
 
 If everything goes well, try to [run unittest](#test-environment) or go back to [run demos](./getting_started.md#inference)
 
@@ -225,6 +237,15 @@ sh scripts/run_docker.sh
 
 To test whether the environment is well installed, please refer to [test doc](./test.md).
 
+### Client only
+
+If you only need to use the client provided by XRMoCap, the installation process will be much simpler. We have increased the compatibility of the client by reducing dependencies, and you only need to execute the commands below.
+
+```bash
+pip install numpy tqdm flask-socketio requests websocket-client
+pip install . --no-deps
+```
+
 ### Frequently Asked Questions
 
 If your environment fails, check our [FAQ](./faq.md) first, it might be helpful to some typical questions.

docs/en/tools/start_service.md

@@ -0,0 +1,40 @@
+# Tool start_service
+
+- [Overview](#overview)
+- [Argument: config_path](#argument-config_path)
+- [Argument: disable_log_file](#argument-disable_log_file)
+- [Example](#example)
+
+### Overview
+
+This tool starts a server in the current console according to the configuration file, and sets up a logger. The logger displays information of no less than `INFO` level in the console, and write information of no less than `DEBUG` level in the log file under the `logs/` directory.
+
+For services that use the `work_dir` parameter, please make sure that the target path can be created correctly. Generally speaking, running `mkdir temp` in advance can ensure that the default configuration file in the repository can be successfully used.
+
+### Argument: config_path
+
+`config_path` is the path to a configuration file for server. Please ensure that all parameters required by `SomeService.__init__()` are specified in the configuration file. An example is provided below. For more details, see the docstring in [code](../../../xrmocap/service/base_flask_service.py).
+
+```python
+type = 'SMPLStreamService'
+name = 'smpl_stream_service'
+work_dir = f'temp/{name}'
+body_model_dir = 'xrmocap_data/body_models'
+device = 'cuda:0'
+enable_cors = True
+port = 29091
+```
+
+Also, you can find our prepared config files in `configs/modules/service/smpl_stream_service.py`.
+
+### Argument: disable_log_file
+
+By default, `disable_log_file` is False and two log files under `logs/f'{service_name}_{time_str}'` will be written. Add `--disable_log_file` makes it True and the tool will only print log to console.
+
+### Example
+
+Run the tool with explicit paths.
+
+```bash
+python tools/start_service.py --config_path configs/modules/service/smpl_stream_service.py
+```

requirements/service.txt

@@ -0,0 +1,6 @@
+flask
+Flask-Caching
+flask-socketio
+flask_api
+flask_cors
+simple-websocket

scripts/start_service_docker.sh

@@ -0,0 +1,15 @@
+#!/usr/bin/env bash
+TAG=openxrlab/xrmocap_runtime:ubuntu1804_x64_cuda116_py38_torch1121_mmcv161_service
+CONFIG_PATH=$1
+PORT=$(grep 'port =' ${CONFIG_PATH} | cut -d "=" -f 2 | tr -d ' ')
+echo "Starting service on port $PORT"
+PORTS="-p $PORT:$PORT"
+WORKSPACE_VOLUMES="-v $PWD:/workspace/xrmocap"
+WORKDIR="-w /workspace/xrmocap"
+MEMORY="--memory=20g"
+docker run --runtime=nvidia -it --rm --entrypoint=/bin/bash $PORTS $WORKSPACE_VOLUMES $WORKDIR $MEMORY $TAG -c "
+  source /opt/miniconda/etc/profile.d/conda.sh
+  conda activate openxrlab
+  pip install .
+  python tools/start_service.py --config_path $CONFIG_PATH
+"

setup.cfg

@@ -12,6 +12,6 @@ multi_line_output = 5
 include_trailing_comma = true
 known_standard_library = pkg_resources,setuptools
 known_first_party = xrmocap
-known_third_party =PIL,cv2,dateutil,filterpy,matplotlib,mmcv,mmhuman3d,numpy,prettytable,pytest,pytorch3d,scipy,smplx,sphinx_rtd_theme,torch,torchvision,tqdm,xrprimer
+known_third_party =PIL,cv2,dateutil,filterpy,flask,flask_socketio,matplotlib,mmcv,mmhuman3d,numpy,prettytable,pytest,pytorch3d,scipy,smplx,socketio,sphinx_rtd_theme,torch,torchvision,tqdm,xrprimer
 no_lines_before = STDLIB,LOCALFOLDER
 default_section = THIRDPARTY

tools/clients/smpl_verts_client.py

@@ -0,0 +1,76 @@
+# yapf: disable
+import argparse
+import logging
+import numpy as np
+import os
+import sys
+import time
+from tqdm import tqdm
+
+from xrmocap.client.smpl_stream_client import SMPLStreamClient
+
+# yapf: enable
+
+
+def main(args) -> int:
+    name = os.path.basename(__file__).split('.')[0]
+    logger = logging.getLogger(name)
+    if args.verbose:
+        logger.setLevel(logging.INFO)
+    else:
+        logger.setLevel(logging.WARNING)
+    if args.smpl_data_path is None:
+        logger.error('Please specify smpl_data_path.')
+        raise ValueError
+    client = SMPLStreamClient(
+        server_ip=args.server_ip, server_port=args.server_port, logger=logger)
+    n_frames = client.upload_smpl_data(args.smpl_data_path)
+    logger.info(f'Motion of {n_frames} frames uploaded.')
+    faces = client.get_faces()
+    faces_np = np.array(faces)
+    logger.info(f'Get faces: {faces_np.shape}')
+    start_time = time.time()
+    for frame_idx in tqdm(range(n_frames)):
+        verts = client.forward(frame_idx)
+        if frame_idx == 0:
+            verts_np = np.array(verts)
+            logger.info(f'Get verts for first frame: {verts_np.shape}')
+    loop_time = time.time() - start_time
+    fps = n_frames / loop_time
+    logger.info(f'Get verts for all frames, average fps: {fps:.2f}')
+    client.close()
+    return 0
+
+
+def setup_parser():
+    parser = argparse.ArgumentParser(
+        description='Send a smpl data file to ' +
+        'SMPLStreamServer and receive faces and verts.')
+    parser.add_argument(
+        '--smpl_data_path',
+        help='Path to a SMPL(X)Data file.',
+        type=str,
+    )
+    parser.add_argument(
+        '--server_ip',
+        help='IP address of the server.',
+        type=str,
+        default='127.0.0.1')
+    parser.add_argument(
+        '--server_port',
+        help='Port number of the server.',
+        type=int,
+        default=29091)
+    parser.add_argument(
+        '--verbose',
+        action='store_true',
+        help='If True, INFO level log will be shown.',
+        default=False)
+    args = parser.parse_args()
+    return args
+
+
+if __name__ == '__main__':
+    args = setup_parser()
+    ret_val = main(args)
+    sys.exit(ret_val)

tools/start_service.py

@@ -0,0 +1,71 @@
+# yapf: disable
+import argparse
+import json
+import mmcv
+import os
+from xrprimer.utils.log_utils import logging, setup_logger
+
+from xrmocap.service.builder import build_service
+from xrmocap.utils.date_utils import get_datetime_local, get_str_from_datetime
+
+# yapf: enable
+
+
+def main(args):
+    # load config
+    service_config = dict(mmcv.Config.fromfile(args.config_path))
+    service_name = service_config['name']
+    # setup logger
+    if not args.disable_log_file:
+        datetime = get_datetime_local()
+        time_str = get_str_from_datetime(datetime)
+        log_dir = os.path.join('logs', f'{service_name}_{time_str}')
+        os.makedirs(log_dir)
+    main_logger_path = None \
+        if args.disable_log_file\
+        else os.path.join(log_dir, f'{service_name}_log.txt')
+    flask_logger_path = None \
+        if args.disable_log_file\
+        else os.path.join(log_dir, 'flask_log.txt')
+    logger = setup_logger(
+        logger_name=service_name,
+        file_level=logging.DEBUG,
+        console_level=logging.INFO,
+        logger_path=main_logger_path)
+    # logger for Flask
+    flask_logger = setup_logger(
+        logger_name='werkzeug',
+        file_level=logging.DEBUG,
+        console_level=logging.INFO,
+        logger_path=flask_logger_path)
+    logger.info('Main logger starts.')
+    flask_logger.info('Flask logger starts.')
+    # build service
+    service_config_str = json.dumps(service_config, indent=4)
+    logger.debug(f'\nservice_config:\n{service_config_str}')
+    service_config['logger'] = logger
+    service = build_service(service_config)
+    service.run()
+
+
+def setup_parser():
+    parser = argparse.ArgumentParser()
+    # input args
+    parser.add_argument(
+        '--config_path',
+        type=str,
+        help='Path to service config file.',
+        default='configs/modules/service/base_service.py')
+    # log args
+    parser.add_argument(
+        '--disable_log_file',
+        action='store_true',
+        help='If checked, log will not be written as file.',
+        default=False)
+    args = parser.parse_args()
+    return args
+
+
+if __name__ == '__main__':
+    args = setup_parser()
+    main(args)

xrmocap/client/__init__.py

@@ -0,0 +1,7 @@
+import logging
+
+# client does not require xrprimer.utils.log_utils
+# logger's level is set to INFO by default
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')