6DRepNet: Head Pose Estimation

Features

• 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.

Performance Metrics

The model achieved the following Mean Absolute Error (MAE) metrics across different pose angles:

Results

Backbone	Epochs	Pitch	Yaw	Roll	Params (M)	FLOPS (B)	Pretrained weights
RepVGG-A0	90	5.4	4.3	3.8	9.1	1.5	model
RepVGG-A1	90	5.2	3.9	3.7	14	2.6	model
RepVGG-A2	90	4.8	3.7	3.4	28.2	5.7	model
RepVGG-B0	90	5.0	3.9	3.5	15.8	3.4	model
RepVGG-B1	90	5.0	3.9	3.5	57.4	13.1	model
RepVGG-B1G2	90	4.9	3.6	3.4	45.7	9.8	model
RepVGG-B1G4	90	4.8	3.6	3.4	39.9	8.1	model
RepVGG-B2	90	4.78	3.4	3.3	89	20.4	model
RepVGG-B2G4	90	4.8	3.6	3.4	61.7	12.6	model
RepVGG-B3	90	4.9	3.6	3.4	123	29.2	model
RepVGG-B3G4	90	4.8	3.7	3.4	83.8	17.9	model

Installation

1. Clone the repository
2. Create a Conda environment using the environment.yml file:

conda env create -f environment.yml

3. Activate the Conda environment:

conda activate HPE

Preparing the Dataset

1. Download the 300W-LP, AFLW2000 Datasets:
2. Download the dataset from the official project page.
3. Place the downloaded dataset into a directory named 'Datasets'.

Training the Model

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

Testing the Model

Configure your dataset path in main.py for testing Run the below command:

bash main.sh $ --test

Inference

• Configure your video path in main.py for visualizing the demo
• Run the below command:

python main.py --demo