This repository is the official implementation of the paper A Skeleton-based Rehabilitation Exercise Assessment System with Rotation Invariance, which has been published at TNSRE 2023.
This project is written in Python 3.7.11
and the following is the list of necessary packages and their versions.
pytorch=1.8.0+cu111
scikit-learn=1.0.2
scipy=0.7.3
opencv-python=4.5.4
matplotlib=3.5.1
seaborn=0.11.2
PyYAML=6.0
We conduct our experiments on three datasets: the IRDS and UI-PRMD datasets, and our push-up dataset.
-
The IRDS dataset can be downloaded from its official website.
-
The UI-PRMD dataset can be downloaded from its official website. As the original data is in the form of
Angles
andPositions
, we transform them intoSkeletons
data and the preprocess code is in preprocess.py -
The push-up dataset can be downloaded from our Google Drive.
The datasets should be downloaded to ./data
directory and should be structured as following.
|-- data
|-- IRDS
|-- Simplified
|-- 101_18_0_1_1_stand.txt
|-- 101_18_0_2_1_stand.txt
|-- ...
|-- 307_18_8_12_1_stand.txt
|-- PushUp
|-- S1
|-- 202112291345
|- keypoints_3d.npy
|- labels.txt
|-- ...
|-- ...
|-- info.CSV
|-- UI-PRMD
|-- Correct
|-- Kinect
|-- Angles/
|-- Positions/
|-- Skeletons
|-- m01_s01_e01.txt
|-- m01_s01_e02.txt
|-- ...
|-- m10_s10_e10.txt
|-- Vicon/
|-- Incorrect/
The pretrained models can be downloaded from our Google Drive and should be put into ./logs
directory. The files should be structured as following.
|-- logs
|-- IRDS_GCN_none_seed1
|-- IRDS_m01
|-- model.pth
|-- ...
|-- result.yaml
|-- ...
To train the model from scratch, you can run the following command.
python train.py [--model MODEL] [--dataset DATASET] [--Pnorm]
--model which model is used choose from [gcn, va-gcn, ri-gcn]
--dataset which dataset is used choose from [IRDS, UIPRMD_Kinect, UIPRMD_Vicon, PushUp]
--Pnorm whether to normalize the orientation according to pelvis
Specifically, here are the commands for running different models.
# vanilla ST-GCN
python train.py --model gcn [--dataset DATASET]
# Pnorm
python train.py --model gcn [--dataset DATASET] --Pnorm
# VA-GCN
python train.py --model va-gcn [--dataset DATASET]
# GCN w.t. RI
python train.py --model ri-gcn [--dataset DATASET]
To inference a pretrained model, run this command.
python inference.py [--model MODEL] [--dataset DATASET] [--aug_angle AUG_ANGLES] [--vis]
# e.g. to evalute model under rotational disturbance
python inference.py --model ri-gcn --dataset PushUp --aug_angle 0 10 20 30 40 50 60 70 80 90
Comparisons on the accuracy, recall, precision and OvR metrics can be plotted using code in plot_result.py. For example, here are the bar charts of accuracy, recall, precision and OvR on the UIPRMD_Kinect respectively, which is the same as is reported in our paper.
Email: zkl22@mails.tsinghua.edu.cn