PersonLab
This is a Keras implementation of PersonLab for Multi-Person Pose Estimation and Instance Segmentation. The model predicts heatmaps and various offsets which allow for computation of joint locations and connections as well as pixel instance ids. See the paper for more details.
Training a model
If you want to use Resnet101 as the base, first download the imagenet initialization weights from here and copy it to your ~/.keras/models/
directory. (Over 100MB files cannot be hosted on github.)
First, construct the dataset in the correct format by running the generate_hdf5.py script. Before running, just set the ANNO_FILE
and IMG_DIR
constants at the top of the script to the paths to the COCO person_keypoints annotation file and the image folder respectively.
Edit the config.py to set options for training, e.g. input resolution, number of GPUs, whether to freeze the batchnorm weights, etc. More advanced options require altering the train.py script. For example, changing the base network can be done by adding an argument to the get_personlab() function, see the documentation there.
After eveything is configured to your liking, go ahead and run the train.py script.
Testing a model
See the demo.ipynb for sample inference and visualizations.
Technical Debts
Several parts of this codebase are borrowed from others. These include:
-
The augmentation code (which is different from the procedure in the PersonLab paper) and data iterator code is heavily borrowed from this fork of the Keras implementation of CMU's "Realtime Multi-Person Pose Estimation". (The pose plotting function is also influenced by the one in that repo.)
-
The Polyak Averaging callback is just a lightly modified version of the EMA callback from here
Environment
This code was tested in the following environment and with the following software versions:
- Ubuntu 16.04
- CUDA 8.0 with cudNN 6.0
- Python 2.7
- Tensorflow 1.7
- Keras 2.1.3
- OpenCV 2.4.9