This repository contains code to train a deep learning model that aims to predict the 3D position of anatomical markers (red dots on the skeleton on the right below) from the 3D position of video keypoints (red dots on the skeleton on the left below). The trained model was deployed as part of OpenCap (source code here).
The 2D position of video keypoints can be identified from videos using pose estimation models (eg, OpenPose). Their 3D position can then be reconstructed using triangulation if videos from multiple views are available. However, the set of markers identified from common pose estimation models is typically sparse (eg, there are no markers between the shoulders and the hips), which makes difficult the robust estimation of 3D joint kinematics (joint angles). By augmenting the video-based marker set into a more comprehensive set of anatomical markers, we address this limitation and improve the accuracy of joint kinematic estimates (by in average 3.4 deg and up to 32.6 deg). More details in this preprint and this conference talk.
To train the model, we used over 100 hours of synthetic data. The dataset was generated by processing motion capture data with OpenSim.
- Install Anaconda
- Clone the repository to your machine.
- Open the Anaconda command prompt and create a conda environment:
conda create -n cs-230-project
- Activate the environment:
conda activate cs-230-project
- Install tensorflow with GPU support (you can probably also use CPU only if you train on a small part of the dataset).
- Install other dependencies. Navigate to the local directory where the repository is cloned, then:
python -m pip install -r requirements.txt
- A subset of the dataset is available here. Download it into augmenter-cs230/Data. The path of the first feature time sequence should be something like /augmenter-cs230/Data/data_CS230/feature_0.npy.
- The full dataset is available here. Download it into augmenter-cs230/Data. The path of the first feature time sequence should be something like /augmenter-cs230/Data/data_CS230/feature_0.npy.
Main files:
trainLSTM.py
: script to train LSTM model.evaluateLSTM.py
: script to evaluate LSTM model.testLSTM.py
: script to test LSTM model on failure example.tuneHyperParametersLSTM.py
: script to tune hyperparameters of LSTM model.myModels.py
: script describing the model architecture.myModelsHyperParameters.py
: script describing the model architecture with support for hyperparameters tuning.myDataGenerator.py
: script describing the data generator.mySettings.py
: script with some tunable model settings.
Other files (you should not need to interact with these files):
splitData.py
: script to split the data into different sets.infoDatasets.py
: some details about the dataset.utilities.py
: various utilities.
- You can also use docker. The following instructions worked for me on a machine where cuda and cudnn were installed, you should adjust the paths and names according to your working environment.
- Build image: eg
docker build -t antoine/tensorflow:latest-gpu-0.1 .
- Train model: eg
sudo docker run --gpus all -it --rm -v /home/clarkadmin/Documents/MyRepositories/augmenter-cs230:/augmenter-cs230 antoine/tensorflow:latest-gpu-0.1 python augmenter-cs230/trainLSTM.py
- Evaluate model: eg
sudo docker run --gpus all -it --rm -v /home/clarkadmin/Documents/MyRepositories/augmenter-cs230:/augmenter-cs230 antoine/tensorflow:latest-gpu-0.1 python augmenter-cs230/evaluateLSTM.py
- Test model: eg
sudo docker run --gpus all -it --rm -v /home/clarkadmin/Documents/MyRepositories/augmenter-cs230:/augmenter-cs230 antoine/tensorflow:latest-gpu-0.1 python augmenter-cs230/testLSTM.py
- Tune hyperparameters: eg
sudo docker run --gpus all -it --rm -v /home/clarkadmin/Documents/MyRepositories/augmenter-cs230:/augmenter-cs230 antoine/tensorflow:latest-gpu-0.1 python augmenter-cs230/tuneHyperParametersLSTM.py
- The files might be locked after training, to unlock run
sudo chown -R $USER: $HOME