TensorRT Implementation of OpenPose

This repo provides the C++ implementation of OpenPose for doing real-time pose estimation on Windows platform. It is based on pose detection program developed by NVIDIA and performs pose estimation on TensorRT framework with a high throughtput. The pose estimation program runs at up to 500 FPS on RTX-3070 GPU using 224x224 ResNet18.

Requirements

The following environment was set for the experiment but if you have a different Graphic Card, you need to download and install TensorRT / CUDA that matches your GPU version.

• Windows 10
• Visual Studio 2017
• RTX 3070 GPU
• TensorRT 7.2.1.6
• CUDA 11.1, Cudnn 8
• Python 3.7
• Torch 1.8.1
• OpenCV 4.5.1 with CUDA

Installation

• Install PyTorch

pip3 install torch==1.8.1+cu111 torchvision==0.9.1+cu111 torchaudio===0.8.1 -f https://download.pytorch.org/whl/torch_stable.html

• Download and install CUDA and CUDNN by following this installation guide.
• Install [OpenCV 4.5.1 or 4.5.0]
• Download the TensorRT zip file that matches your Windows version.
• Install TensorRT by copying the DLL files from <tensorrt_path>/lib to your CUDA installation directory. For more information, refer to TensorRT installation guideline.

move <TensorRT_Installpath>/lib/*.dll "C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v11.1/bin"

• Install trt_pose for parsing the trained pytorch model to an onnx graph model. For more information, refer to trt_pose installation guide.

git clone https://github.com/NVIDIA-AI-IOT/trt_pose
cd trt_pose
python setup.py install

• Clone this github repository

cd ..
git clone https://github.com/spacewalk01/TensorRT-Openpose
cd Human_pose_detection

• Download a pretrained resnet model, resnet18_baseline_att_224x224_A and put it in the current project folder.
• Create an onnx model, optimize the model using TensorRT and build a runtime inference engine.

python convert2onnx.py -i resnet18_baseline_att_224x224_A_epoch_249.pth -o trt_pose.onnx
<tensorrt_path>/bin/trtexec.exe --onnx=trt_pose.onnx --explicitBatch --saveEngine=trt_pose_fp16.engine --fp16

• Open the solution with Visual Studio. Select x64 and Release for the configuration and start building the project.

Usage

#include <opencv2/core.hpp>
#include "opencv2/imgcodecs.hpp"
#include "opencv2/highgui.hpp"
#include "opencv2/imgproc.hpp"
#include "Openpose.h"
#include "TensorrtPoseNet.h"

int main() {

    std::string filepath = "your image path";
   
    // Initialize pose network
    TensorrtPoseNet posenet("trt_pose_fp16.engine", 0.9, 0.1);
    
    // Initialize openpose parser
    Openpose openpose(posenet.outputDims[0]);

    cv::Mat image = cv::imread(filepath);

    // First, openpose uses resnet to predict a heatmap for each keypoint on the image
    net.infer(image);  

    // Then, it groups the keypoints to estimate a skeleton pose for each person
    openpose.detect(posenet.cpuCmapBuffer, posenet.cpuPafBuffer, image); 
    
    cv::imshow("Result", image);
    cv::waitKey(0);
      
    return 0;
}

Training

• For training a larger model, you may refer to link

References

• https://github.com/NVIDIA-AI-IOT/trt_pose - Real-time pose estimation (Python and C++)
• https://github.com/CaoWGG/TensorRT-YOLOv4 - Object detection based on Tensorrt (C++)