This repository contains code for training neural network using Pytorch framework. Neural network model is based on Convolutional Pose Machines paper, but instead of body parts network is trained to recognize facial keypoints and output probability map for each keypoint.
Dataset used is Facial Keypints (68). It consists of two folders, test and train, and .csv files with information about each image and it's keypoints locations. During training, train folder is split in two parts (split size defined in config.py). One for training and the other for validation. Test folder can be used after training is complete to test model once more.
Requires: PyTorch 1.9.0
- Download dataset and place it inside project folder.
- Inside config.py file change DATA_ROOT_PATH to coresponding folder name. For example: "data". Rest of the file leave at default, change batch size and number of workers depending on hardware used.
- Run main.py to start training. Model, loss graph and copy of configuration parameters are saved in "output/current_date_and_time" folder in project root. Along with it, examples of some iterations in each epoch are saved, also each new best model.
Example of folder structure:
├── ...
├── data/
│ ├── test/
│ ├── training/
| ├── ...
└── ...
Trained model can be downloaded from this link. Model was trained on image resolution of 320x320.
Inside camera_test.py change:
...
# path to the trained model
model_load_path = './weights/cpm_net_ep60.pt'
...