Skip to content

Head Pose Estimation

In this application, we will use both algorithms (6dRepNet and DirectMHP) for head pose estimation Go to /projects/head_pose_estimation

1. Prepare your dataset

For 6DRepNet

Data can be download from: https://drive.google.com/file/d/0B7OEHD3T4eCkVGs0TkhUWFN6N1k/view?resourcekey=0-WT5tO4TOCbNZY6r6z6WmOA
# your dataset structure should be like this (300W_LP is your dataset name)
datasets/300W_LP/
    -AFW/
        -*jpg
        -*.mat
    -HELEN/
        -*jpg
        -*.mat
    ...
    -files.txt

files.txt is a file contain files name from datasets/300W_LP dir. Example:

./AFW/AFW_134212_1_0
./HELEN/HELEN_232194_1_0

For DirectMHP

# your dataset structure should be like this
datasets/
    -images/
        -train/
            -*.jpg
        -validation/
            -*.jpg
    -yolov5_labels/
        -img_txt/
            -train.txt
            -validation.txt
        -train/
            -*.txt
        -validation/
            -*.txt
    ...
    -data.yaml

Data yaml like this:

path: ./headpose/DirectMHP/data/
labels: yolov5_labels
train: yolov5_labels/img_txt/train.txt
val: yolov5_labels/img_txt/validation.txt

nc: 1  # number of classes (only one class: human head)
num_angles: 3  # number of Euler angles is 3 (pitch, yaw, roll)
names: [ 'person' ]  # class names. We still use 'person' in json file

Note

  • train.txt and validation.txt contains the path to the training or validation image
  • labels *.txt is format class_index x y w h pitch yaw roll

2. Usages

  1. For training, run: 
python train.py --opt "<PATH_TO_OPT_FILE>"  --weight-path <PATH_TO_CUSTOM_WEIGHT>

Example:

python train.py --opt opt.json
  • Notes:
  • PATH_TO_OPT_FILE is path to file opt.json

  • PATH_TO_CUSTOM_WEIGHT is path to .pt file

  • For testing, run: 

python predict.py --opt "<PATH_TO_OPT_FILE>" --output "<PATH_TO_OUTPUT_FOLDER>" --input-path "<PATH_TO_TEST_IMG>" --weight-path "<PATH_TO_CUSTOM_WEIGHT>"

Example

python predict.py --input-path .\datasets\300W_LP\AFW\AFW_134212_1_0.jpg --opt opt.json --output 123.png --weight-path .\results\SixDRepNet_1686737564_bs4\epoch_10.pth
  • Notes:
  • PATH_TO_TEST_IMG should be where the image file are
  • PATH_TO_OUTPUT_FOLDER where result will be saved
  • PATH_TO_OPT_FILE is path to file opt.json
  • PATH_TO_CUSTOM_WEIGHT is path to custom weight
  • If the prediction is not good, please consider adding more data

Example opt.json like this:

{   
    "model_type": "sixd_repnet",
    "sixd_repnet":{
        "num_epochs": 10,
        "batch_size": 4,
        "scheduler": false,
        "lr": 0.0001,
        "num_workers": 1,
        "output_path": "results/",
        "dataset": "Pose_300W_LP",
        "data_dir": "datasets/300W_LP",
        "filename_list": "datasets/300W_LP/files.txt"
    },
    "directmhp":{
        "imgsz": 1280,
        "conf_thres": 0.7,
        "iou_thres": 0.45,
        "epochs": 2,
        "batch_size": 1,
        "save_dir": "./train_directmhp",
        "data": "./headpose/DirectMHP/data/300w_lp_aflw2000.yaml",
        "workers": 1
    }
}

You can modify data_dir, output_path, batch_size, num_epochs, lr, filename_list and num_workers. model_type can be "sixd_repnet" or "directmhp"