yakhyo Face Parsing

The model expects a specific directory structure for training datasets. You should organize your images and their corresponding segmentation masks as follows:

• dataset/images/: Contains the input face images (e.g., .jpg).
• dataset/labels/: Contains the corresponding segmentation masks (e.g., .png).

Note: The model was originally trained on the CelebAMask-HQ dataset.

dataset/
├── images/           # Input face images
│   ├── image1.jpg
│   ├── image2.jpg
│   └── ...
└── labels/           # Corresponding segmentation masks
    ├── image1.png
    ├── image2.png
    └── ...

Use train.py to train the BiSeNet model on your dataset. Ensure your dataset is prepared according to the project's dataset requirements before starting. You can specify the backbone (e.g., resnet18 or resnet34), batch size, epochs, and other hyperparameters via CLI arguments.

# Start training with default parameters
python train.py --data-root /path/to/your/dataset

# Custom training example
python train.py --data-root /path/to/dataset --backbone resnet34 --batch-size 8 --epochs 200

Use onnx_export.py to convert trained PyTorch .pt models into ONNX format for cross-platform deployment. The exported .onnx file will be created in the same directory as the source weight file.

# Export ResNet18 model
python onnx_export.py --model resnet18 --weight ./weights/resnet18.pt

# Export ResNet34 model
python onnx_export.py --model resnet34 --weight ./weights/resnet34.pt

Use inference.py to perform face parsing on a single image or a folder of images using PyTorch weights.

# Single image inference
python inference.py --model resnet18 --weight ./weights/resnet18.pt --input ./assets/images/1.jpg --output ./results

# Batch processing (folder of images)
python inference.py --model resnet18 --weight ./weights/resnet18.pt --input ./assets/images --output ./assets/results

To use the Face Parsing Model, clone the repository and install the necessary Python dependencies using pip.

git clone https://github.com/yakhyo/face-parsing.git
cd face-parsing
pip install -r requirements.txt

Use onnx_inference.py to perform face parsing using an exported ONNX model. This is useful for deployment in environments where PyTorch is not available.

python onnx_inference.py --model ./weights/resnet18.onnx --input ./assets/images --output ./assets/results/resnet18onnx

Arguments for inference.py.

options:
  -h, --help       show this help message and exit
  --model MODEL    model name, i.e resnet18, resnet34
  --weight WEIGHT  path to trained model, i.e resnet18/34
  --input INPUT    path to an image or a folder of images
  --output OUTPUT  path to save model outputs

Arguments for onnx_inference.py.

options:
  -h, --help       show this help message and exit
  --model MODEL    path to ONNX model file
  --input INPUT    path to an image or a folder of images
  --output OUTPUT  path to save model outputs

Arguments for onnx_export.py.

options:
  -h, --help       show this help message and exit
  --model MODEL    model name, i.e resnet18, resnet34 (default: resnet18)
  --weight WEIGHT   path to trained PyTorch model (default: ./weights/resnet18.pt)

The repository provides two backbone models for face parsing. Choose between them based on your requirements for parameter count and model size: