/tf2-pose-estimation

:dancer: Single pose estimation for iOS and android using TensorFlow 2.0

Primary LanguagePythonApache License 2.0Apache-2.0

💃 Pose Estimation for TensorFlow 2.0

This repository is forked from tucan9389/tf2-mobile-pose-estimation when the original repository doesn't work and the author have no plan to update it, so i fork it and do some research for study purposes.

Even though I made the orginal model worked. But the final result was not as good as I expected. So I import another model from yuanyuanli85/Stacked_Hourglass_Network_Keras. This model can train a better model than previous one. Model details can be found in the file models/hourglass_model_v2.py.

This repository currently implemented the Hourglass model using TensorFlow 2.0 (and Keras API). Instead of normal convolution, inverted residuals (also known as Mobilenet V2) module has been used inside the model for real-time inference.

Table of contents

Goals

  • 📚 Easy to train
  • 🏃‍ Easy to use the model on mobile device

Getting Started

Install Anaconda (~10 min)

Create Virtual Environment (~2 min)

Create new environment.

conda create -n {env_name} python={python_version} anaconda
# in my case
# conda create -n mpe-env-tf2-alpha0 python=3.7 anaconda

Start the environment.

source activate {env_name}
# in my case
# source activate mpe-env-tf2-alpha0
# the newest conda version may use command(on my mac): conda activate {env_name} ## from 2019-12-31

Install the requirements (~1 min)

cd {tf2-mobile-pose-estimation_path}
pip install -r requirements.txt

Run The Project

In order to use the project you have to:

  1. Prepare the dataset(ai_challenger dataset) and locate the dataset on ./datasets.
  2. Run the model using:
python train.py

or run HourglassModel version 2 using:

python train.py --model=hourglass_v2
  1. Monitoring with TensorBoard:
tensorboard --logdir="./outputs/logs"
  1. Test model training result(change model_name in demo.py):
python demo.py

demo-01

Results

hourglass_model.py

  • loss reduced to around 0.0030, it's working, but not good enough.

pred-img-01 loss-img-01 acc-img-01

hourglass_model_v2.py

  • loss reduced to around 0.0016, this model performs better.

pred-img-02 loss-img-02 acc-img-02

Converting To Mobile Model

TensorFLow Lite

  1. Prepare the trained model(.hdf5) from section below Run The Project
  2. Modify trained_model_file_name variable in convert_to_tflite.py to your own trained model.
trained_model_file_name = "{model_file_name}.hdf5"
  1. Run the convert_to_tflite.py:
python convert_to_tflite.py
  1. And then, you can find the .tflite model on {PROJECT_PATH}/outputs/models/tflite/{model_file_name}.tflite.

Details

Folder Structure

├── train.py            - the main script file
├── data_loader.py
├── data_augment.py
├── data_prepare.py
├── model_config.py
├── network_base.py
├── path_manage.py
├── train_config.py
├── requirements.txt
├── models
|   ├── hourglass_model.py
|   └── hourglass_model_v2.py
├── datasets            - this folder contain the datasets of the project.
|   └── ai_challenger
|       ├── ai_challenger_train.json
|       ├── ai_challenger_valid.json
|       ├── train
|       └── valid
└── outputs             - this folder will be generated automatically when start training
    ├── models
    └── logs

TODO

  • Save model(.hdf5 or .ckpt)
  • Convert the model(.hdf5 or .ckpt) to TFLite model(.tflite)
  • Add a demo to check if the model is ok to use.
  • Run the model on Android
  • Run the model on iOS
  • Make DEMO gif running on mobile device
  • multi-person estimation

Acknowledgements

This project is based on tucan9389/tf2-mobile-pose-estimation.

Reference

[1]paper of convolutional pose machines
[2]paper of stack hourglass
[3]paper of mobilenet v2
[4]repository poseestimation-coreml
[5]repository of tf-pose-estimation
[6]devlope guide of tensorflow lite
[7]mace documentation

Related Projects

Other Pose Estimation Projects

Contributing

Any contributions are welcome including improving the project.

License

Apache License 2.0