The architecture was proposed by Alex L. Cheng C, etc. 'Real-time deep hair matting on mobile devices'. This repository is based on https://github.com/aobo-y/hair-dye.
Download dataset from https://github.com/switchablenorms/CelebAMask-HQ and split it in train/val/test folds using https://github.com/switchablenorms/CelebAMask-HQ/tree/master/face_parsing#preprocessing. You should get 23608 in train and 2931 in validation.
Data structure training
├── hair_seg/data/dataset_celeba
│ ├── train
│ │ ├──images
│ │ │ ├── 1.jpg
│ │ │ ├── 2.jpg
...
│ │ ├──masks
│ │ │ ├── 1.png
│ │ │ ├── 2.png
...
│ ├── val
...
$ CUDA_VISIBLE_DEVICES=0 python -u main.py --mode=train --model_name default --print_freq 5 --optimizer adam --lr 1e-3 --wup 2000 --ep 16 --lr_schedule cosine
The checkpoint and sample images are saved in hair_seg/checkpoint/default/
by default.
Checkpoints, logs and tensorboards you can download here.
Trained model gives around 0.89 IoU.
Plot a groundtruth image, the predicted segmentation and the hue adjusted result. With specified image:
$ python main.py --mode=run --checkpoint train_16 --model_name default --image=./path/to/the/image.png
With train or test datasets images:
$ python main.py --mode=run --set=test --num=4 --checkpoint train_16 --model_name default
set
can be one train
and test
, default is test
num
is the random number of images from the set, default is 4