Skip to content

Multi-person Human Pose Estimation with HigherHRNet in Pytorch, with TensorRT support

License

GPL-3.0, MIT licenses found

Licenses found

GPL-3.0
LICENSE
MIT
LICENSE.MIT
Notifications You must be signed in to change notification settings

stefanopini/simple-HigherHRNet

Repository files navigation

Multi-person Human Pose Estimation with HigherHRNet in PyTorch

This is an unofficial implementation of the paper HigherHRNet: Scale-Aware Representation Learning for Bottom-Up Human Pose Estimation.
The code is a simplified version of the official code with the ease-of-use in mind.

The code is fully compatible with the official pre-trained weights. It supports both Windows and Linux.

This repository currently provides:

  • A slightly simpler implementation of HigherHRNet in PyTorch (>=1.0) - compatible with official weights (pose_higher_hrnet_*).
  • A simple class (SimpleHigherHRNet) that loads the HigherHRNet network for the bottom-up human pose estimation, loads the pre-trained weights, and make human predictions on a single image or a batch of images.
  • Support for multi-GPU inference.
  • Multi-person support by design (HigherHRNet is a bottom-up approach).
  • A reference code that runs a live demo reading frames from a webcam or a video file.
  • NEW Support for TensorRT (thanks to @gpastal24, see #14 and #15).
  • NEW A Jupyter Notebook compatible with Google Colab showcasing how to use this repository.

This repository is built along the lines of the repository simple-HRNet.
Unfortunately, compared to HRNet, results and performance of HigherHRNet are somewhat disappointing: the network and the required post-processing are slower and the predictions does not look more precise. Moreover, multiple skeletons are often predicted for the same person, requiring additional steps to filter out the redundant poses.
On the other hand, being a bottom-up approach, HigherHRNet does not rely on any person detection algorithm like Yolo-v3 and can be used for person detection too.

Examples

Class usage

import cv2
from SimpleHigherHRNet import SimpleHigherHRNet

model = SimpleHigherHRNet(32, 17, "./weights/pose_higher_hrnet_w32_512.pth")
image = cv2.imread("image.png", cv2.IMREAD_COLOR)

joints = model.predict(image)

The most useful parameters of the __init__ function are:

cnumber of channels (HRNet: 32, 48)
nof_jointsnumber of joints (COCO: 17, CrowdPose: 14)
checkpoint_pathpath of the (official) weights to be loaded
resolutionimage resolution (min side), it depends on the loaded weights
return_heatmapsthe `predict` method returns also the heatmaps
return_bounding_boxesthe `predict` method returns also the bounding boxes
filter_redundant_posesredundant poses (poses being almost identical) are filtered out
max_nof_peoplemaximum number of people in the scene
max_batch_sizemaximum batch size used in hrnet inference
devicedevice (cpu or cuda)

Running the live demo

From a connected camera:

python scripts/live-demo.py --camera_id 0

From a saved video:

python scripts/live-demo.py --filename video.mp4

For help:

python scripts/live-demo.py --help

Extracting keypoints:

From a saved video:

python scripts/extract-keypoints.py --format csv --filename video.mp4

For help:

python scripts/extract-keypoints.py --help

Converting the model to TensorRT:

Warning: require the installation of TensorRT (see Nvidia website) and onnx. On some platforms, they can be installed with

pip install tensorrt onnx

Converting in FP16:

python scripts/export-tensorrt-model.py --device 0 --half

For help:

python scripts/export-tensorrt-model.py --help

Installation instructions

ToDos

  • Add keypoint extraction script (thanks to @wuyenlin)
  • Optimize the post-processing steps
  • Add COCO dataset and evaluation
  • Add Train/Test scripts
  • Add TensorRT support
  • Add notebook compatible with Colab