- 6D Rotation Matrix Representation
- High Performance.
- Easy Integration
- Customizability
*** The project is structured into four main files:**
main.py
: The entry point of the project, orchestrating the training, evaluation, and prediction processes.nets.py
: Contains the definitions of neural network models or architectures.datasets.py
: Manages dataset handling, including loading, preprocessing, and augmentations.util.py
: Provides utility functions for data manipulation, visualization, logging, and other support tasks.
The model achieved the following Mean Absolute Error (MAE) metrics across different pose angles:
Backbone | Epochs | Pitch | Yaw | Roll | Params (M) | FLOPS (M) | Pretrained weights |
---|---|---|---|---|---|---|---|
MobileOne_s0 | 120 | 5.4 | 4.2 | 3.8 | 2.1 | 275 | model |
MobileOne_s1 | 120 | 5.2 | 4.0 | 3.7 | 4.8 | 825 | model |
MobileOne_s2 | 120 | 4.9 | 3.7 | 3.4 | 7.8 | 1299 | model |
MobileOne_s3 | 120 | 4.7 | 3.6 | 3.3 | 10.1 | 1896 | model |
MobileOne_s4 | 120 | 4.6 | 3.5 | 3.2 | 14.8 | 2978 | model |
- Clone the repository
- Create a Conda environment using the environment.yml file:
conda env create -f environment.yml
- Activate the Conda environment:
conda activate HPE
- Download the 300W-LP, AFLW2000 Datasets:
- Download the dataset from the official project page.
- Place the downloaded dataset into a directory named 'Datasets'.
To initiate the training process, use the following command:
- Configure your dataset path in main.py for training
- Configure Model name (default A2) in main.py for training
- Run the below command for Single-GPU training
python main.py --train
- Run the below command for Multi-GPU training $ is number of GPUs
bash main.sh $ --train
Configure your dataset path in main.py for testing Run the below command:
bash main.sh $ --test
- Configure your video path in main.py for visualizing the demo
- Run the below command:
python main.py --demo