Home

Awesome

SPARNet: Learning Spatial Attention for Face Super-Resolution in PyTorch

Learning Spatial Attention for Face Super-Resolution
Chaofeng Chen, Dihong Gong, Hao Wang, Zhifeng Li, Kwan-Yee K. Wong

example result

Installation and Requirements

Clone this repository

git clone https://github.com/chaofengc/Face-SPARNet.git
cd Face-SPARNet

I have tested the codes on

Getting Started

Download Pretrain Models and Dataset

Download the pretrained models and data from the following link and put them to ./pretrain_models and ./test_dirs respectively

Test with Pretrained Models

We provide example test commands in script test.sh for both SPARNet and SPARNetHD. Two models with difference configurations are provided for each of them, refer to section below to see the differences. Here are some test tips:

We also provide command to crop and align faces from single image, and then paste them back, the same as PSFRGAN

python test_enhance_single_unalign.py --gpus 1 --model sparnethd --name SPARNetHD_V4_Attn2D \
    --res_depth 10 --att_name spar --Gnorm 'in' \
    --pretrain_model_path [./path/to/model/SPARNetHD_V4_Attn2D_net_H-epoch10.pth] \
    --test_img_path ./test_images/test_hzgg.jpg --results_dir test_hzgg_results

Train the Model

The commands used to train the released models are provided in script train.sh. Here are some train tips:

Differences with the Paper

Since the original codes are messed up, we rewrite the codes and retrain all models. This leads to slightly different results between the released model and those reported in the paper. Besides, we also extend the 2D spatial attention to 3D attention, and release some models with 3D attention. We list all of them below

SPARNet

We found that extending 2D spatial attention to 3D attention improves the performance a lot. We trained a light model with half parameter number by reducing the number of FAU blocks, denoted as SPARNet-Light-Attn3D. SPARNet-Light-Attn3D shows similar performance with SPARNet. We also released the model for your reference.

ModelDICNetSPARNet (in paper)SPARNet (Released)SPARNet-Light-Attn3D (Released)
#Params(M)22.89.8610.525.24
PSNR (↑)26.7326.9727.4327.39
SSIM (↑)0.79550.80260.82010.8189

All models are trained with CelebA and tested on Helen test set provided by DICNet

example result

SPARNetHD

We also provide network with 2D and 3D attention for SPARNetHD. For the test dataset, we clean up non-face images, add some extra test images from internet, and obtain a new CelebA-TestN dataset with 1117 images. We test the retrained model on the new dataset and recalculate the FID scores.

Similar as StyleGAN, we use the exponential moving average weight as the final model, which shows slightly better results.

ModelSPARNetHD (in paper)SPARNetHD-Attn2D (Released)SPARNetHD-Attn3D (Released)
FID (↓)27.1626.7228.42

Citation

@InProceedings{ChenSPARNet,
    author = {Chen, Chaofeng and Gong, Dihong and Wang, Hao and Li, Zhifeng and Wong, Kwan-Yee~K.},
    title = {Learning Spatial Attention for Face Super-Resolution},
    Journal = {IEEE Transactions on Image Processing (TIP)},
    year = {2020}
}

License

<a rel="license" href="http://creativecommons.org/licenses/by-nc-sa/4.0/"><img alt="Creative Commons License" style="border-width:0" src="https://i.creativecommons.org/l/by-nc-sa/4.0/88x31.png" /></a><br />This work is licensed under a <a rel="license" href="http://creativecommons.org/licenses/by-nc-sa/4.0/">Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International License</a>.

Acknowledgement

The codes are based on CycleGAN. The project also benefits from DICNet.