6DRepNet Documentation

You can install the SixDRepNet package directly from PyPI for easy access and usage.

pip3 install sixdrepnet

To train or test the model, you must download the 300W-LP, AFLW2000, or BIWI datasets and store them in a datasets directory.

• For 300W-LP and AFLW2000: Generate a filename list using create_filename_list.py.
• For BIWI: Preprocess the dataset by cropping faces using a face detector (recommended crop size is 256).

# Create filename list for 300W-LP
python create_filename_list.py --root_dir datasets/300W_LP

To train the model, first download the pre-trained RepVGG model 'RepVGG-B1g2-train.pth' and place it in the root directory. Then run the training script.

python sixdrepnet/train.py

To optimize trained models for inference, use the convert.py script to reparameterize them. When loading the resulting model, set deploy=True in the SixDRepNet constructor.

# Convert model
python convert.py input-model.tar output-model.pth

# Load deployed model
model = SixDRepNet(backbone_name='RepVGG-B1g2',
                    backbone_file='',
                    deploy=True,
                    pretrained=False)

To set up the repository manually, clone the repository, create a virtual environment, and install the required dependencies. Note that running demo scripts requires an additional face detector installation.

git clone https://github.com/thohemp/6DRepNet
cd 6DRepNet

# Set up a virtual environment
python3 -m venv venv
source venv/bin/activate
pip install -r requirements.txt

# Install face detector for demo scripts
pip install git+https://github.com/elliottzheng/face-detection.git@master

Run the testing script by providing the batch size, dataset name, data directory, filename list, and the model snapshot path.

python test.py  --batch_size 64 \ 
                --dataset AFLW2000 \ 
                --data_dir datasets/AFLW2000 \ 
                --filename_list datasets/AFLW2000/files.txt \ 
                --snapshot output/snapshots/1.pth \ 
                --show_viz False

Use the SixDRepNet class to perform head pose estimation. Weights are automatically downloaded upon initialization. The predict method returns the pitch, yaw, and roll angles, and draw_axis can be used to visualize the pose on the image.

# Import SixDRepNet
from sixdrepnet import SixDRepNet
import cv2


# Create model
# Weights are automatically downloaded
model = SixDRepNet()

img = cv2.imread('/path/to/image.jpg')

pitch, yaw, roll = model.predict(img)

model.draw_axis(img, yaw, pitch, roll)

cv2.imshow("test_window", img)
cv2.waitKey(0)

Run the camera demo script by specifying a model snapshot and the camera index (e.g., 0 for the default webcam).

python ./sixdrepnet/demo.py  --snapshot 6DRepNet_300W_LP_AFLW2000.pth \ 
                             --cam 0