How to Train the PFLD Model on a Custom Dataset: A Complete Guide

To train the PFLD model on a custom dataset, format your annotations to match the WFLW dataset structure with 98 facial landmarks, run data/SetPreparation.py to preprocess and augment the data, then execute train_model.py pointing to the generated list.txt files.

The guoqiangqi/pfld repository implements a lightweight, real-time facial landmark detector using a MobileNet-V2 backbone. When you train the PFLD model on a custom dataset, you must adhere to the specific data formatting and preprocessing pipeline originally designed for the WFLW dataset to ensure compatibility with the loss functions defined in model2.py.

Understanding PFLD Data Format Requirements

The PFLD training pipeline expects input data that mirrors the WFLW dataset format exactly. Your custom dataset must provide raw images and a structured annotation file with specific numerical fields.

Required Directory Structure

Organize your raw data as follows before running any preprocessing:

my_dataset/
├── images/               # Raw images readable by OpenCV (.jpg, .png)

│   ├── img_001.jpg
│   └── img_002.jpg
└── landmarks.txt         # WFLW-style annotation

Annotation File Format (WFLW-Style)

The annotation file must contain one line per image with exactly 207 fields:

1. Image path – Relative or absolute path to the image file.
2. 196 landmark values – 98 facial landmarks, each with normalized x and y coordinates in the range [0, 1].
3. 6 attribute flags – Binary indicators for pose, expression, illumination, make_up, occlusion, and blur (0 or 1).
4. 3 Euler angles – Float values for pitch, yaw, and roll.

Example line format:

<full_image_path> <x1> <y1> <x2> <y2> ... <x98> <y98> <pose> <expression> <illumination> <make_up> <occlusion> <blur> <pitch> <yaw> <roll>

Preprocessing Your Custom Dataset with SetPreparation.py

The repository includes data/SetPreparation.py to handle the heavy lifting of converting raw WFLW-style annotations into the exact format required by the training loop.

What the Conversion Script Does

When you run SetPreparation.py, it performs the following operations defined in the source:

• Face cropping – Expands the tight bounding box around landmarks by a factor of 1.2, pads when the crop exceeds image borders, and resizes to the target image_size (default 112×112).
• Landmark normalization – Converts absolute pixel coordinates to the [0, 1] range by dividing by the crop box size (boxsize).
• Data augmentation – Generates up to repeat = 10 augmented samples per image, applying random rotations (± 20°), random center jitter, and optional horizontal mirroring (using the Mirror98.txt index file).
• Euler angle computation – Extracts 14 tracked landmarks and passes them to euler_angles_utils.calculate_pitch_yaw_roll to generate the three angle values required for training.
• Attribute preservation – Copies the six binary attribute flags unchanged into the output label file.

Running the Preprocessing Command

Execute the script from the repository root, pointing to your custom dataset:

python data/SetPreparation.py \
    --image_dir my_dataset/images \
    --landmark_file my_dataset/landmarks.txt \
    --output_dir my_dataset/prepared

This creates my_dataset/prepared/train_data/ and my_dataset/prepared/test_data/, each containing an imgs/ folder and a list.txt file.

Verify the generated labels contain 207 fields:

head -n 3 my_dataset/prepared/train_data/list.txt

Each line should follow this pattern:

/abs/path/to/imgs/0_0.png 0.12 0.34 ... 0 1 0 0 1 0 0.1 -0.2 0.0

Launching Training on Custom Data

With preprocessed data ready, use train_model.py to start training the MobileNet-V2 backbone defined in model2.py.

Command Line Configuration

The training script accepts several critical arguments:

• --file_list – Path to the training list.txt (e.g., my_dataset/prepared/train_data/list.txt).
• --test_list – Path to the validation list.txt.
• --model_dir – Directory to save checkpoints (e.g., my_custom_model).
• --max_epoch – Total training epochs (default 200).
• --batch_size – Batch size (e.g., 64 or 32 depending on GPU memory).
• --image_size – Input resolution (default 112).
• --pretrained_model – Path to a checkpoint to resume from (optional).

Example training command:

python train_model.py \
    --file_list my_dataset/prepared/train_data/list.txt \
    --test_list my_dataset/prepared/test_data/list.txt \
    --model_dir my_custom_model \
    --max_epoch 200 \
    --batch_size 64 \
    --image_size 112

Using the Train Shell Wrapper

Alternatively, edit the provided train.sh wrapper to point to your custom paths:

# train.sh (excerpt)

python train_model.py \
    --file_list my_dataset/prepared/train_data/list.txt \
    --test_list my_dataset/prepared/test_data/list.txt \
    --model_dir my_custom_model \
    "$@"

Then launch with:

bash train.sh --max_epoch 200 --batch_size 64

Monitoring and Evaluation

The training script writes TensorBoard logs to ./tensorboard. Launch the dashboard to visualize loss curves, learning rate schedules from utils.py, and evaluation metrics:

tensorboard --logdir=./tensorboard

Navigate to http://localhost:6006 to monitor loss, learning rate, and validation accuracy in real time.

After training completes, checkpoints are saved in my_custom_model as model.ckpt-<epoch>. Run inference on new images using test_model.py:

python test_model.py \
    --model_dir my_custom_model \
    --image_path my_dataset/images/example.jpg

This loads the trained graph and outputs the 98 facial landmarks predicted by the PFLD model.

Summary

• Format your annotations to match the WFLW structure: 196 normalized landmark values (98 points), 6 binary attributes, and 3 Euler angles (207 total fields per line).
• Run data/SetPreparation.py to crop faces (1.2× bounding box expansion), normalize coordinates to [0, 1], augment data (±20° rotation, mirroring via Mirror98.txt), and compute angles via euler_angles_utils.calculate_pitch_yaw_roll.
• Execute train_model.py with --file_list and --test_list pointing to the generated list.txt files, adjusting --batch_size and --max_epoch for your hardware.
• Monitor training via TensorBoard logs in ./tensorboard and evaluate using test_model.py.

Frequently Asked Questions

How do I prepare annotations if my dataset has fewer than 98 landmarks?

The PFLD architecture in model2.py expects exactly 98 landmarks (196 values). If your dataset contains fewer points, you must map your existing landmarks to the 98-point format by interpolating missing points, or modify the final output layer in model2.py and recompile. The preprocessing script SetPreparation.py will raise an error if the landmark count does not match 196 values per sample.

Can I train PFLD without the 6 attribute flags or Euler angles?

Yes, but you must modify the data pipeline. The train_model.py script expects all 207 fields by default. If you lack attributes or angles, modify SetPreparation.py to output dummy zero values for the missing fields, or edit the dataset loading logic in train_model.py to parse fewer columns. Note that the Euler angle loss in model2.py contributes to pose-aware training stability; zeroing these labels may reduce accuracy for extreme head poses.

What image resolution should I use for custom training?

The default image_size in both SetPreparation.py and train_model.py is 112×112 pixels, matching the MobileNet-V2 backbone input in model2.py. You can train at higher resolutions (e.g., 224×224) by passing --image_size 224 to both the preprocessing and training scripts, but you must reduce --batch_size proportionally to avoid out-of-memory errors and may need to adjust the learning rate schedule defined in utils.py.

How do I resume training from a checkpoint?

Pass the --pretrained_model argument to train_model.py pointing to the specific checkpoint file. For example:

python train_model.py \
    --pretrained_model my_custom_model/model.ckpt-150 \
    --file_list my_dataset/prepared/train_data/list.txt \
    --model_dir my_custom_model_resumed

The script loads the TensorFlow checkpoint and continues optimization from epoch 150, preserving the learned weights from the MobileNet-V2 backbone and auxiliary head defined in model2.py.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}