FaceMask Removal (FaceMask Inpainting)

Environments

• Windows 10
• Pytorch 1.6

CelebA data preparation:

• Download CelebA dataset then crop the image while keeping ratio with here
• Create synthesis facemask segmentation dataset on the cropped set with here (orginal code does not provide binary masks, add it yourself)

• Folder structure:

this repo
│   train.py
│   trainer.py
│   unet_trainer.py
│
└───configs
│      facemask.yaml
│      segm.yaml
│
└───datasets  
│   └───celeba
│       └───images
│           └───celeba512_30k
│           └───celeba512_30k_binary
│           └───celeba512_30k_masked
│       └───annotations
│           |  train.csv
│           |  val.csv

• Put unmasked images in celeba512_30k, facemasked images in celeba512_30k_masked, and binary masks in celeba512_30k_binary
• Split train/validation then save filepaths to .csv. Example:

  ,img_name,mask_name
  0,celeba512_30k_masked\012653_masked.jpg,celeba512_30k_binary\012653_binary.jpg
  1,celeba512_30k_masked\016162_masked.jpg,celeba512_30k_binary\016162_binary.jpg
  2,celeba512_30k_masked\011913_masked.jpg,celeba512_30k_binary\011913_binary.jpg
  

• Edit configs on both segm.yaml and facemask.yaml
• Follow the same steps above when applying custom dataset

Training steps:

• Train segmentation model in unet_trainer.py
• Train inpainting model in trainer.py
• Using infer.py with 2 checkpoints from above tasks to do inference

Train facemask segmentation

python train.py segm --resume=<resume checkpoint>

Train facemask inpainting

python train.py facemask --resume=<resume checkpoint>

Results (100,000 iterations with batch size = 2):

Inpainting results on Masked CelebA-512 (from left to right: FaceMasked - Segmented - Inpainted - Ground Truth)

Free-Form Inpainting results on Places365-256 (from left to right: Ground Truth - Masked - Inpainted )

Paper References:

• Idea and training process from A Novel GAN-Based Network for Unmasking of Masked Face
• Base model from Free-Form Image Inpainting with Gated Convolution

Code References

• Generator from https://github.com/zhaoyuzhi/deepfillv2
• Discriminator from https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix
• https://github.com/avalonstrel/GatedConvolution_pytorch
• https://github.com/LynnHo/HD-CelebA-Cropper
• https://github.com/aqeelanwar/MaskTheFace