├── .gitignore ├── GPEN.ipynb ├── README.md ├── __init_paths.py ├── align_faces.py ├── demo.py ├── distributed.py ├── examples ├── ffhq-10 │ ├── 00000.png │ ├── 00001.png │ ├── 00002.png │ ├── 00003.png │ ├── 00004.png │ ├── 00005.png │ ├── 00006.png │ ├── 00007.png │ ├── 00008.png │ └── 00009.png ├── grays │ ├── 106000_gray.png │ ├── 762000_gray.png │ ├── 809000_gray.png │ ├── 813000_gray.png │ ├── 827000_gray.png │ └── 836000_gray.png ├── imgs │ └── Solvay_conference_1927.png ├── outs-bfr │ ├── Solvay_conference_1927_COMP.jpg │ ├── Solvay_conference_1927_GPEN.jpg │ ├── Solvay_conference_1927_face00.jpg │ ├── Solvay_conference_1927_face01.jpg │ ├── Solvay_conference_1927_face02.jpg │ ├── Solvay_conference_1927_face03.jpg │ ├── Solvay_conference_1927_face04.jpg │ ├── Solvay_conference_1927_face05.jpg │ ├── Solvay_conference_1927_face06.jpg │ ├── Solvay_conference_1927_face07.jpg │ ├── Solvay_conference_1927_face08.jpg │ ├── Solvay_conference_1927_face09.jpg │ ├── Solvay_conference_1927_face10.jpg │ ├── Solvay_conference_1927_face11.jpg │ ├── Solvay_conference_1927_face12.jpg │ ├── Solvay_conference_1927_face13.jpg │ ├── Solvay_conference_1927_face14.jpg │ ├── Solvay_conference_1927_face15.jpg │ ├── Solvay_conference_1927_face16.jpg │ ├── Solvay_conference_1927_face17.jpg │ ├── Solvay_conference_1927_face18.jpg │ ├── Solvay_conference_1927_face19.jpg │ ├── Solvay_conference_1927_face20.jpg │ ├── Solvay_conference_1927_face21.jpg │ ├── Solvay_conference_1927_face22.jpg │ ├── Solvay_conference_1927_face23.jpg │ ├── Solvay_conference_1927_face24.jpg │ ├── Solvay_conference_1927_face25.jpg │ ├── Solvay_conference_1927_face26.jpg │ ├── Solvay_conference_1927_face27.jpg │ └── Solvay_conference_1927_face28.jpg ├── outs-colorization │ ├── 106000_gray_COMP.jpg │ ├── 106000_gray_GPEN.jpg │ ├── 762000_gray_COMP.jpg │ ├── 762000_gray_GPEN.jpg │ ├── 809000_gray_COMP.jpg │ ├── 809000_gray_GPEN.jpg │ ├── 813000_gray_COMP.jpg │ ├── 813000_gray_GPEN.jpg │ ├── 827000_gray_COMP.jpg │ ├── 827000_gray_GPEN.jpg │ ├── 836000_gray_COMP.jpg │ └── 836000_gray_GPEN.jpg ├── outs-inpainting │ ├── 00000_COMP.jpg │ ├── 00000_GPEN.jpg │ ├── 00001_COMP.jpg │ ├── 00001_GPEN.jpg │ ├── 00002_COMP.jpg │ ├── 00002_GPEN.jpg │ ├── 00003_COMP.jpg │ ├── 00003_GPEN.jpg │ ├── 00004_COMP.jpg │ ├── 00004_GPEN.jpg │ ├── 00005_COMP.jpg │ ├── 00005_GPEN.jpg │ ├── 00006_COMP.jpg │ ├── 00006_GPEN.jpg │ ├── 00007_COMP.jpg │ ├── 00007_GPEN.jpg │ ├── 00008_COMP.jpg │ ├── 00008_GPEN.jpg │ ├── 00009_COMP.jpg │ └── 00009_GPEN.jpg ├── outs-seg2face │ ├── 28000_COMP.jpg │ ├── 28000_GPEN.jpg │ ├── 28001_COMP.jpg │ ├── 28001_GPEN.jpg │ ├── 28002_COMP.jpg │ ├── 28002_GPEN.jpg │ ├── 28003_COMP.jpg │ ├── 28003_GPEN.jpg │ ├── 28004_COMP.jpg │ ├── 28004_GPEN.jpg │ ├── 28005_COMP.jpg │ ├── 28005_GPEN.jpg │ ├── 28006_COMP.jpg │ ├── 28006_GPEN.jpg │ ├── 28007_COMP.jpg │ ├── 28007_GPEN.jpg │ ├── 28008_COMP.jpg │ ├── 28008_GPEN.jpg │ ├── 28009_COMP.jpg │ ├── 28009_GPEN.jpg │ ├── 28010_COMP.jpg │ ├── 28010_GPEN.jpg │ ├── 28011_COMP.jpg │ ├── 28011_GPEN.jpg │ ├── 28012_COMP.jpg │ ├── 28012_GPEN.jpg │ ├── 28013_COMP.jpg │ ├── 28013_GPEN.jpg │ ├── 28014_COMP.jpg │ ├── 28014_GPEN.jpg │ ├── 28015_COMP.jpg │ ├── 28015_GPEN.jpg │ ├── 28016_COMP.jpg │ ├── 28016_GPEN.jpg │ ├── 28017_COMP.jpg │ ├── 28017_GPEN.jpg │ ├── 28018_COMP.jpg │ ├── 28018_GPEN.jpg │ ├── 28019_COMP.jpg │ └── 28019_GPEN.jpg ├── outs-selfie │ ├── 2020-12-10 175750_COMP.jpg │ ├── 2020-12-10 175750_GPEN.jpg │ ├── 2020-12-10 175750_face00.jpg │ ├── 2020-12-17 185036_COMP.jpg │ ├── 2020-12-17 185036_GPEN.jpg │ ├── 2020-12-17 185036_face00.jpg │ ├── 20200320.101220_COMP.jpg │ ├── 20200320.101220_GPEN.jpg │ ├── 20200320.101220_face00.jpg │ ├── 2021-03-10 160843(29)_COMP.jpg │ ├── 2021-03-10 160843(29)_GPEN.jpg │ ├── 2021-03-10 160843(29)_face00.jpg │ ├── IMG20180716171517_COMP.jpg │ ├── IMG20180716171517_GPEN.jpg │ └── IMG20180716171517_face00.jpg ├── segs │ ├── 28000.png │ ├── 28001.png │ ├── 28002.png │ ├── 28003.png │ ├── 28004.png │ ├── 28005.png │ ├── 28006.png │ ├── 28007.png │ ├── 28008.png │ ├── 28009.png │ ├── 28010.png │ ├── 28011.png │ ├── 28012.png │ ├── 28013.png │ ├── 28014.png │ ├── 28015.png │ ├── 28016.png │ ├── 28017.png │ ├── 28018.png │ └── 28019.png └── selfie │ ├── 2020-12-10 175750.jpg │ ├── 2020-12-17 185036.jpg │ ├── 20200320.101220.jpg │ ├── 2021-03-10 160843(29).jpg │ └── IMG20180716171517.jpg ├── face_colorization.py ├── face_detect ├── .DS_Store ├── data │ ├── FDDB │ │ └── img_list.txt │ ├── __init__.py │ ├── config.py │ ├── data_augment.py │ └── wider_face.py ├── facemodels │ ├── __init__.py │ ├── net.py │ └── retinaface.py ├── layers │ ├── __init__.py │ ├── functions │ │ └── prior_box.py │ └── modules │ │ ├── __init__.py │ │ └── multibox_loss.py ├── retinaface_detection.py └── utils │ ├── __init__.py │ ├── box_utils.py │ ├── nms │ ├── __init__.py │ └── py_cpu_nms.py │ └── timer.py ├── face_enhancement.py ├── face_inpainting.py ├── face_model ├── face_gan.py ├── gpen_model.py └── op │ ├── __init__.py │ ├── fused_act.py │ ├── fused_bias_act.cpp │ ├── fused_bias_act_kernel.cu │ ├── upfirdn2d.cpp │ ├── upfirdn2d.py │ └── upfirdn2d_kernel.cu ├── face_parse ├── blocks.py ├── face_parsing.py ├── mask.png ├── parse_model.py └── test.png ├── figs ├── Solvay_conference_1927_comp.jpg ├── architecture.png ├── colorization_00.jpg ├── colorization_01.jpg ├── inpainting_00.jpg ├── inpainting_01.jpg ├── real_00.png ├── real_01.png ├── real_02.png ├── real_03.png ├── seg2face_00.jpg ├── seg2face_01.jpg ├── selfie_00.jpg └── selfie_01.jpg ├── misc ├── cog.yaml ├── onnx_export.py └── predict.py ├── requirements.txt ├── run.sh ├── segmentation2face.py ├── sr_model ├── arch_util.py ├── real_esrnet.py └── rrdbnet_arch.py ├── train_simple.py ├── training ├── data_loader │ ├── dataset_face.py │ └── degradations.py ├── loss │ ├── helpers.py │ ├── id_loss.py │ └── model_irse.py └── lpips │ ├── __init__.py │ ├── lpips.py │ ├── pretrained_networks.py │ ├── trainer.py │ └── weights │ ├── v0.0 │ ├── alex.pth │ ├── squeeze.pth │ └── vgg.pth │ └── v0.1 │ ├── alex.pth │ ├── squeeze.pth │ └── vgg.pth ├── val ├── hq │ ├── ffhq_00000.png │ ├── ffhq_00001.png │ ├── ffhq_00002.png │ └── ffhq_00003.png └── lq │ ├── ffhq_00000.png │ ├── ffhq_00001.png │ ├── ffhq_00002.png │ └── ffhq_00003.png └── weights └── README.md /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | pip-wheel-metadata/ 24 | share/python-wheels/ 25 | *.egg-info/ 26 | .installed.cfg 27 | *.egg 28 | MANIFEST 29 | 30 | # PyInstaller 31 | # Usually these files are written by a python script from a template 32 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 33 | *.manifest 34 | *.spec 35 | 36 | # Installer logs 37 | pip-log.txt 38 | pip-delete-this-directory.txt 39 | 40 | # Unit test / coverage reports 41 | htmlcov/ 42 | .tox/ 43 | .nox/ 44 | .coverage 45 | .coverage.* 46 | .cache 47 | nosetests.xml 48 | coverage.xml 49 | *.cover 50 | *.py,cover 51 | .hypothesis/ 52 | .pytest_cache/ 53 | 54 | # Translations 55 | *.mo 56 | *.pot 57 | 58 | # Django stuff: 59 | *.log 60 | local_settings.py 61 | db.sqlite3 62 | db.sqlite3-journal 63 | 64 | # Flask stuff: 65 | instance/ 66 | .webassets-cache 67 | 68 | # Scrapy stuff: 69 | .scrapy 70 | 71 | # Sphinx documentation 72 | docs/_build/ 73 | 74 | # PyBuilder 75 | target/ 76 | 77 | # Jupyter Notebook 78 | .ipynb_checkpoints 79 | 80 | # IPython 81 | profile_default/ 82 | ipython_config.py 83 | 84 | # pyenv 85 | .python-version 86 | 87 | # pipenv 88 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 89 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 90 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 91 | # install all needed dependencies. 92 | #Pipfile.lock 93 | 94 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 95 | __pypackages__/ 96 | 97 | # Celery stuff 98 | celerybeat-schedule 99 | celerybeat.pid 100 | 101 | # SageMath parsed files 102 | *.sage.py 103 | 104 | # Environments 105 | .env 106 | .venv 107 | env/ 108 | venv/ 109 | ENV/ 110 | env.bak/ 111 | venv.bak/ 112 | 113 | # Spyder project settings 114 | .spyderproject 115 | .spyproject 116 | 117 | # Rope project settings 118 | .ropeproject 119 | 120 | # mkdocs documentation 121 | /site 122 | 123 | # mypy 124 | .mypy_cache/ 125 | .dmypy.json 126 | dmypy.json 127 | 128 | # Pyre type checker 129 | .pyre/ 130 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # GAN Prior Embedded Network for Blind Face Restoration in the Wild 2 | 3 | [Paper](https://arxiv.org/abs/2105.06070) | [Supplementary](https://www4.comp.polyu.edu.hk/~cslzhang/paper/GPEN-cvpr21-supp.pdf) | [Demo](https://vision.aliyun.com/experience/detail?spm=a211p3.14020179.J_7524944390.17.66cd4850wVDkUQ&tagName=facebody&children=EnhanceFace) | [ModelScope](https://www.modelscope.cn/models/damo/cv_gpen_image-portrait-enhancement/summary) 4 | 5 | [![Hugging Face Spaces](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue)](https://huggingface.co/spaces/akhaliq/GPEN) 6 | 7 | [Tao Yang](https://cg.cs.tsinghua.edu.cn/people/~tyang)1, Peiran Ren1, Xuansong Xie1, [Lei Zhang](https://www4.comp.polyu.edu.hk/~cslzhang)1,2 8 | _1[DAMO Academy, Alibaba Group](https://damo.alibaba.com), Hangzhou, China_ 9 | _2[Department of Computing, The Hong Kong Polytechnic University](http://www.comp.polyu.edu.hk), Hong Kong, China_ 10 | 11 | #### Face Restoration 12 | 13 | 14 | 15 | 16 | 17 | 18 | #### Selfie Restoration 19 | 20 | 21 | 22 | #### Face Colorization 23 | 24 | 25 | 26 | #### Face Inpainting 27 | 28 | 29 | 30 | #### Conditional Image Synthesis (Seg2Face) 31 | 32 | 33 | 34 | ## News 35 | (2023-02-15) **GPEN-BFR-1024** and **GPEN-BFR-2048** are now publicly available. Please download them via \[[ModelScope2](https://www.modelscope.cn/models/damo/cv_gpen_image-portrait-enhancement-hires/summary)\]. 36 | 37 | (2023-02-15) We provide online demos via \[[ModelScope1](https://www.modelscope.cn/models/damo/cv_gpen_image-portrait-enhancement/summary)\] and \[[ModelScope2](https://www.modelscope.cn/models/damo/cv_gpen_image-portrait-enhancement-hires/summary)\]. 38 | 39 | (2022-05-16) Add x1 sr model. Add ``--tile_size`` to avoid OOM. 40 | 41 | (2022-03-15) Add x4 sr model. Try ``--sr_scale``. 42 | 43 | (2022-03-09) Add GPEN-BFR-2048 for selfies. I have to take it down due to commercial issues. Sorry about that. 44 | 45 | (2021-12-29) Add online demos [![Hugging Face Spaces](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue)](https://huggingface.co/spaces/akhaliq/GPEN). Many thanks to [CJWBW](https://github.com/CJWBW) and [AK391](https://github.com/AK391). 46 | 47 | (2021-12-16) Release a simplified training code of GPEN. It differs from our implementation in the paper, but could achieve comparable performance. We strongly recommend to change the degradation model. 48 | 49 | (2021-12-09) Add face parsing to better paste restored faces back. 50 | 51 | (2021-12-09) GPEN can run on CPU now by simply discarding ``--use_cuda``. 52 | 53 | (2021-12-01) GPEN can now work on a Windows machine without compiling cuda codes. Please check it out. Thanks to [Animadversio](https://github.com/rosinality/stylegan2-pytorch/issues/81). Alternatively, you can try [GPEN-Windows](https://drive.google.com/file/d/1YJJVnPGq90e_mWZxSGGTptNQilZNfOEO/view?usp=drivesdk). Many thanks to [Cioscos](https://github.com/yangxy/GPEN/issues/74). 54 | 55 | (2021-10-22) GPEN can now work with SR methods. A SR model trained by myself is provided. Replace it with your own model if necessary. 56 | 57 | (2021-10-11) The Colab demo for GPEN is available now google colab logo. 58 | 59 | ## Download models from Modelscope 60 | 61 | - Install modelscope: 62 | ```bash 63 | pip install "modelscope[cv]" -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html 64 | ``` 65 | 66 | - Run the following codes: 67 | ```python 68 | import cv2 69 | from modelscope.pipelines import pipeline 70 | from modelscope.utils.constant import Tasks 71 | from modelscope.outputs import OutputKeys 72 | 73 | portrait_enhancement = pipeline(Tasks.image_portrait_enhancement, model='damo/cv_gpen_image-portrait-enhancement-hires') 74 | result = portrait_enhancement('https://modelscope.oss-cn-beijing.aliyuncs.com/test/images/marilyn_monroe_4.jpg') 75 | cv2.imwrite('result.png', result[OutputKeys.OUTPUT_IMG]) 76 | ``` 77 | 78 | It will automatically download the GPEN models. You can find the model in the local path ``~/.cache/modelscope/hub/damo``. Please note pytorch_model.pt, pytorch_model-2048.pt are respectively the 1024 and 2048 versions. 79 | 80 | ## Usage 81 | 82 | ![python](https://img.shields.io/badge/python-v3.7.4-green.svg?style=plastic) 83 | ![pytorch](https://img.shields.io/badge/pytorch-v1.7.0-green.svg?style=plastic) 84 | ![cuda](https://img.shields.io/badge/cuda-v10.2.89-green.svg?style=plastic) 85 | ![driver](https://img.shields.io/badge/driver-v460.73.01-green.svg?style=plastic) 86 | ![gcc](https://img.shields.io/badge/gcc-v7.5.0-green.svg?style=plastic) 87 | 88 | - Clone this repository: 89 | ```bash 90 | git clone https://github.com/yangxy/GPEN.git 91 | cd GPEN 92 | ``` 93 | - Download RetinaFace model and our pre-trained model (not our best model due to commercial issues) and put them into ``weights/``. 94 | 95 | [RetinaFace-R50](https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/RetinaFace-R50.pth) | [ParseNet-latest](https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/ParseNet-latest.pth) | [model_ir_se50](https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/model_ir_se50.pth) | [GPEN-BFR-512](https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/GPEN-BFR-512.pth) | [GPEN-BFR-512-D](https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/GPEN-BFR-512-D.pth) | [GPEN-BFR-256](https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/GPEN-BFR-256.pth) | [GPEN-BFR-256-D](https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/GPEN-BFR-256-D.pth) | [GPEN-Colorization-1024](https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/GPEN-Colorization-1024.pth) | [GPEN-Inpainting-1024](https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/GPEN-Inpainting-1024.pth) | [GPEN-Seg2face-512](https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/GPEN-Seg2face-512.pth) | [realesrnet_x1](https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/realesrnet_x1.pth) | [realesrnet_x2](https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/realesrnet_x2.pth) | [realesrnet_x4](https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/realesrnet_x4.pth) 96 | 97 | - Restore face images: 98 | ```bash 99 | python demo.py --task FaceEnhancement --model GPEN-BFR-512 --in_size 512 --channel_multiplier 2 --narrow 1 --use_sr --sr_scale 4 --use_cuda --save_face --indir examples/imgs --outdir examples/outs-bfr 100 | ``` 101 | 102 | - Colorize faces: 103 | ```bash 104 | python demo.py --task FaceColorization --model GPEN-Colorization-1024 --in_size 1024 --use_cuda --indir examples/grays --outdir examples/outs-colorization 105 | ``` 106 | 107 | - Complete faces: 108 | ```bash 109 | python demo.py --task FaceInpainting --model GPEN-Inpainting-1024 --in_size 1024 --use_cuda --indir examples/ffhq-10 --outdir examples/outs-inpainting 110 | ``` 111 | 112 | - Synthesize faces: 113 | ```bash 114 | python demo.py --task Segmentation2Face --model GPEN-Seg2face-512 --in_size 512 --use_cuda --indir examples/segs --outdir examples/outs-seg2face 115 | ``` 116 | 117 | - Train GPEN for BFR with 4 GPUs: 118 | ```bash 119 | CUDA_VISIBLE_DEVICES='0,1,2,3' python -m torch.distributed.launch --nproc_per_node=4 --master_port=4321 train_simple.py --size 1024 --channel_multiplier 2 --narrow 1 --ckpt weights --sample results --batch 2 --path your_path_of_croped+aligned_hq_faces (e.g., FFHQ) 120 | 121 | ``` 122 | When testing your own model, set ``--key g_ema``. 123 | 124 | Please check out ``run.sh`` for more details. 125 | 126 | ## Main idea 127 | 128 | 129 | ## Citation 130 | If our work is useful for your research, please consider citing: 131 | 132 | @inproceedings{Yang2021GPEN, 133 | title={GAN Prior Embedded Network for Blind Face Restoration in the Wild}, 134 | author={Tao Yang, Peiran Ren, Xuansong Xie, and Lei Zhang}, 135 | booktitle={IEEE Conference on Computer Vision and Pattern Recognition (CVPR)}, 136 | year={2021} 137 | } 138 | 139 | ## License 140 | © Alibaba, 2021. For academic and non-commercial use only. 141 | 142 | ## Acknowledgments 143 | We borrow some codes from [Pytorch_Retinaface](https://github.com/biubug6/Pytorch_Retinaface), [stylegan2-pytorch](https://github.com/rosinality/stylegan2-pytorch), [Real-ESRGAN](https://github.com/xinntao/Real-ESRGAN), and [GFPGAN](https://github.com/TencentARC/GFPGAN). 144 | 145 | ## Contact 146 | If you have any questions or suggestions about this paper, feel free to reach me at yangtao9009@gmail.com. 147 | -------------------------------------------------------------------------------- /__init_paths.py: -------------------------------------------------------------------------------- 1 | ''' 2 | @paper: GAN Prior Embedded Network for Blind Face Restoration in the Wild (CVPR2021) 3 | @author: yangxy (yangtao9009@gmail.com) 4 | ''' 5 | import os.path as osp 6 | import sys 7 | 8 | def add_path(path): 9 | if path not in sys.path: 10 | sys.path.insert(0, path) 11 | 12 | this_dir = osp.dirname(__file__) 13 | 14 | path = osp.join(this_dir, 'face_detect') 15 | add_path(path) 16 | 17 | path = osp.join(this_dir, 'face_parse') 18 | add_path(path) 19 | 20 | path = osp.join(this_dir, 'face_model') 21 | add_path(path) 22 | 23 | path = osp.join(this_dir, 'sr_model') 24 | add_path(path) 25 | 26 | path = osp.join(this_dir, 'training') 27 | add_path(path) 28 | 29 | path = osp.join(this_dir, 'training/loss') 30 | add_path(path) 31 | 32 | path = osp.join(this_dir, 'training/data_loader') 33 | add_path(path) 34 | -------------------------------------------------------------------------------- /demo.py: -------------------------------------------------------------------------------- 1 | ''' 2 | @paper: GAN Prior Embedded Network for Blind Face Restoration in the Wild (CVPR2021) 3 | @author: yangxy (yangtao9009@gmail.com) 4 | ''' 5 | import os 6 | import cv2 7 | import glob 8 | import time 9 | import math 10 | import argparse 11 | import numpy as np 12 | from PIL import Image, ImageDraw 13 | import __init_paths 14 | from face_enhancement import FaceEnhancement 15 | from face_colorization import FaceColorization 16 | from face_inpainting import FaceInpainting 17 | from segmentation2face import Segmentation2Face 18 | 19 | def brush_stroke_mask(img, color=(255,255,255)): 20 | min_num_vertex = 8 21 | max_num_vertex = 28 22 | mean_angle = 2*math.pi / 5 23 | angle_range = 2*math.pi / 15 24 | min_width = 12 25 | max_width = 80 26 | def generate_mask(H, W, img=None): 27 | average_radius = math.sqrt(H*H+W*W) / 8 28 | mask = Image.new('RGB', (W, H), 0) 29 | if img is not None: mask = img #Image.fromarray(img) 30 | 31 | for _ in range(np.random.randint(1, 4)): 32 | num_vertex = np.random.randint(min_num_vertex, max_num_vertex) 33 | angle_min = mean_angle - np.random.uniform(0, angle_range) 34 | angle_max = mean_angle + np.random.uniform(0, angle_range) 35 | angles = [] 36 | vertex = [] 37 | for i in range(num_vertex): 38 | if i % 2 == 0: 39 | angles.append(2*math.pi - np.random.uniform(angle_min, angle_max)) 40 | else: 41 | angles.append(np.random.uniform(angle_min, angle_max)) 42 | 43 | h, w = mask.size 44 | vertex.append((int(np.random.randint(0, w)), int(np.random.randint(0, h)))) 45 | for i in range(num_vertex): 46 | r = np.clip( 47 | np.random.normal(loc=average_radius, scale=average_radius//2), 48 | 0, 2*average_radius) 49 | new_x = np.clip(vertex[-1][0] + r * math.cos(angles[i]), 0, w) 50 | new_y = np.clip(vertex[-1][1] + r * math.sin(angles[i]), 0, h) 51 | vertex.append((int(new_x), int(new_y))) 52 | 53 | draw = ImageDraw.Draw(mask) 54 | width = int(np.random.uniform(min_width, max_width)) 55 | draw.line(vertex, fill=color, width=width) 56 | for v in vertex: 57 | draw.ellipse((v[0] - width//2, 58 | v[1] - width//2, 59 | v[0] + width//2, 60 | v[1] + width//2), 61 | fill=color) 62 | 63 | return mask 64 | 65 | width, height = img.size 66 | mask = generate_mask(height, width, img) 67 | return mask 68 | 69 | if __name__=='__main__': 70 | parser = argparse.ArgumentParser() 71 | parser.add_argument('--model', type=str, default='GPEN-BFR-512', help='GPEN model') 72 | parser.add_argument('--task', type=str, default='FaceEnhancement', help='task of GPEN model') 73 | parser.add_argument('--key', type=str, default=None, help='key of GPEN model') 74 | parser.add_argument('--in_size', type=int, default=512, help='in resolution of GPEN') 75 | parser.add_argument('--out_size', type=int, default=None, help='out resolution of GPEN') 76 | parser.add_argument('--channel_multiplier', type=int, default=2, help='channel multiplier of GPEN') 77 | parser.add_argument('--narrow', type=float, default=1, help='channel narrow scale') 78 | parser.add_argument('--alpha', type=float, default=1, help='blending the results') 79 | parser.add_argument('--use_sr', action='store_true', help='use sr or not') 80 | parser.add_argument('--use_cuda', action='store_true', help='use cuda or not') 81 | parser.add_argument('--save_face', action='store_true', help='save face or not') 82 | parser.add_argument('--aligned', action='store_true', help='input are aligned faces or not') 83 | parser.add_argument('--sr_model', type=str, default='realesrnet', help='SR model') 84 | parser.add_argument('--sr_scale', type=int, default=2, help='SR scale') 85 | parser.add_argument('--tile_size', type=int, default=0, help='tile size for SR to avoid OOM') 86 | parser.add_argument('--indir', type=str, default='examples/imgs', help='input folder') 87 | parser.add_argument('--outdir', type=str, default='results/outs-BFR', help='output folder') 88 | parser.add_argument('--ext', type=str, default='.jpg', help='extension of output') 89 | args = parser.parse_args() 90 | 91 | #model = {'name':'GPEN-BFR-512', 'size':512, 'channel_multiplier':2, 'narrow':1} 92 | #model = {'name':'GPEN-BFR-256', 'size':256, 'channel_multiplier':1, 'narrow':0.5} 93 | 94 | os.makedirs(args.outdir, exist_ok=True) 95 | 96 | if args.task == 'FaceEnhancement': 97 | processer = FaceEnhancement(args, in_size=args.in_size, model=args.model, use_sr=args.use_sr, device='cuda' if args.use_cuda else 'cpu') 98 | elif args.task == 'FaceColorization': 99 | processer = FaceColorization(in_size=args.in_size, model=args.model, device='cuda' if args.use_cuda else 'cpu') 100 | elif args.task == 'FaceInpainting': 101 | processer = FaceInpainting(in_size=args.in_size, model=args.model, device='cuda' if args.use_cuda else 'cpu') 102 | elif args.task == 'Segmentation2Face': 103 | processer = Segmentation2Face(in_size=args.in_size, model=args.model, is_norm=False, device='cuda' if args.use_cuda else 'cpu') 104 | 105 | 106 | files = sorted(glob.glob(os.path.join(args.indir, '*.*g'))) 107 | for n, file in enumerate(files[:]): 108 | filename, ext = os.path.splitext(os.path.basename(file)) 109 | 110 | img = cv2.imread(file, cv2.IMREAD_COLOR) # BGR 111 | if not isinstance(img, np.ndarray): print(filename, 'error'); continue 112 | #img = cv2.resize(img, (0,0), fx=2, fy=2) # optional 113 | 114 | if args.task == 'FaceInpainting': 115 | img = np.asarray(brush_stroke_mask(Image.fromarray(img))) 116 | 117 | img_out, orig_faces, enhanced_faces = processer.process(img, aligned=args.aligned) 118 | 119 | img = cv2.resize(img, img_out.shape[:2][::-1]) 120 | cv2.imwrite(f'{args.outdir}/{filename}_COMP{args.ext}', np.hstack((img, img_out))) 121 | cv2.imwrite(f'{args.outdir}/{filename}_GPEN{args.ext}', img_out) 122 | 123 | if args.save_face: 124 | for m, (ef, of) in enumerate(zip(enhanced_faces, orig_faces)): 125 | of = cv2.resize(of, ef.shape[:2]) 126 | cv2.imwrite(f'{args.outdir}/{filename}_face{m:02d}{args.ext}', np.hstack((of, ef))) 127 | 128 | if n%10==0: print(n, filename) 129 | -------------------------------------------------------------------------------- /distributed.py: -------------------------------------------------------------------------------- 1 | import math 2 | import pickle 3 | 4 | import torch 5 | from torch import distributed as dist 6 | from torch.utils.data.sampler import Sampler 7 | 8 | 9 | def get_rank(): 10 | if not dist.is_available(): 11 | return 0 12 | 13 | if not dist.is_initialized(): 14 | return 0 15 | 16 | return dist.get_rank() 17 | 18 | 19 | def synchronize(): 20 | if not dist.is_available(): 21 | return 22 | 23 | if not dist.is_initialized(): 24 | return 25 | 26 | world_size = dist.get_world_size() 27 | 28 | if world_size == 1: 29 | return 30 | 31 | dist.barrier() 32 | 33 | 34 | def get_world_size(): 35 | if not dist.is_available(): 36 | return 1 37 | 38 | if not dist.is_initialized(): 39 | return 1 40 | 41 | return dist.get_world_size() 42 | 43 | 44 | def reduce_sum(tensor): 45 | if not dist.is_available(): 46 | return tensor 47 | 48 | if not dist.is_initialized(): 49 | return tensor 50 | 51 | tensor = tensor.clone() 52 | dist.all_reduce(tensor, op=dist.ReduceOp.SUM) 53 | 54 | return tensor 55 | 56 | 57 | def gather_grad(params): 58 | world_size = get_world_size() 59 | 60 | if world_size == 1: 61 | return 62 | 63 | for param in params: 64 | if param.grad is not None: 65 | dist.all_reduce(param.grad.data, op=dist.ReduceOp.SUM) 66 | param.grad.data.div_(world_size) 67 | 68 | 69 | def all_gather(data): 70 | world_size = get_world_size() 71 | 72 | if world_size == 1: 73 | return [data] 74 | 75 | buffer = pickle.dumps(data) 76 | storage = torch.ByteStorage.from_buffer(buffer) 77 | tensor = torch.ByteTensor(storage).to('cuda') 78 | 79 | local_size = torch.IntTensor([tensor.numel()]).to('cuda') 80 | size_list = [torch.IntTensor([0]).to('cuda') for _ in range(world_size)] 81 | dist.all_gather(size_list, local_size) 82 | size_list = [int(size.item()) for size in size_list] 83 | max_size = max(size_list) 84 | 85 | tensor_list = [] 86 | for _ in size_list: 87 | tensor_list.append(torch.ByteTensor(size=(max_size,)).to('cuda')) 88 | 89 | if local_size != max_size: 90 | padding = torch.ByteTensor(size=(max_size - local_size,)).to('cuda') 91 | tensor = torch.cat((tensor, padding), 0) 92 | 93 | dist.all_gather(tensor_list, tensor) 94 | 95 | data_list = [] 96 | 97 | for size, tensor in zip(size_list, tensor_list): 98 | buffer = tensor.cpu().numpy().tobytes()[:size] 99 | data_list.append(pickle.loads(buffer)) 100 | 101 | return data_list 102 | 103 | 104 | def reduce_loss_dict(loss_dict): 105 | world_size = get_world_size() 106 | 107 | if world_size < 2: 108 | return loss_dict 109 | 110 | with torch.no_grad(): 111 | keys = [] 112 | losses = [] 113 | 114 | for k in sorted(loss_dict.keys()): 115 | keys.append(k) 116 | losses.append(loss_dict[k]) 117 | 118 | losses = torch.stack(losses, 0) 119 | dist.reduce(losses, dst=0) 120 | 121 | if dist.get_rank() == 0: 122 | losses /= world_size 123 | 124 | reduced_losses = {k: v for k, v in zip(keys, losses)} 125 | 126 | return reduced_losses 127 | -------------------------------------------------------------------------------- /examples/ffhq-10/00000.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/ffhq-10/00000.png -------------------------------------------------------------------------------- /examples/ffhq-10/00001.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/ffhq-10/00001.png -------------------------------------------------------------------------------- /examples/ffhq-10/00002.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/ffhq-10/00002.png -------------------------------------------------------------------------------- /examples/ffhq-10/00003.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/ffhq-10/00003.png -------------------------------------------------------------------------------- /examples/ffhq-10/00004.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/ffhq-10/00004.png -------------------------------------------------------------------------------- /examples/ffhq-10/00005.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/ffhq-10/00005.png -------------------------------------------------------------------------------- /examples/ffhq-10/00006.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/ffhq-10/00006.png -------------------------------------------------------------------------------- /examples/ffhq-10/00007.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/ffhq-10/00007.png -------------------------------------------------------------------------------- /examples/ffhq-10/00008.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/ffhq-10/00008.png -------------------------------------------------------------------------------- /examples/ffhq-10/00009.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/ffhq-10/00009.png -------------------------------------------------------------------------------- /examples/grays/106000_gray.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/grays/106000_gray.png -------------------------------------------------------------------------------- /examples/grays/762000_gray.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/grays/762000_gray.png -------------------------------------------------------------------------------- /examples/grays/809000_gray.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/grays/809000_gray.png -------------------------------------------------------------------------------- /examples/grays/813000_gray.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/grays/813000_gray.png -------------------------------------------------------------------------------- /examples/grays/827000_gray.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/grays/827000_gray.png -------------------------------------------------------------------------------- /examples/grays/836000_gray.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/grays/836000_gray.png -------------------------------------------------------------------------------- /examples/imgs/Solvay_conference_1927.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/imgs/Solvay_conference_1927.png -------------------------------------------------------------------------------- /examples/outs-bfr/Solvay_conference_1927_COMP.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-bfr/Solvay_conference_1927_COMP.jpg -------------------------------------------------------------------------------- /examples/outs-bfr/Solvay_conference_1927_GPEN.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-bfr/Solvay_conference_1927_GPEN.jpg -------------------------------------------------------------------------------- /examples/outs-bfr/Solvay_conference_1927_face00.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-bfr/Solvay_conference_1927_face00.jpg -------------------------------------------------------------------------------- /examples/outs-bfr/Solvay_conference_1927_face01.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-bfr/Solvay_conference_1927_face01.jpg -------------------------------------------------------------------------------- /examples/outs-bfr/Solvay_conference_1927_face02.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-bfr/Solvay_conference_1927_face02.jpg -------------------------------------------------------------------------------- /examples/outs-bfr/Solvay_conference_1927_face03.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-bfr/Solvay_conference_1927_face03.jpg -------------------------------------------------------------------------------- /examples/outs-bfr/Solvay_conference_1927_face04.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-bfr/Solvay_conference_1927_face04.jpg -------------------------------------------------------------------------------- /examples/outs-bfr/Solvay_conference_1927_face05.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-bfr/Solvay_conference_1927_face05.jpg -------------------------------------------------------------------------------- /examples/outs-bfr/Solvay_conference_1927_face06.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-bfr/Solvay_conference_1927_face06.jpg -------------------------------------------------------------------------------- /examples/outs-bfr/Solvay_conference_1927_face07.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-bfr/Solvay_conference_1927_face07.jpg -------------------------------------------------------------------------------- /examples/outs-bfr/Solvay_conference_1927_face08.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-bfr/Solvay_conference_1927_face08.jpg -------------------------------------------------------------------------------- /examples/outs-bfr/Solvay_conference_1927_face09.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-bfr/Solvay_conference_1927_face09.jpg -------------------------------------------------------------------------------- /examples/outs-bfr/Solvay_conference_1927_face10.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-bfr/Solvay_conference_1927_face10.jpg -------------------------------------------------------------------------------- /examples/outs-bfr/Solvay_conference_1927_face11.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-bfr/Solvay_conference_1927_face11.jpg -------------------------------------------------------------------------------- /examples/outs-bfr/Solvay_conference_1927_face12.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-bfr/Solvay_conference_1927_face12.jpg -------------------------------------------------------------------------------- /examples/outs-bfr/Solvay_conference_1927_face13.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-bfr/Solvay_conference_1927_face13.jpg -------------------------------------------------------------------------------- /examples/outs-bfr/Solvay_conference_1927_face14.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-bfr/Solvay_conference_1927_face14.jpg -------------------------------------------------------------------------------- /examples/outs-bfr/Solvay_conference_1927_face15.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-bfr/Solvay_conference_1927_face15.jpg -------------------------------------------------------------------------------- /examples/outs-bfr/Solvay_conference_1927_face16.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-bfr/Solvay_conference_1927_face16.jpg -------------------------------------------------------------------------------- /examples/outs-bfr/Solvay_conference_1927_face17.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-bfr/Solvay_conference_1927_face17.jpg -------------------------------------------------------------------------------- /examples/outs-bfr/Solvay_conference_1927_face18.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-bfr/Solvay_conference_1927_face18.jpg -------------------------------------------------------------------------------- /examples/outs-bfr/Solvay_conference_1927_face19.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-bfr/Solvay_conference_1927_face19.jpg -------------------------------------------------------------------------------- /examples/outs-bfr/Solvay_conference_1927_face20.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-bfr/Solvay_conference_1927_face20.jpg -------------------------------------------------------------------------------- /examples/outs-bfr/Solvay_conference_1927_face21.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-bfr/Solvay_conference_1927_face21.jpg -------------------------------------------------------------------------------- /examples/outs-bfr/Solvay_conference_1927_face22.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-bfr/Solvay_conference_1927_face22.jpg -------------------------------------------------------------------------------- /examples/outs-bfr/Solvay_conference_1927_face23.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-bfr/Solvay_conference_1927_face23.jpg -------------------------------------------------------------------------------- /examples/outs-bfr/Solvay_conference_1927_face24.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-bfr/Solvay_conference_1927_face24.jpg -------------------------------------------------------------------------------- /examples/outs-bfr/Solvay_conference_1927_face25.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-bfr/Solvay_conference_1927_face25.jpg -------------------------------------------------------------------------------- /examples/outs-bfr/Solvay_conference_1927_face26.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-bfr/Solvay_conference_1927_face26.jpg -------------------------------------------------------------------------------- /examples/outs-bfr/Solvay_conference_1927_face27.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-bfr/Solvay_conference_1927_face27.jpg -------------------------------------------------------------------------------- /examples/outs-bfr/Solvay_conference_1927_face28.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-bfr/Solvay_conference_1927_face28.jpg -------------------------------------------------------------------------------- /examples/outs-colorization/106000_gray_COMP.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-colorization/106000_gray_COMP.jpg -------------------------------------------------------------------------------- /examples/outs-colorization/106000_gray_GPEN.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-colorization/106000_gray_GPEN.jpg -------------------------------------------------------------------------------- /examples/outs-colorization/762000_gray_COMP.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-colorization/762000_gray_COMP.jpg -------------------------------------------------------------------------------- /examples/outs-colorization/762000_gray_GPEN.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-colorization/762000_gray_GPEN.jpg -------------------------------------------------------------------------------- /examples/outs-colorization/809000_gray_COMP.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-colorization/809000_gray_COMP.jpg -------------------------------------------------------------------------------- /examples/outs-colorization/809000_gray_GPEN.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-colorization/809000_gray_GPEN.jpg -------------------------------------------------------------------------------- /examples/outs-colorization/813000_gray_COMP.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-colorization/813000_gray_COMP.jpg -------------------------------------------------------------------------------- /examples/outs-colorization/813000_gray_GPEN.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-colorization/813000_gray_GPEN.jpg -------------------------------------------------------------------------------- /examples/outs-colorization/827000_gray_COMP.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-colorization/827000_gray_COMP.jpg -------------------------------------------------------------------------------- /examples/outs-colorization/827000_gray_GPEN.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-colorization/827000_gray_GPEN.jpg -------------------------------------------------------------------------------- /examples/outs-colorization/836000_gray_COMP.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-colorization/836000_gray_COMP.jpg -------------------------------------------------------------------------------- /examples/outs-colorization/836000_gray_GPEN.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-colorization/836000_gray_GPEN.jpg -------------------------------------------------------------------------------- /examples/outs-inpainting/00000_COMP.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-inpainting/00000_COMP.jpg -------------------------------------------------------------------------------- /examples/outs-inpainting/00000_GPEN.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-inpainting/00000_GPEN.jpg -------------------------------------------------------------------------------- /examples/outs-inpainting/00001_COMP.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-inpainting/00001_COMP.jpg -------------------------------------------------------------------------------- /examples/outs-inpainting/00001_GPEN.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-inpainting/00001_GPEN.jpg -------------------------------------------------------------------------------- /examples/outs-inpainting/00002_COMP.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-inpainting/00002_COMP.jpg -------------------------------------------------------------------------------- /examples/outs-inpainting/00002_GPEN.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-inpainting/00002_GPEN.jpg -------------------------------------------------------------------------------- /examples/outs-inpainting/00003_COMP.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-inpainting/00003_COMP.jpg -------------------------------------------------------------------------------- /examples/outs-inpainting/00003_GPEN.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-inpainting/00003_GPEN.jpg -------------------------------------------------------------------------------- /examples/outs-inpainting/00004_COMP.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-inpainting/00004_COMP.jpg -------------------------------------------------------------------------------- /examples/outs-inpainting/00004_GPEN.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-inpainting/00004_GPEN.jpg -------------------------------------------------------------------------------- /examples/outs-inpainting/00005_COMP.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-inpainting/00005_COMP.jpg -------------------------------------------------------------------------------- /examples/outs-inpainting/00005_GPEN.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-inpainting/00005_GPEN.jpg -------------------------------------------------------------------------------- /examples/outs-inpainting/00006_COMP.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-inpainting/00006_COMP.jpg -------------------------------------------------------------------------------- /examples/outs-inpainting/00006_GPEN.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-inpainting/00006_GPEN.jpg -------------------------------------------------------------------------------- /examples/outs-inpainting/00007_COMP.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-inpainting/00007_COMP.jpg -------------------------------------------------------------------------------- /examples/outs-inpainting/00007_GPEN.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-inpainting/00007_GPEN.jpg -------------------------------------------------------------------------------- /examples/outs-inpainting/00008_COMP.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-inpainting/00008_COMP.jpg -------------------------------------------------------------------------------- /examples/outs-inpainting/00008_GPEN.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-inpainting/00008_GPEN.jpg -------------------------------------------------------------------------------- /examples/outs-inpainting/00009_COMP.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-inpainting/00009_COMP.jpg -------------------------------------------------------------------------------- /examples/outs-inpainting/00009_GPEN.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-inpainting/00009_GPEN.jpg -------------------------------------------------------------------------------- /examples/outs-seg2face/28000_COMP.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-seg2face/28000_COMP.jpg -------------------------------------------------------------------------------- /examples/outs-seg2face/28000_GPEN.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-seg2face/28000_GPEN.jpg -------------------------------------------------------------------------------- /examples/outs-seg2face/28001_COMP.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-seg2face/28001_COMP.jpg -------------------------------------------------------------------------------- /examples/outs-seg2face/28001_GPEN.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-seg2face/28001_GPEN.jpg -------------------------------------------------------------------------------- /examples/outs-seg2face/28002_COMP.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-seg2face/28002_COMP.jpg -------------------------------------------------------------------------------- /examples/outs-seg2face/28002_GPEN.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-seg2face/28002_GPEN.jpg -------------------------------------------------------------------------------- /examples/outs-seg2face/28003_COMP.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-seg2face/28003_COMP.jpg -------------------------------------------------------------------------------- /examples/outs-seg2face/28003_GPEN.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-seg2face/28003_GPEN.jpg -------------------------------------------------------------------------------- /examples/outs-seg2face/28004_COMP.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-seg2face/28004_COMP.jpg -------------------------------------------------------------------------------- /examples/outs-seg2face/28004_GPEN.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-seg2face/28004_GPEN.jpg -------------------------------------------------------------------------------- /examples/outs-seg2face/28005_COMP.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-seg2face/28005_COMP.jpg -------------------------------------------------------------------------------- /examples/outs-seg2face/28005_GPEN.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-seg2face/28005_GPEN.jpg -------------------------------------------------------------------------------- /examples/outs-seg2face/28006_COMP.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-seg2face/28006_COMP.jpg -------------------------------------------------------------------------------- /examples/outs-seg2face/28006_GPEN.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-seg2face/28006_GPEN.jpg -------------------------------------------------------------------------------- /examples/outs-seg2face/28007_COMP.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-seg2face/28007_COMP.jpg -------------------------------------------------------------------------------- /examples/outs-seg2face/28007_GPEN.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-seg2face/28007_GPEN.jpg -------------------------------------------------------------------------------- /examples/outs-seg2face/28008_COMP.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-seg2face/28008_COMP.jpg -------------------------------------------------------------------------------- /examples/outs-seg2face/28008_GPEN.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-seg2face/28008_GPEN.jpg -------------------------------------------------------------------------------- /examples/outs-seg2face/28009_COMP.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-seg2face/28009_COMP.jpg -------------------------------------------------------------------------------- /examples/outs-seg2face/28009_GPEN.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-seg2face/28009_GPEN.jpg -------------------------------------------------------------------------------- /examples/outs-seg2face/28010_COMP.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-seg2face/28010_COMP.jpg -------------------------------------------------------------------------------- /examples/outs-seg2face/28010_GPEN.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-seg2face/28010_GPEN.jpg -------------------------------------------------------------------------------- /examples/outs-seg2face/28011_COMP.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-seg2face/28011_COMP.jpg -------------------------------------------------------------------------------- /examples/outs-seg2face/28011_GPEN.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-seg2face/28011_GPEN.jpg -------------------------------------------------------------------------------- /examples/outs-seg2face/28012_COMP.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-seg2face/28012_COMP.jpg -------------------------------------------------------------------------------- /examples/outs-seg2face/28012_GPEN.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-seg2face/28012_GPEN.jpg -------------------------------------------------------------------------------- /examples/outs-seg2face/28013_COMP.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-seg2face/28013_COMP.jpg -------------------------------------------------------------------------------- /examples/outs-seg2face/28013_GPEN.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-seg2face/28013_GPEN.jpg -------------------------------------------------------------------------------- /examples/outs-seg2face/28014_COMP.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-seg2face/28014_COMP.jpg -------------------------------------------------------------------------------- /examples/outs-seg2face/28014_GPEN.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-seg2face/28014_GPEN.jpg -------------------------------------------------------------------------------- /examples/outs-seg2face/28015_COMP.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-seg2face/28015_COMP.jpg -------------------------------------------------------------------------------- /examples/outs-seg2face/28015_GPEN.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-seg2face/28015_GPEN.jpg -------------------------------------------------------------------------------- /examples/outs-seg2face/28016_COMP.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-seg2face/28016_COMP.jpg -------------------------------------------------------------------------------- /examples/outs-seg2face/28016_GPEN.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-seg2face/28016_GPEN.jpg -------------------------------------------------------------------------------- /examples/outs-seg2face/28017_COMP.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-seg2face/28017_COMP.jpg -------------------------------------------------------------------------------- /examples/outs-seg2face/28017_GPEN.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-seg2face/28017_GPEN.jpg -------------------------------------------------------------------------------- /examples/outs-seg2face/28018_COMP.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-seg2face/28018_COMP.jpg -------------------------------------------------------------------------------- /examples/outs-seg2face/28018_GPEN.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-seg2face/28018_GPEN.jpg -------------------------------------------------------------------------------- /examples/outs-seg2face/28019_COMP.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-seg2face/28019_COMP.jpg -------------------------------------------------------------------------------- /examples/outs-seg2face/28019_GPEN.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-seg2face/28019_GPEN.jpg -------------------------------------------------------------------------------- /examples/outs-selfie/2020-12-10 175750_COMP.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-selfie/2020-12-10 175750_COMP.jpg -------------------------------------------------------------------------------- /examples/outs-selfie/2020-12-10 175750_GPEN.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-selfie/2020-12-10 175750_GPEN.jpg -------------------------------------------------------------------------------- /examples/outs-selfie/2020-12-10 175750_face00.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-selfie/2020-12-10 175750_face00.jpg -------------------------------------------------------------------------------- /examples/outs-selfie/2020-12-17 185036_COMP.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-selfie/2020-12-17 185036_COMP.jpg -------------------------------------------------------------------------------- /examples/outs-selfie/2020-12-17 185036_GPEN.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-selfie/2020-12-17 185036_GPEN.jpg -------------------------------------------------------------------------------- /examples/outs-selfie/2020-12-17 185036_face00.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-selfie/2020-12-17 185036_face00.jpg -------------------------------------------------------------------------------- /examples/outs-selfie/20200320.101220_COMP.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-selfie/20200320.101220_COMP.jpg -------------------------------------------------------------------------------- /examples/outs-selfie/20200320.101220_GPEN.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-selfie/20200320.101220_GPEN.jpg -------------------------------------------------------------------------------- /examples/outs-selfie/20200320.101220_face00.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-selfie/20200320.101220_face00.jpg -------------------------------------------------------------------------------- /examples/outs-selfie/2021-03-10 160843(29)_COMP.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-selfie/2021-03-10 160843(29)_COMP.jpg -------------------------------------------------------------------------------- /examples/outs-selfie/2021-03-10 160843(29)_GPEN.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-selfie/2021-03-10 160843(29)_GPEN.jpg -------------------------------------------------------------------------------- /examples/outs-selfie/2021-03-10 160843(29)_face00.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-selfie/2021-03-10 160843(29)_face00.jpg -------------------------------------------------------------------------------- /examples/outs-selfie/IMG20180716171517_COMP.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-selfie/IMG20180716171517_COMP.jpg -------------------------------------------------------------------------------- /examples/outs-selfie/IMG20180716171517_GPEN.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-selfie/IMG20180716171517_GPEN.jpg -------------------------------------------------------------------------------- /examples/outs-selfie/IMG20180716171517_face00.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/outs-selfie/IMG20180716171517_face00.jpg -------------------------------------------------------------------------------- /examples/segs/28000.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/segs/28000.png -------------------------------------------------------------------------------- /examples/segs/28001.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/segs/28001.png -------------------------------------------------------------------------------- /examples/segs/28002.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/segs/28002.png -------------------------------------------------------------------------------- /examples/segs/28003.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/segs/28003.png -------------------------------------------------------------------------------- /examples/segs/28004.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/segs/28004.png -------------------------------------------------------------------------------- /examples/segs/28005.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/segs/28005.png -------------------------------------------------------------------------------- /examples/segs/28006.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/segs/28006.png -------------------------------------------------------------------------------- /examples/segs/28007.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/segs/28007.png -------------------------------------------------------------------------------- /examples/segs/28008.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/segs/28008.png -------------------------------------------------------------------------------- /examples/segs/28009.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/segs/28009.png -------------------------------------------------------------------------------- /examples/segs/28010.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/segs/28010.png -------------------------------------------------------------------------------- /examples/segs/28011.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/segs/28011.png -------------------------------------------------------------------------------- /examples/segs/28012.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/segs/28012.png -------------------------------------------------------------------------------- /examples/segs/28013.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/segs/28013.png -------------------------------------------------------------------------------- /examples/segs/28014.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/segs/28014.png -------------------------------------------------------------------------------- /examples/segs/28015.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/segs/28015.png -------------------------------------------------------------------------------- /examples/segs/28016.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/segs/28016.png -------------------------------------------------------------------------------- /examples/segs/28017.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/segs/28017.png -------------------------------------------------------------------------------- /examples/segs/28018.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/segs/28018.png -------------------------------------------------------------------------------- /examples/segs/28019.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/segs/28019.png -------------------------------------------------------------------------------- /examples/selfie/2020-12-10 175750.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/selfie/2020-12-10 175750.jpg -------------------------------------------------------------------------------- /examples/selfie/2020-12-17 185036.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/selfie/2020-12-17 185036.jpg -------------------------------------------------------------------------------- /examples/selfie/20200320.101220.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/selfie/20200320.101220.jpg -------------------------------------------------------------------------------- /examples/selfie/2021-03-10 160843(29).jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/selfie/2021-03-10 160843(29).jpg -------------------------------------------------------------------------------- /examples/selfie/IMG20180716171517.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/examples/selfie/IMG20180716171517.jpg -------------------------------------------------------------------------------- /face_colorization.py: -------------------------------------------------------------------------------- 1 | ''' 2 | @paper: GAN Prior Embedded Network for Blind Face Restoration in the Wild (CVPR2021) 3 | @author: yangxy (yangtao9009@gmail.com) 4 | ''' 5 | import cv2 6 | from face_model.face_gan import FaceGAN 7 | 8 | class FaceColorization(object): 9 | def __init__(self, base_dir='./', in_size=1024, out_size=1024, model=None, channel_multiplier=2, narrow=1, key=None, device='cuda'): 10 | self.facegan = FaceGAN(base_dir, in_size, out_size, model, channel_multiplier, narrow, key, device=device) 11 | 12 | def post_process(self, gray, out): 13 | out_rs = cv2.resize(out, gray.shape[:2][::-1]) 14 | gray_yuv = cv2.cvtColor(gray, cv2.COLOR_BGR2YUV) 15 | out_yuv = cv2.cvtColor(out_rs, cv2.COLOR_BGR2YUV) 16 | 17 | out_yuv[:, :, 0] = gray_yuv[:, :, 0] 18 | final = cv2.cvtColor(out_yuv, cv2.COLOR_YUV2BGR) 19 | 20 | return final 21 | 22 | # make sure the face image is well aligned. Please refer to face_enhancement.py 23 | def process(self, gray, aligned=True): 24 | # colorize the face 25 | out = self.facegan.process(gray) 26 | 27 | if gray.shape[:2] != out.shape[:2]: 28 | out = self.post_process(gray, out) 29 | 30 | return out, [gray], [out] 31 | 32 | 33 | -------------------------------------------------------------------------------- /face_detect/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/face_detect/.DS_Store -------------------------------------------------------------------------------- /face_detect/data/__init__.py: -------------------------------------------------------------------------------- 1 | from .wider_face import WiderFaceDetection, detection_collate 2 | from .data_augment import * 3 | from .config import * 4 | -------------------------------------------------------------------------------- /face_detect/data/config.py: -------------------------------------------------------------------------------- 1 | # config.py 2 | 3 | cfg_mnet = { 4 | 'name': 'mobilenet0.25', 5 | 'min_sizes': [[16, 32], [64, 128], [256, 512]], 6 | 'steps': [8, 16, 32], 7 | 'variance': [0.1, 0.2], 8 | 'clip': False, 9 | 'loc_weight': 2.0, 10 | 'gpu_train': True, 11 | 'batch_size': 32, 12 | 'ngpu': 1, 13 | 'epoch': 250, 14 | 'decay1': 190, 15 | 'decay2': 220, 16 | 'image_size': 640, 17 | 'pretrain': False, 18 | 'return_layers': {'stage1': 1, 'stage2': 2, 'stage3': 3}, 19 | 'in_channel': 32, 20 | 'out_channel': 64 21 | } 22 | 23 | cfg_re50 = { 24 | 'name': 'Resnet50', 25 | 'min_sizes': [[16, 32], [64, 128], [256, 512]], 26 | 'steps': [8, 16, 32], 27 | 'variance': [0.1, 0.2], 28 | 'clip': False, 29 | 'loc_weight': 2.0, 30 | 'gpu_train': True, 31 | 'batch_size': 24, 32 | 'ngpu': 4, 33 | 'epoch': 100, 34 | 'decay1': 70, 35 | 'decay2': 90, 36 | 'image_size': 840, 37 | 'pretrain': False, 38 | 'return_layers': {'layer2': 1, 'layer3': 2, 'layer4': 3}, 39 | 'in_channel': 256, 40 | 'out_channel': 256 41 | } 42 | 43 | -------------------------------------------------------------------------------- /face_detect/data/data_augment.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import numpy as np 3 | import random 4 | from utils.box_utils import matrix_iof 5 | 6 | 7 | def _crop(image, boxes, labels, landm, img_dim): 8 | height, width, _ = image.shape 9 | pad_image_flag = True 10 | 11 | for _ in range(250): 12 | """ 13 | if random.uniform(0, 1) <= 0.2: 14 | scale = 1.0 15 | else: 16 | scale = random.uniform(0.3, 1.0) 17 | """ 18 | PRE_SCALES = [0.3, 0.45, 0.6, 0.8, 1.0] 19 | scale = random.choice(PRE_SCALES) 20 | short_side = min(width, height) 21 | w = int(scale * short_side) 22 | h = w 23 | 24 | if width == w: 25 | l = 0 26 | else: 27 | l = random.randrange(width - w) 28 | if height == h: 29 | t = 0 30 | else: 31 | t = random.randrange(height - h) 32 | roi = np.array((l, t, l + w, t + h)) 33 | 34 | value = matrix_iof(boxes, roi[np.newaxis]) 35 | flag = (value >= 1) 36 | if not flag.any(): 37 | continue 38 | 39 | centers = (boxes[:, :2] + boxes[:, 2:]) / 2 40 | mask_a = np.logical_and(roi[:2] < centers, centers < roi[2:]).all(axis=1) 41 | boxes_t = boxes[mask_a].copy() 42 | labels_t = labels[mask_a].copy() 43 | landms_t = landm[mask_a].copy() 44 | landms_t = landms_t.reshape([-1, 5, 2]) 45 | 46 | if boxes_t.shape[0] == 0: 47 | continue 48 | 49 | image_t = image[roi[1]:roi[3], roi[0]:roi[2]] 50 | 51 | boxes_t[:, :2] = np.maximum(boxes_t[:, :2], roi[:2]) 52 | boxes_t[:, :2] -= roi[:2] 53 | boxes_t[:, 2:] = np.minimum(boxes_t[:, 2:], roi[2:]) 54 | boxes_t[:, 2:] -= roi[:2] 55 | 56 | # landm 57 | landms_t[:, :, :2] = landms_t[:, :, :2] - roi[:2] 58 | landms_t[:, :, :2] = np.maximum(landms_t[:, :, :2], np.array([0, 0])) 59 | landms_t[:, :, :2] = np.minimum(landms_t[:, :, :2], roi[2:] - roi[:2]) 60 | landms_t = landms_t.reshape([-1, 10]) 61 | 62 | 63 | # make sure that the cropped image contains at least one face > 16 pixel at training image scale 64 | b_w_t = (boxes_t[:, 2] - boxes_t[:, 0] + 1) / w * img_dim 65 | b_h_t = (boxes_t[:, 3] - boxes_t[:, 1] + 1) / h * img_dim 66 | mask_b = np.minimum(b_w_t, b_h_t) > 0.0 67 | boxes_t = boxes_t[mask_b] 68 | labels_t = labels_t[mask_b] 69 | landms_t = landms_t[mask_b] 70 | 71 | if boxes_t.shape[0] == 0: 72 | continue 73 | 74 | pad_image_flag = False 75 | 76 | return image_t, boxes_t, labels_t, landms_t, pad_image_flag 77 | return image, boxes, labels, landm, pad_image_flag 78 | 79 | 80 | def _distort(image): 81 | 82 | def _convert(image, alpha=1, beta=0): 83 | tmp = image.astype(float) * alpha + beta 84 | tmp[tmp < 0] = 0 85 | tmp[tmp > 255] = 255 86 | image[:] = tmp 87 | 88 | image = image.copy() 89 | 90 | if random.randrange(2): 91 | 92 | #brightness distortion 93 | if random.randrange(2): 94 | _convert(image, beta=random.uniform(-32, 32)) 95 | 96 | #contrast distortion 97 | if random.randrange(2): 98 | _convert(image, alpha=random.uniform(0.5, 1.5)) 99 | 100 | image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV) 101 | 102 | #saturation distortion 103 | if random.randrange(2): 104 | _convert(image[:, :, 1], alpha=random.uniform(0.5, 1.5)) 105 | 106 | #hue distortion 107 | if random.randrange(2): 108 | tmp = image[:, :, 0].astype(int) + random.randint(-18, 18) 109 | tmp %= 180 110 | image[:, :, 0] = tmp 111 | 112 | image = cv2.cvtColor(image, cv2.COLOR_HSV2BGR) 113 | 114 | else: 115 | 116 | #brightness distortion 117 | if random.randrange(2): 118 | _convert(image, beta=random.uniform(-32, 32)) 119 | 120 | image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV) 121 | 122 | #saturation distortion 123 | if random.randrange(2): 124 | _convert(image[:, :, 1], alpha=random.uniform(0.5, 1.5)) 125 | 126 | #hue distortion 127 | if random.randrange(2): 128 | tmp = image[:, :, 0].astype(int) + random.randint(-18, 18) 129 | tmp %= 180 130 | image[:, :, 0] = tmp 131 | 132 | image = cv2.cvtColor(image, cv2.COLOR_HSV2BGR) 133 | 134 | #contrast distortion 135 | if random.randrange(2): 136 | _convert(image, alpha=random.uniform(0.5, 1.5)) 137 | 138 | return image 139 | 140 | 141 | def _expand(image, boxes, fill, p): 142 | if random.randrange(2): 143 | return image, boxes 144 | 145 | height, width, depth = image.shape 146 | 147 | scale = random.uniform(1, p) 148 | w = int(scale * width) 149 | h = int(scale * height) 150 | 151 | left = random.randint(0, w - width) 152 | top = random.randint(0, h - height) 153 | 154 | boxes_t = boxes.copy() 155 | boxes_t[:, :2] += (left, top) 156 | boxes_t[:, 2:] += (left, top) 157 | expand_image = np.empty( 158 | (h, w, depth), 159 | dtype=image.dtype) 160 | expand_image[:, :] = fill 161 | expand_image[top:top + height, left:left + width] = image 162 | image = expand_image 163 | 164 | return image, boxes_t 165 | 166 | 167 | def _mirror(image, boxes, landms): 168 | _, width, _ = image.shape 169 | if random.randrange(2): 170 | image = image[:, ::-1] 171 | boxes = boxes.copy() 172 | boxes[:, 0::2] = width - boxes[:, 2::-2] 173 | 174 | # landm 175 | landms = landms.copy() 176 | landms = landms.reshape([-1, 5, 2]) 177 | landms[:, :, 0] = width - landms[:, :, 0] 178 | tmp = landms[:, 1, :].copy() 179 | landms[:, 1, :] = landms[:, 0, :] 180 | landms[:, 0, :] = tmp 181 | tmp1 = landms[:, 4, :].copy() 182 | landms[:, 4, :] = landms[:, 3, :] 183 | landms[:, 3, :] = tmp1 184 | landms = landms.reshape([-1, 10]) 185 | 186 | return image, boxes, landms 187 | 188 | 189 | def _pad_to_square(image, rgb_mean, pad_image_flag): 190 | if not pad_image_flag: 191 | return image 192 | height, width, _ = image.shape 193 | long_side = max(width, height) 194 | image_t = np.empty((long_side, long_side, 3), dtype=image.dtype) 195 | image_t[:, :] = rgb_mean 196 | image_t[0:0 + height, 0:0 + width] = image 197 | return image_t 198 | 199 | 200 | def _resize_subtract_mean(image, insize, rgb_mean): 201 | interp_methods = [cv2.INTER_LINEAR, cv2.INTER_CUBIC, cv2.INTER_AREA, cv2.INTER_NEAREST, cv2.INTER_LANCZOS4] 202 | interp_method = interp_methods[random.randrange(5)] 203 | image = cv2.resize(image, (insize, insize), interpolation=interp_method) 204 | image = image.astype(np.float32) 205 | image -= rgb_mean 206 | return image.transpose(2, 0, 1) 207 | 208 | 209 | class preproc(object): 210 | 211 | def __init__(self, img_dim, rgb_means): 212 | self.img_dim = img_dim 213 | self.rgb_means = rgb_means 214 | 215 | def __call__(self, image, targets): 216 | assert targets.shape[0] > 0, "this image does not have gt" 217 | 218 | boxes = targets[:, :4].copy() 219 | labels = targets[:, -1].copy() 220 | landm = targets[:, 4:-1].copy() 221 | 222 | image_t, boxes_t, labels_t, landm_t, pad_image_flag = _crop(image, boxes, labels, landm, self.img_dim) 223 | image_t = _distort(image_t) 224 | image_t = _pad_to_square(image_t,self.rgb_means, pad_image_flag) 225 | image_t, boxes_t, landm_t = _mirror(image_t, boxes_t, landm_t) 226 | height, width, _ = image_t.shape 227 | image_t = _resize_subtract_mean(image_t, self.img_dim, self.rgb_means) 228 | boxes_t[:, 0::2] /= width 229 | boxes_t[:, 1::2] /= height 230 | 231 | landm_t[:, 0::2] /= width 232 | landm_t[:, 1::2] /= height 233 | 234 | labels_t = np.expand_dims(labels_t, 1) 235 | targets_t = np.hstack((boxes_t, landm_t, labels_t)) 236 | 237 | return image_t, targets_t 238 | -------------------------------------------------------------------------------- /face_detect/data/wider_face.py: -------------------------------------------------------------------------------- 1 | import os 2 | import os.path 3 | import sys 4 | import torch 5 | import torch.utils.data as data 6 | import cv2 7 | import numpy as np 8 | 9 | class WiderFaceDetection(data.Dataset): 10 | def __init__(self, txt_path, preproc=None): 11 | self.preproc = preproc 12 | self.imgs_path = [] 13 | self.words = [] 14 | f = open(txt_path,'r') 15 | lines = f.readlines() 16 | isFirst = True 17 | labels = [] 18 | for line in lines: 19 | line = line.rstrip() 20 | if line.startswith('#'): 21 | if isFirst is True: 22 | isFirst = False 23 | else: 24 | labels_copy = labels.copy() 25 | self.words.append(labels_copy) 26 | labels.clear() 27 | path = line[2:] 28 | path = txt_path.replace('label.txt','images/') + path 29 | self.imgs_path.append(path) 30 | else: 31 | line = line.split(' ') 32 | label = [float(x) for x in line] 33 | labels.append(label) 34 | 35 | self.words.append(labels) 36 | 37 | def __len__(self): 38 | return len(self.imgs_path) 39 | 40 | def __getitem__(self, index): 41 | img = cv2.imread(self.imgs_path[index]) 42 | height, width, _ = img.shape 43 | 44 | labels = self.words[index] 45 | annotations = np.zeros((0, 15)) 46 | if len(labels) == 0: 47 | return annotations 48 | for idx, label in enumerate(labels): 49 | annotation = np.zeros((1, 15)) 50 | # bbox 51 | annotation[0, 0] = label[0] # x1 52 | annotation[0, 1] = label[1] # y1 53 | annotation[0, 2] = label[0] + label[2] # x2 54 | annotation[0, 3] = label[1] + label[3] # y2 55 | 56 | # landmarks 57 | annotation[0, 4] = label[4] # l0_x 58 | annotation[0, 5] = label[5] # l0_y 59 | annotation[0, 6] = label[7] # l1_x 60 | annotation[0, 7] = label[8] # l1_y 61 | annotation[0, 8] = label[10] # l2_x 62 | annotation[0, 9] = label[11] # l2_y 63 | annotation[0, 10] = label[13] # l3_x 64 | annotation[0, 11] = label[14] # l3_y 65 | annotation[0, 12] = label[16] # l4_x 66 | annotation[0, 13] = label[17] # l4_y 67 | if (annotation[0, 4]<0): 68 | annotation[0, 14] = -1 69 | else: 70 | annotation[0, 14] = 1 71 | 72 | annotations = np.append(annotations, annotation, axis=0) 73 | target = np.array(annotations) 74 | if self.preproc is not None: 75 | img, target = self.preproc(img, target) 76 | 77 | return torch.from_numpy(img), target 78 | 79 | def detection_collate(batch): 80 | """Custom collate fn for dealing with batches of images that have a different 81 | number of associated object annotations (bounding boxes). 82 | 83 | Arguments: 84 | batch: (tuple) A tuple of tensor images and lists of annotations 85 | 86 | Return: 87 | A tuple containing: 88 | 1) (tensor) batch of images stacked on their 0 dim 89 | 2) (list of tensors) annotations for a given image are stacked on 0 dim 90 | """ 91 | targets = [] 92 | imgs = [] 93 | for _, sample in enumerate(batch): 94 | for _, tup in enumerate(sample): 95 | if torch.is_tensor(tup): 96 | imgs.append(tup) 97 | elif isinstance(tup, type(np.empty(0))): 98 | annos = torch.from_numpy(tup).float() 99 | targets.append(annos) 100 | 101 | return (torch.stack(imgs, 0), targets) 102 | -------------------------------------------------------------------------------- /face_detect/facemodels/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/face_detect/facemodels/__init__.py -------------------------------------------------------------------------------- /face_detect/facemodels/net.py: -------------------------------------------------------------------------------- 1 | import time 2 | import torch 3 | import torch.nn as nn 4 | import torchvision.models._utils as _utils 5 | import torchvision.models as models 6 | import torch.nn.functional as F 7 | from torch.autograd import Variable 8 | 9 | def conv_bn(inp, oup, stride = 1, leaky = 0): 10 | return nn.Sequential( 11 | nn.Conv2d(inp, oup, 3, stride, 1, bias=False), 12 | nn.BatchNorm2d(oup), 13 | nn.LeakyReLU(negative_slope=leaky, inplace=True) 14 | ) 15 | 16 | def conv_bn_no_relu(inp, oup, stride): 17 | return nn.Sequential( 18 | nn.Conv2d(inp, oup, 3, stride, 1, bias=False), 19 | nn.BatchNorm2d(oup), 20 | ) 21 | 22 | def conv_bn1X1(inp, oup, stride, leaky=0): 23 | return nn.Sequential( 24 | nn.Conv2d(inp, oup, 1, stride, padding=0, bias=False), 25 | nn.BatchNorm2d(oup), 26 | nn.LeakyReLU(negative_slope=leaky, inplace=True) 27 | ) 28 | 29 | def conv_dw(inp, oup, stride, leaky=0.1): 30 | return nn.Sequential( 31 | nn.Conv2d(inp, inp, 3, stride, 1, groups=inp, bias=False), 32 | nn.BatchNorm2d(inp), 33 | nn.LeakyReLU(negative_slope= leaky,inplace=True), 34 | 35 | nn.Conv2d(inp, oup, 1, 1, 0, bias=False), 36 | nn.BatchNorm2d(oup), 37 | nn.LeakyReLU(negative_slope= leaky,inplace=True), 38 | ) 39 | 40 | class SSH(nn.Module): 41 | def __init__(self, in_channel, out_channel): 42 | super(SSH, self).__init__() 43 | assert out_channel % 4 == 0 44 | leaky = 0 45 | if (out_channel <= 64): 46 | leaky = 0.1 47 | self.conv3X3 = conv_bn_no_relu(in_channel, out_channel//2, stride=1) 48 | 49 | self.conv5X5_1 = conv_bn(in_channel, out_channel//4, stride=1, leaky = leaky) 50 | self.conv5X5_2 = conv_bn_no_relu(out_channel//4, out_channel//4, stride=1) 51 | 52 | self.conv7X7_2 = conv_bn(out_channel//4, out_channel//4, stride=1, leaky = leaky) 53 | self.conv7x7_3 = conv_bn_no_relu(out_channel//4, out_channel//4, stride=1) 54 | 55 | def forward(self, input): 56 | conv3X3 = self.conv3X3(input) 57 | 58 | conv5X5_1 = self.conv5X5_1(input) 59 | conv5X5 = self.conv5X5_2(conv5X5_1) 60 | 61 | conv7X7_2 = self.conv7X7_2(conv5X5_1) 62 | conv7X7 = self.conv7x7_3(conv7X7_2) 63 | 64 | out = torch.cat([conv3X3, conv5X5, conv7X7], dim=1) 65 | out = F.relu(out) 66 | return out 67 | 68 | class FPN(nn.Module): 69 | def __init__(self,in_channels_list,out_channels): 70 | super(FPN,self).__init__() 71 | leaky = 0 72 | if (out_channels <= 64): 73 | leaky = 0.1 74 | self.output1 = conv_bn1X1(in_channels_list[0], out_channels, stride = 1, leaky = leaky) 75 | self.output2 = conv_bn1X1(in_channels_list[1], out_channels, stride = 1, leaky = leaky) 76 | self.output3 = conv_bn1X1(in_channels_list[2], out_channels, stride = 1, leaky = leaky) 77 | 78 | self.merge1 = conv_bn(out_channels, out_channels, leaky = leaky) 79 | self.merge2 = conv_bn(out_channels, out_channels, leaky = leaky) 80 | 81 | def forward(self, input): 82 | # names = list(input.keys()) 83 | input = list(input.values()) 84 | 85 | output1 = self.output1(input[0]) 86 | output2 = self.output2(input[1]) 87 | output3 = self.output3(input[2]) 88 | 89 | up3 = F.interpolate(output3, size=[output2.size(2), output2.size(3)], mode="nearest") 90 | output2 = output2 + up3 91 | output2 = self.merge2(output2) 92 | 93 | up2 = F.interpolate(output2, size=[output1.size(2), output1.size(3)], mode="nearest") 94 | output1 = output1 + up2 95 | output1 = self.merge1(output1) 96 | 97 | out = [output1, output2, output3] 98 | return out 99 | 100 | 101 | 102 | class MobileNetV1(nn.Module): 103 | def __init__(self): 104 | super(MobileNetV1, self).__init__() 105 | self.stage1 = nn.Sequential( 106 | conv_bn(3, 8, 2, leaky = 0.1), # 3 107 | conv_dw(8, 16, 1), # 7 108 | conv_dw(16, 32, 2), # 11 109 | conv_dw(32, 32, 1), # 19 110 | conv_dw(32, 64, 2), # 27 111 | conv_dw(64, 64, 1), # 43 112 | ) 113 | self.stage2 = nn.Sequential( 114 | conv_dw(64, 128, 2), # 43 + 16 = 59 115 | conv_dw(128, 128, 1), # 59 + 32 = 91 116 | conv_dw(128, 128, 1), # 91 + 32 = 123 117 | conv_dw(128, 128, 1), # 123 + 32 = 155 118 | conv_dw(128, 128, 1), # 155 + 32 = 187 119 | conv_dw(128, 128, 1), # 187 + 32 = 219 120 | ) 121 | self.stage3 = nn.Sequential( 122 | conv_dw(128, 256, 2), # 219 +3 2 = 241 123 | conv_dw(256, 256, 1), # 241 + 64 = 301 124 | ) 125 | self.avg = nn.AdaptiveAvgPool2d((1,1)) 126 | self.fc = nn.Linear(256, 1000) 127 | 128 | def forward(self, x): 129 | x = self.stage1(x) 130 | x = self.stage2(x) 131 | x = self.stage3(x) 132 | x = self.avg(x) 133 | # x = self.model(x) 134 | x = x.view(-1, 256) 135 | x = self.fc(x) 136 | return x 137 | 138 | -------------------------------------------------------------------------------- /face_detect/facemodels/retinaface.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torchvision.models.detection.backbone_utils as backbone_utils 4 | import torchvision.models._utils as _utils 5 | import torch.nn.functional as F 6 | from collections import OrderedDict 7 | 8 | from facemodels.net import MobileNetV1 as MobileNetV1 9 | from facemodels.net import FPN as FPN 10 | from facemodels.net import SSH as SSH 11 | 12 | 13 | 14 | class ClassHead(nn.Module): 15 | def __init__(self,inchannels=512,num_anchors=3): 16 | super(ClassHead,self).__init__() 17 | self.num_anchors = num_anchors 18 | self.conv1x1 = nn.Conv2d(inchannels,self.num_anchors*2,kernel_size=(1,1),stride=1,padding=0) 19 | 20 | def forward(self,x): 21 | out = self.conv1x1(x) 22 | out = out.permute(0,2,3,1).contiguous() 23 | 24 | return out.view(out.shape[0], -1, 2) 25 | 26 | class BboxHead(nn.Module): 27 | def __init__(self,inchannels=512,num_anchors=3): 28 | super(BboxHead,self).__init__() 29 | self.conv1x1 = nn.Conv2d(inchannels,num_anchors*4,kernel_size=(1,1),stride=1,padding=0) 30 | 31 | def forward(self,x): 32 | out = self.conv1x1(x) 33 | out = out.permute(0,2,3,1).contiguous() 34 | 35 | return out.view(out.shape[0], -1, 4) 36 | 37 | class LandmarkHead(nn.Module): 38 | def __init__(self,inchannels=512,num_anchors=3): 39 | super(LandmarkHead,self).__init__() 40 | self.conv1x1 = nn.Conv2d(inchannels,num_anchors*10,kernel_size=(1,1),stride=1,padding=0) 41 | 42 | def forward(self,x): 43 | out = self.conv1x1(x) 44 | out = out.permute(0,2,3,1).contiguous() 45 | 46 | return out.view(out.shape[0], -1, 10) 47 | 48 | class RetinaFace(nn.Module): 49 | def __init__(self, cfg = None, phase = 'train'): 50 | """ 51 | :param cfg: Network related settings. 52 | :param phase: train or test. 53 | """ 54 | super(RetinaFace,self).__init__() 55 | self.phase = phase 56 | backbone = None 57 | if cfg['name'] == 'mobilenet0.25': 58 | backbone = MobileNetV1() 59 | if cfg['pretrain']: 60 | checkpoint = torch.load("./weights/mobilenetV1X0.25_pretrain.tar", map_location=torch.device('cpu')) 61 | from collections import OrderedDict 62 | new_state_dict = OrderedDict() 63 | for k, v in checkpoint['state_dict'].items(): 64 | name = k[7:] # remove module. 65 | new_state_dict[name] = v 66 | # load params 67 | backbone.load_state_dict(new_state_dict) 68 | elif cfg['name'] == 'Resnet50': 69 | import torchvision.models as models 70 | backbone = models.resnet50(pretrained=cfg['pretrain']) 71 | 72 | self.body = _utils.IntermediateLayerGetter(backbone, cfg['return_layers']) 73 | in_channels_stage2 = cfg['in_channel'] 74 | in_channels_list = [ 75 | in_channels_stage2 * 2, 76 | in_channels_stage2 * 4, 77 | in_channels_stage2 * 8, 78 | ] 79 | out_channels = cfg['out_channel'] 80 | self.fpn = FPN(in_channels_list,out_channels) 81 | self.ssh1 = SSH(out_channels, out_channels) 82 | self.ssh2 = SSH(out_channels, out_channels) 83 | self.ssh3 = SSH(out_channels, out_channels) 84 | 85 | self.ClassHead = self._make_class_head(fpn_num=3, inchannels=cfg['out_channel']) 86 | self.BboxHead = self._make_bbox_head(fpn_num=3, inchannels=cfg['out_channel']) 87 | self.LandmarkHead = self._make_landmark_head(fpn_num=3, inchannels=cfg['out_channel']) 88 | 89 | def _make_class_head(self,fpn_num=3,inchannels=64,anchor_num=2): 90 | classhead = nn.ModuleList() 91 | for i in range(fpn_num): 92 | classhead.append(ClassHead(inchannels,anchor_num)) 93 | return classhead 94 | 95 | def _make_bbox_head(self,fpn_num=3,inchannels=64,anchor_num=2): 96 | bboxhead = nn.ModuleList() 97 | for i in range(fpn_num): 98 | bboxhead.append(BboxHead(inchannels,anchor_num)) 99 | return bboxhead 100 | 101 | def _make_landmark_head(self,fpn_num=3,inchannels=64,anchor_num=2): 102 | landmarkhead = nn.ModuleList() 103 | for i in range(fpn_num): 104 | landmarkhead.append(LandmarkHead(inchannels,anchor_num)) 105 | return landmarkhead 106 | 107 | def forward(self,inputs): 108 | out = self.body(inputs) 109 | 110 | # FPN 111 | fpn = self.fpn(out) 112 | 113 | # SSH 114 | feature1 = self.ssh1(fpn[0]) 115 | feature2 = self.ssh2(fpn[1]) 116 | feature3 = self.ssh3(fpn[2]) 117 | features = [feature1, feature2, feature3] 118 | 119 | bbox_regressions = torch.cat([self.BboxHead[i](feature) for i, feature in enumerate(features)], dim=1) 120 | classifications = torch.cat([self.ClassHead[i](feature) for i, feature in enumerate(features)],dim=1) 121 | ldm_regressions = torch.cat([self.LandmarkHead[i](feature) for i, feature in enumerate(features)], dim=1) 122 | 123 | if self.phase == 'train': 124 | output = (bbox_regressions, classifications, ldm_regressions) 125 | else: 126 | output = (bbox_regressions, F.softmax(classifications, dim=-1), ldm_regressions) 127 | return output -------------------------------------------------------------------------------- /face_detect/layers/__init__.py: -------------------------------------------------------------------------------- 1 | from .functions import * 2 | from .modules import * 3 | -------------------------------------------------------------------------------- /face_detect/layers/functions/prior_box.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from itertools import product as product 3 | import numpy as np 4 | from math import ceil 5 | 6 | 7 | class PriorBox(object): 8 | def __init__(self, cfg, image_size=None, phase='train'): 9 | super(PriorBox, self).__init__() 10 | self.min_sizes = cfg['min_sizes'] 11 | self.steps = cfg['steps'] 12 | self.clip = cfg['clip'] 13 | self.image_size = image_size 14 | self.feature_maps = [[ceil(self.image_size[0]/step), ceil(self.image_size[1]/step)] for step in self.steps] 15 | self.name = "s" 16 | 17 | def forward(self): 18 | anchors = [] 19 | for k, f in enumerate(self.feature_maps): 20 | min_sizes = self.min_sizes[k] 21 | for i, j in product(range(f[0]), range(f[1])): 22 | for min_size in min_sizes: 23 | s_kx = min_size / self.image_size[1] 24 | s_ky = min_size / self.image_size[0] 25 | dense_cx = [x * self.steps[k] / self.image_size[1] for x in [j + 0.5]] 26 | dense_cy = [y * self.steps[k] / self.image_size[0] for y in [i + 0.5]] 27 | for cy, cx in product(dense_cy, dense_cx): 28 | anchors += [cx, cy, s_kx, s_ky] 29 | 30 | # back to torch land 31 | output = torch.Tensor(anchors).view(-1, 4) 32 | if self.clip: 33 | output.clamp_(max=1, min=0) 34 | return output 35 | -------------------------------------------------------------------------------- /face_detect/layers/modules/__init__.py: -------------------------------------------------------------------------------- 1 | from .multibox_loss import MultiBoxLoss 2 | 3 | __all__ = ['MultiBoxLoss'] 4 | -------------------------------------------------------------------------------- /face_detect/layers/modules/multibox_loss.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | from torch.autograd import Variable 5 | from utils.box_utils import match, log_sum_exp 6 | from data import cfg_mnet 7 | GPU = cfg_mnet['gpu_train'] 8 | 9 | class MultiBoxLoss(nn.Module): 10 | """SSD Weighted Loss Function 11 | Compute Targets: 12 | 1) Produce Confidence Target Indices by matching ground truth boxes 13 | with (default) 'priorboxes' that have jaccard index > threshold parameter 14 | (default threshold: 0.5). 15 | 2) Produce localization target by 'encoding' variance into offsets of ground 16 | truth boxes and their matched 'priorboxes'. 17 | 3) Hard negative mining to filter the excessive number of negative examples 18 | that comes with using a large number of default bounding boxes. 19 | (default negative:positive ratio 3:1) 20 | Objective Loss: 21 | L(x,c,l,g) = (Lconf(x, c) + αLloc(x,l,g)) / N 22 | Where, Lconf is the CrossEntropy Loss and Lloc is the SmoothL1 Loss 23 | weighted by α which is set to 1 by cross val. 24 | Args: 25 | c: class confidences, 26 | l: predicted boxes, 27 | g: ground truth boxes 28 | N: number of matched default boxes 29 | See: https://arxiv.org/pdf/1512.02325.pdf for more details. 30 | """ 31 | 32 | def __init__(self, num_classes, overlap_thresh, prior_for_matching, bkg_label, neg_mining, neg_pos, neg_overlap, encode_target): 33 | super(MultiBoxLoss, self).__init__() 34 | self.num_classes = num_classes 35 | self.threshold = overlap_thresh 36 | self.background_label = bkg_label 37 | self.encode_target = encode_target 38 | self.use_prior_for_matching = prior_for_matching 39 | self.do_neg_mining = neg_mining 40 | self.negpos_ratio = neg_pos 41 | self.neg_overlap = neg_overlap 42 | self.variance = [0.1, 0.2] 43 | 44 | def forward(self, predictions, priors, targets): 45 | """Multibox Loss 46 | Args: 47 | predictions (tuple): A tuple containing loc preds, conf preds, 48 | and prior boxes from SSD net. 49 | conf shape: torch.size(batch_size,num_priors,num_classes) 50 | loc shape: torch.size(batch_size,num_priors,4) 51 | priors shape: torch.size(num_priors,4) 52 | 53 | ground_truth (tensor): Ground truth boxes and labels for a batch, 54 | shape: [batch_size,num_objs,5] (last idx is the label). 55 | """ 56 | 57 | loc_data, conf_data, landm_data = predictions 58 | priors = priors 59 | num = loc_data.size(0) 60 | num_priors = (priors.size(0)) 61 | 62 | # match priors (default boxes) and ground truth boxes 63 | loc_t = torch.Tensor(num, num_priors, 4) 64 | landm_t = torch.Tensor(num, num_priors, 10) 65 | conf_t = torch.LongTensor(num, num_priors) 66 | for idx in range(num): 67 | truths = targets[idx][:, :4].data 68 | labels = targets[idx][:, -1].data 69 | landms = targets[idx][:, 4:14].data 70 | defaults = priors.data 71 | match(self.threshold, truths, defaults, self.variance, labels, landms, loc_t, conf_t, landm_t, idx) 72 | if GPU: 73 | loc_t = loc_t.cuda() 74 | conf_t = conf_t.cuda() 75 | landm_t = landm_t.cuda() 76 | 77 | zeros = torch.tensor(0).cuda() 78 | # landm Loss (Smooth L1) 79 | # Shape: [batch,num_priors,10] 80 | pos1 = conf_t > zeros 81 | num_pos_landm = pos1.long().sum(1, keepdim=True) 82 | N1 = max(num_pos_landm.data.sum().float(), 1) 83 | pos_idx1 = pos1.unsqueeze(pos1.dim()).expand_as(landm_data) 84 | landm_p = landm_data[pos_idx1].view(-1, 10) 85 | landm_t = landm_t[pos_idx1].view(-1, 10) 86 | loss_landm = F.smooth_l1_loss(landm_p, landm_t, reduction='sum') 87 | 88 | 89 | pos = conf_t != zeros 90 | conf_t[pos] = 1 91 | 92 | # Localization Loss (Smooth L1) 93 | # Shape: [batch,num_priors,4] 94 | pos_idx = pos.unsqueeze(pos.dim()).expand_as(loc_data) 95 | loc_p = loc_data[pos_idx].view(-1, 4) 96 | loc_t = loc_t[pos_idx].view(-1, 4) 97 | loss_l = F.smooth_l1_loss(loc_p, loc_t, reduction='sum') 98 | 99 | # Compute max conf across batch for hard negative mining 100 | batch_conf = conf_data.view(-1, self.num_classes) 101 | loss_c = log_sum_exp(batch_conf) - batch_conf.gather(1, conf_t.view(-1, 1)) 102 | 103 | # Hard Negative Mining 104 | loss_c[pos.view(-1, 1)] = 0 # filter out pos boxes for now 105 | loss_c = loss_c.view(num, -1) 106 | _, loss_idx = loss_c.sort(1, descending=True) 107 | _, idx_rank = loss_idx.sort(1) 108 | num_pos = pos.long().sum(1, keepdim=True) 109 | num_neg = torch.clamp(self.negpos_ratio*num_pos, max=pos.size(1)-1) 110 | neg = idx_rank < num_neg.expand_as(idx_rank) 111 | 112 | # Confidence Loss Including Positive and Negative Examples 113 | pos_idx = pos.unsqueeze(2).expand_as(conf_data) 114 | neg_idx = neg.unsqueeze(2).expand_as(conf_data) 115 | conf_p = conf_data[(pos_idx+neg_idx).gt(0)].view(-1,self.num_classes) 116 | targets_weighted = conf_t[(pos+neg).gt(0)] 117 | loss_c = F.cross_entropy(conf_p, targets_weighted, reduction='sum') 118 | 119 | # Sum of losses: L(x,c,l,g) = (Lconf(x, c) + αLloc(x,l,g)) / N 120 | N = max(num_pos.data.sum().float(), 1) 121 | loss_l /= N 122 | loss_c /= N 123 | loss_landm /= N1 124 | 125 | return loss_l, loss_c, loss_landm 126 | -------------------------------------------------------------------------------- /face_detect/retinaface_detection.py: -------------------------------------------------------------------------------- 1 | ''' 2 | @paper: GAN Prior Embedded Network for Blind Face Restoration in the Wild (CVPR2021) 3 | @author: yangxy (yangtao9009@gmail.com) 4 | ''' 5 | import os 6 | import torch 7 | import torch.backends.cudnn as cudnn 8 | import numpy as np 9 | from data import cfg_re50 10 | from layers.functions.prior_box import PriorBox 11 | from utils.nms.py_cpu_nms import py_cpu_nms 12 | import cv2 13 | from facemodels.retinaface import RetinaFace 14 | from utils.box_utils import decode, decode_landm 15 | import time 16 | import torch.nn.functional as F 17 | 18 | 19 | class RetinaFaceDetection(object): 20 | def __init__(self, base_dir, device='cuda', network='RetinaFace-R50'): 21 | torch.set_grad_enabled(False) 22 | cudnn.benchmark = True 23 | self.pretrained_path = os.path.join(base_dir, 'weights', network+'.pth') 24 | self.device = device #torch.cuda.current_device() 25 | self.cfg = cfg_re50 26 | self.net = RetinaFace(cfg=self.cfg, phase='test') 27 | self.load_model() 28 | self.net = self.net.to(device) 29 | 30 | self.mean = torch.tensor([[[[104]], [[117]], [[123]]]]).to(device) 31 | 32 | def check_keys(self, pretrained_state_dict): 33 | ckpt_keys = set(pretrained_state_dict.keys()) 34 | model_keys = set(self.net.state_dict().keys()) 35 | used_pretrained_keys = model_keys & ckpt_keys 36 | unused_pretrained_keys = ckpt_keys - model_keys 37 | missing_keys = model_keys - ckpt_keys 38 | assert len(used_pretrained_keys) > 0, 'load NONE from pretrained checkpoint' 39 | return True 40 | 41 | def remove_prefix(self, state_dict, prefix): 42 | ''' Old style model is stored with all names of parameters sharing common prefix 'module.' ''' 43 | f = lambda x: x.split(prefix, 1)[-1] if x.startswith(prefix) else x 44 | return {f(key): value for key, value in state_dict.items()} 45 | 46 | def load_model(self, load_to_cpu=False): 47 | #if load_to_cpu: 48 | # pretrained_dict = torch.load(self.pretrained_path, map_location=lambda storage, loc: storage) 49 | #else: 50 | # pretrained_dict = torch.load(self.pretrained_path, map_location=lambda storage, loc: storage.cuda()) 51 | pretrained_dict = torch.load(self.pretrained_path, map_location=torch.device('cpu')) 52 | if "state_dict" in pretrained_dict.keys(): 53 | pretrained_dict = self.remove_prefix(pretrained_dict['state_dict'], 'module.') 54 | else: 55 | pretrained_dict = self.remove_prefix(pretrained_dict, 'module.') 56 | self.check_keys(pretrained_dict) 57 | self.net.load_state_dict(pretrained_dict, strict=False) 58 | self.net.eval() 59 | 60 | def detect(self, img_raw, resize=1, confidence_threshold=0.9, nms_threshold=0.4, top_k=5000, keep_top_k=750, save_image=False): 61 | img = np.float32(img_raw) 62 | 63 | im_height, im_width = img.shape[:2] 64 | ss = 1.0 65 | # tricky 66 | if max(im_height, im_width) > 1500: 67 | ss = 1000.0/max(im_height, im_width) 68 | img = cv2.resize(img, (0,0), fx=ss, fy=ss) 69 | im_height, im_width = img.shape[:2] 70 | 71 | scale = torch.Tensor([img.shape[1], img.shape[0], img.shape[1], img.shape[0]]) 72 | img -= (104, 117, 123) 73 | img = img.transpose(2, 0, 1) 74 | img = torch.from_numpy(img).unsqueeze(0) 75 | img = img.to(self.device) 76 | scale = scale.to(self.device) 77 | 78 | loc, conf, landms = self.net(img) # forward pass 79 | del img 80 | 81 | priorbox = PriorBox(self.cfg, image_size=(im_height, im_width)) 82 | priors = priorbox.forward() 83 | priors = priors.to(self.device) 84 | prior_data = priors.data 85 | boxes = decode(loc.data.squeeze(0), prior_data, self.cfg['variance']) 86 | boxes = boxes * scale / resize 87 | boxes = boxes.cpu().numpy() 88 | scores = conf.squeeze(0).data.cpu().numpy()[:, 1] 89 | landms = decode_landm(landms.data.squeeze(0), prior_data, self.cfg['variance']) 90 | scale1 = torch.Tensor([im_width, im_height, im_width, im_height, 91 | im_width, im_height, im_width, im_height, 92 | im_width, im_height]) 93 | scale1 = scale1.to(self.device) 94 | landms = landms * scale1 / resize 95 | landms = landms.cpu().numpy() 96 | 97 | # ignore low scores 98 | inds = np.where(scores > confidence_threshold)[0] 99 | boxes = boxes[inds] 100 | landms = landms[inds] 101 | scores = scores[inds] 102 | 103 | # keep top-K before NMS 104 | order = scores.argsort()[::-1][:top_k] 105 | boxes = boxes[order] 106 | landms = landms[order] 107 | scores = scores[order] 108 | 109 | # do NMS 110 | dets = np.hstack((boxes, scores[:, np.newaxis])).astype(np.float32, copy=False) 111 | keep = py_cpu_nms(dets, nms_threshold) 112 | # keep = nms(dets, nms_threshold,force_cpu=args.cpu) 113 | dets = dets[keep, :] 114 | landms = landms[keep] 115 | 116 | # keep top-K faster NMS 117 | dets = dets[:keep_top_k, :] 118 | landms = landms[:keep_top_k, :] 119 | 120 | # sort faces(delete) 121 | ''' 122 | fscores = [det[4] for det in dets] 123 | sorted_idx = sorted(range(len(fscores)), key=lambda k:fscores[k], reverse=False) # sort index 124 | tmp = [landms[idx] for idx in sorted_idx] 125 | landms = np.asarray(tmp) 126 | ''' 127 | 128 | landms = landms.reshape((-1, 5, 2)) 129 | landms = landms.transpose((0, 2, 1)) 130 | landms = landms.reshape(-1, 10, ) 131 | return dets/ss, landms/ss 132 | 133 | def detect_tensor(self, img, resize=1, confidence_threshold=0.9, nms_threshold=0.4, top_k=5000, keep_top_k=750, save_image=False): 134 | im_height, im_width = img.shape[-2:] 135 | ss = 1000/max(im_height, im_width) 136 | img = F.interpolate(img, scale_factor=ss) 137 | im_height, im_width = img.shape[-2:] 138 | scale = torch.Tensor([im_width, im_height, im_width, im_height]).to(self.device) 139 | img -= self.mean 140 | 141 | loc, conf, landms = self.net(img) # forward pass 142 | 143 | priorbox = PriorBox(self.cfg, image_size=(im_height, im_width)) 144 | priors = priorbox.forward() 145 | priors = priors.to(self.device) 146 | prior_data = priors.data 147 | boxes = decode(loc.data.squeeze(0), prior_data, self.cfg['variance']) 148 | boxes = boxes * scale / resize 149 | boxes = boxes.cpu().numpy() 150 | scores = conf.squeeze(0).data.cpu().numpy()[:, 1] 151 | landms = decode_landm(landms.data.squeeze(0), prior_data, self.cfg['variance']) 152 | scale1 = torch.Tensor([img.shape[3], img.shape[2], img.shape[3], img.shape[2], 153 | img.shape[3], img.shape[2], img.shape[3], img.shape[2], 154 | img.shape[3], img.shape[2]]) 155 | scale1 = scale1.to(self.device) 156 | landms = landms * scale1 / resize 157 | landms = landms.cpu().numpy() 158 | 159 | # ignore low scores 160 | inds = np.where(scores > confidence_threshold)[0] 161 | boxes = boxes[inds] 162 | landms = landms[inds] 163 | scores = scores[inds] 164 | 165 | # keep top-K before NMS 166 | order = scores.argsort()[::-1][:top_k] 167 | boxes = boxes[order] 168 | landms = landms[order] 169 | scores = scores[order] 170 | 171 | # do NMS 172 | dets = np.hstack((boxes, scores[:, np.newaxis])).astype(np.float32, copy=False) 173 | keep = py_cpu_nms(dets, nms_threshold) 174 | # keep = nms(dets, nms_threshold,force_cpu=args.cpu) 175 | dets = dets[keep, :] 176 | landms = landms[keep] 177 | 178 | # keep top-K faster NMS 179 | dets = dets[:keep_top_k, :] 180 | landms = landms[:keep_top_k, :] 181 | 182 | # sort faces(delete) 183 | ''' 184 | fscores = [det[4] for det in dets] 185 | sorted_idx = sorted(range(len(fscores)), key=lambda k:fscores[k], reverse=False) # sort index 186 | tmp = [landms[idx] for idx in sorted_idx] 187 | landms = np.asarray(tmp) 188 | ''' 189 | 190 | landms = landms.reshape((-1, 5, 2)) 191 | landms = landms.transpose((0, 2, 1)) 192 | landms = landms.reshape(-1, 10, ) 193 | return dets/ss, landms/ss 194 | -------------------------------------------------------------------------------- /face_detect/utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/face_detect/utils/__init__.py -------------------------------------------------------------------------------- /face_detect/utils/nms/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/face_detect/utils/nms/__init__.py -------------------------------------------------------------------------------- /face_detect/utils/nms/py_cpu_nms.py: -------------------------------------------------------------------------------- 1 | # -------------------------------------------------------- 2 | # Fast R-CNN 3 | # Copyright (c) 2015 Microsoft 4 | # Licensed under The MIT License [see LICENSE for details] 5 | # Written by Ross Girshick 6 | # -------------------------------------------------------- 7 | 8 | import numpy as np 9 | 10 | def py_cpu_nms(dets, thresh): 11 | """Pure Python NMS baseline.""" 12 | x1 = dets[:, 0] 13 | y1 = dets[:, 1] 14 | x2 = dets[:, 2] 15 | y2 = dets[:, 3] 16 | scores = dets[:, 4] 17 | 18 | areas = (x2 - x1 + 1) * (y2 - y1 + 1) 19 | order = scores.argsort()[::-1] 20 | 21 | keep = [] 22 | while order.size > 0: 23 | i = order[0] 24 | keep.append(i) 25 | xx1 = np.maximum(x1[i], x1[order[1:]]) 26 | yy1 = np.maximum(y1[i], y1[order[1:]]) 27 | xx2 = np.minimum(x2[i], x2[order[1:]]) 28 | yy2 = np.minimum(y2[i], y2[order[1:]]) 29 | 30 | w = np.maximum(0.0, xx2 - xx1 + 1) 31 | h = np.maximum(0.0, yy2 - yy1 + 1) 32 | inter = w * h 33 | ovr = inter / (areas[i] + areas[order[1:]] - inter) 34 | 35 | inds = np.where(ovr <= thresh)[0] 36 | order = order[inds + 1] 37 | 38 | return keep 39 | -------------------------------------------------------------------------------- /face_detect/utils/timer.py: -------------------------------------------------------------------------------- 1 | # -------------------------------------------------------- 2 | # Fast R-CNN 3 | # Copyright (c) 2015 Microsoft 4 | # Licensed under The MIT License [see LICENSE for details] 5 | # Written by Ross Girshick 6 | # -------------------------------------------------------- 7 | 8 | import time 9 | 10 | 11 | class Timer(object): 12 | """A simple timer.""" 13 | def __init__(self): 14 | self.total_time = 0. 15 | self.calls = 0 16 | self.start_time = 0. 17 | self.diff = 0. 18 | self.average_time = 0. 19 | 20 | def tic(self): 21 | # using time.time instead of time.clock because time time.clock 22 | # does not normalize for multithreading 23 | self.start_time = time.time() 24 | 25 | def toc(self, average=True): 26 | self.diff = time.time() - self.start_time 27 | self.total_time += self.diff 28 | self.calls += 1 29 | self.average_time = self.total_time / self.calls 30 | if average: 31 | return self.average_time 32 | else: 33 | return self.diff 34 | 35 | def clear(self): 36 | self.total_time = 0. 37 | self.calls = 0 38 | self.start_time = 0. 39 | self.diff = 0. 40 | self.average_time = 0. 41 | -------------------------------------------------------------------------------- /face_enhancement.py: -------------------------------------------------------------------------------- 1 | ''' 2 | @paper: GAN Prior Embedded Network for Blind Face Restoration in the Wild (CVPR2021) 3 | @author: yangxy (yangtao9009@gmail.com) 4 | ''' 5 | import cv2 6 | import time 7 | import numpy as np 8 | import __init_paths 9 | from face_detect.retinaface_detection import RetinaFaceDetection 10 | from face_parse.face_parsing import FaceParse 11 | from face_model.face_gan import FaceGAN 12 | from sr_model.real_esrnet import RealESRNet 13 | from align_faces import warp_and_crop_face, get_reference_facial_points 14 | 15 | class FaceEnhancement(object): 16 | def __init__(self, args, base_dir='./', in_size=512, out_size=None, model=None, use_sr=True, device='cuda'): 17 | self.facedetector = RetinaFaceDetection(base_dir, device) 18 | self.facegan = FaceGAN(base_dir, in_size, out_size, model, args.channel_multiplier, args.narrow, args.key, device=device) 19 | self.srmodel = RealESRNet(base_dir, args.sr_model, args.sr_scale, args.tile_size, device=device) 20 | self.faceparser = FaceParse(base_dir, device=device) 21 | self.use_sr = use_sr 22 | self.in_size = in_size 23 | self.out_size = in_size if out_size is None else out_size 24 | self.threshold = 0.9 25 | self.alpha = args.alpha 26 | 27 | # the mask for pasting restored faces back 28 | self.mask = np.zeros((512, 512), np.float32) 29 | cv2.rectangle(self.mask, (26, 26), (486, 486), (1, 1, 1), -1, cv2.LINE_AA) 30 | self.mask = cv2.GaussianBlur(self.mask, (101, 101), 4) 31 | self.mask = cv2.GaussianBlur(self.mask, (101, 101), 4) 32 | 33 | self.kernel = np.array(( 34 | [0.0625, 0.125, 0.0625], 35 | [0.125, 0.25, 0.125], 36 | [0.0625, 0.125, 0.0625]), dtype="float32") 37 | 38 | # get the reference 5 landmarks position in the crop settings 39 | default_square = True 40 | inner_padding_factor = 0.25 41 | outer_padding = (0, 0) 42 | self.reference_5pts = get_reference_facial_points( 43 | (self.in_size, self.in_size), inner_padding_factor, outer_padding, default_square) 44 | 45 | def mask_postprocess(self, mask, thres=26): 46 | mask[:thres, :] = 0; mask[-thres:, :] = 0 47 | mask[:, :thres] = 0; mask[:, -thres:] = 0 48 | mask = cv2.GaussianBlur(mask, (101, 101), 4) 49 | mask = cv2.GaussianBlur(mask, (101, 101), 4) 50 | return mask.astype(np.float32) 51 | 52 | def process(self, img, aligned=False): 53 | orig_faces, enhanced_faces = [], [] 54 | if aligned: 55 | ef = self.facegan.process(img) 56 | orig_faces.append(img) 57 | enhanced_faces.append(ef) 58 | 59 | if self.use_sr: 60 | ef = self.srmodel.process(ef) 61 | 62 | return ef, orig_faces, enhanced_faces 63 | 64 | if self.use_sr: 65 | img_sr = self.srmodel.process(img) 66 | if img_sr is not None: 67 | img = cv2.resize(img, img_sr.shape[:2][::-1]) 68 | 69 | facebs, landms = self.facedetector.detect(img) 70 | 71 | height, width = img.shape[:2] 72 | full_mask = np.zeros((height, width), dtype=np.float32) 73 | full_img = np.zeros(img.shape, dtype=np.uint8) 74 | 75 | for i, (faceb, facial5points) in enumerate(zip(facebs, landms)): 76 | if faceb[4]0)] = tmp_mask[np.where(mask>0)] 105 | full_img[np.where(mask>0)] = tmp_img[np.where(mask>0)] 106 | 107 | full_mask = full_mask[:, :, np.newaxis] 108 | if self.use_sr and img_sr is not None: 109 | img = cv2.convertScaleAbs(img_sr*(1-full_mask) + full_img*full_mask) 110 | else: 111 | img = cv2.convertScaleAbs(img*(1-full_mask) + full_img*full_mask) 112 | 113 | return img, orig_faces, enhanced_faces 114 | 115 | 116 | -------------------------------------------------------------------------------- /face_inpainting.py: -------------------------------------------------------------------------------- 1 | ''' 2 | @paper: GAN Prior Embedded Network for Blind Face Restoration in the Wild (CVPR2021) 3 | @author: yangxy (yangtao9009@gmail.com) 4 | ''' 5 | from face_model.face_gan import FaceGAN 6 | 7 | class FaceInpainting(object): 8 | def __init__(self, base_dir='./', in_size=1024, out_size=1024, model=None, channel_multiplier=2, narrow=1, key=None, device='cuda'): 9 | self.facegan = FaceGAN(base_dir, in_size, out_size, model, channel_multiplier, narrow, key, device=device) 10 | 11 | # make sure the face image is well aligned. Please refer to face_enhancement.py 12 | def process(self, brokenf, aligned=True): 13 | # complete the face 14 | out = self.facegan.process(brokenf) 15 | 16 | return out, [brokenf], [out] 17 | 18 | 19 | -------------------------------------------------------------------------------- /face_model/face_gan.py: -------------------------------------------------------------------------------- 1 | ''' 2 | @paper: GAN Prior Embedded Network for Blind Face Restoration in the Wild (CVPR2021) 3 | @author: yangxy (yangtao9009@gmail.com) 4 | ''' 5 | import torch 6 | import os 7 | import cv2 8 | import glob 9 | import numpy as np 10 | from torch import nn 11 | import torch.nn.functional as F 12 | from torchvision import transforms, utils 13 | from gpen_model import FullGenerator, FullGenerator_SR 14 | 15 | class FaceGAN(object): 16 | def __init__(self, base_dir='./', in_size=512, out_size=None, model=None, channel_multiplier=2, narrow=1, key=None, is_norm=True, device='cuda'): 17 | self.mfile = os.path.join(base_dir, 'weights', model+'.pth') 18 | self.n_mlp = 8 19 | self.device = device 20 | self.is_norm = is_norm 21 | self.in_resolution = in_size 22 | self.out_resolution = in_size if out_size is None else out_size 23 | self.key = key 24 | self.load_model(channel_multiplier, narrow) 25 | 26 | def load_model(self, channel_multiplier=2, narrow=1): 27 | if self.in_resolution == self.out_resolution: 28 | self.model = FullGenerator(self.in_resolution, 512, self.n_mlp, channel_multiplier, narrow=narrow, device=self.device) 29 | else: 30 | self.model = FullGenerator_SR(self.in_resolution, self.out_resolution, 512, self.n_mlp, channel_multiplier, narrow=narrow, device=self.device) 31 | pretrained_dict = torch.load(self.mfile, map_location=torch.device('cpu')) 32 | if self.key is not None: pretrained_dict = pretrained_dict[self.key] 33 | self.model.load_state_dict(pretrained_dict) 34 | self.model.to(self.device) 35 | self.model.eval() 36 | 37 | def process(self, img): 38 | img = cv2.resize(img, (self.in_resolution, self.in_resolution)) 39 | img_t = self.img2tensor(img) 40 | 41 | with torch.no_grad(): 42 | out, __ = self.model(img_t) 43 | del img_t 44 | 45 | out = self.tensor2img(out) 46 | 47 | return out 48 | 49 | def img2tensor(self, img): 50 | img_t = torch.from_numpy(img).to(self.device)/255. 51 | if self.is_norm: 52 | img_t = (img_t - 0.5) / 0.5 53 | img_t = img_t.permute(2, 0, 1).unsqueeze(0).flip(1) # BGR->RGB 54 | return img_t 55 | 56 | def tensor2img(self, img_t, pmax=255.0, imtype=np.uint8): 57 | if self.is_norm: 58 | img_t = img_t * 0.5 + 0.5 59 | img_t = img_t.squeeze(0).permute(1, 2, 0).flip(2) # RGB->BGR 60 | img_np = np.clip(img_t.float().cpu().numpy(), 0, 1) * pmax 61 | 62 | return img_np.astype(imtype) 63 | -------------------------------------------------------------------------------- /face_model/op/__init__.py: -------------------------------------------------------------------------------- 1 | from .fused_act import FusedLeakyReLU, fused_leaky_relu 2 | from .upfirdn2d import upfirdn2d 3 | -------------------------------------------------------------------------------- /face_model/op/fused_act.py: -------------------------------------------------------------------------------- 1 | import os 2 | import platform 3 | 4 | import torch 5 | from torch import nn 6 | import torch.nn.functional as F 7 | from torch.autograd import Function 8 | from torch.utils.cpp_extension import load, _import_module_from_library 9 | 10 | # if running GPEN without cuda, please comment line 11-19 11 | if platform.system() == 'Linux' and torch.cuda.is_available(): 12 | module_path = os.path.dirname(__file__) 13 | fused = load( 14 | 'fused', 15 | sources=[ 16 | os.path.join(module_path, 'fused_bias_act.cpp'), 17 | os.path.join(module_path, 'fused_bias_act_kernel.cu'), 18 | ], 19 | ) 20 | 21 | 22 | #fused = _import_module_from_library('fused', '/tmp/torch_extensions/fused', True) 23 | 24 | 25 | class FusedLeakyReLUFunctionBackward(Function): 26 | @staticmethod 27 | def forward(ctx, grad_output, out, negative_slope, scale): 28 | ctx.save_for_backward(out) 29 | ctx.negative_slope = negative_slope 30 | ctx.scale = scale 31 | 32 | empty = grad_output.new_empty(0) 33 | 34 | grad_input = fused.fused_bias_act( 35 | grad_output, empty, out, 3, 1, negative_slope, scale 36 | ) 37 | 38 | dim = [0] 39 | 40 | if grad_input.ndim > 2: 41 | dim += list(range(2, grad_input.ndim)) 42 | 43 | grad_bias = grad_input.sum(dim).detach() 44 | 45 | return grad_input, grad_bias 46 | 47 | @staticmethod 48 | def backward(ctx, gradgrad_input, gradgrad_bias): 49 | out, = ctx.saved_tensors 50 | gradgrad_out = fused.fused_bias_act( 51 | gradgrad_input, gradgrad_bias, out, 3, 1, ctx.negative_slope, ctx.scale 52 | ) 53 | 54 | return gradgrad_out, None, None, None 55 | 56 | 57 | class FusedLeakyReLUFunction(Function): 58 | @staticmethod 59 | def forward(ctx, input, bias, negative_slope, scale): 60 | empty = input.new_empty(0) 61 | out = fused.fused_bias_act(input, bias, empty, 3, 0, negative_slope, scale) 62 | ctx.save_for_backward(out) 63 | ctx.negative_slope = negative_slope 64 | ctx.scale = scale 65 | 66 | return out 67 | 68 | @staticmethod 69 | def backward(ctx, grad_output): 70 | out, = ctx.saved_tensors 71 | 72 | grad_input, grad_bias = FusedLeakyReLUFunctionBackward.apply( 73 | grad_output, out, ctx.negative_slope, ctx.scale 74 | ) 75 | 76 | return grad_input, grad_bias, None, None 77 | 78 | 79 | class FusedLeakyReLU(nn.Module): 80 | def __init__(self, channel, negative_slope=0.2, scale=2 ** 0.5, device='cpu'): 81 | super().__init__() 82 | 83 | self.bias = nn.Parameter(torch.zeros(channel)) 84 | self.negative_slope = negative_slope 85 | self.scale = scale 86 | self.device = device 87 | 88 | def forward(self, input): 89 | return fused_leaky_relu(input, self.bias, self.negative_slope, self.scale, self.device) 90 | 91 | 92 | def fused_leaky_relu(input, bias, negative_slope=0.2, scale=2 ** 0.5, device='cpu'): 93 | if platform.system() == 'Linux' and torch.cuda.is_available() and device != 'cpu': 94 | return FusedLeakyReLUFunction.apply(input, bias, negative_slope, scale) 95 | else: 96 | return scale * F.leaky_relu(input + bias.view((1, -1)+(1,)*(len(input.shape)-2)), negative_slope=negative_slope) 97 | -------------------------------------------------------------------------------- /face_model/op/fused_bias_act.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | 4 | torch::Tensor fused_bias_act_op(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer, 5 | int act, int grad, float alpha, float scale); 6 | 7 | #define CHECK_CUDA(x) TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor") 8 | #define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous") 9 | #define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) 10 | 11 | torch::Tensor fused_bias_act(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer, 12 | int act, int grad, float alpha, float scale) { 13 | CHECK_CUDA(input); 14 | CHECK_CUDA(bias); 15 | 16 | return fused_bias_act_op(input, bias, refer, act, grad, alpha, scale); 17 | } 18 | 19 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 20 | m.def("fused_bias_act", &fused_bias_act, "fused bias act (CUDA)"); 21 | } -------------------------------------------------------------------------------- /face_model/op/fused_bias_act_kernel.cu: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2019, NVIDIA Corporation. All rights reserved. 2 | // 3 | // This work is made available under the Nvidia Source Code License-NC. 4 | // To view a copy of this license, visit 5 | // https://nvlabs.github.io/stylegan2/license.html 6 | 7 | #include 8 | 9 | #include 10 | #include 11 | #include 12 | #include 13 | 14 | #include 15 | #include 16 | 17 | 18 | template 19 | static __global__ void fused_bias_act_kernel(scalar_t* out, const scalar_t* p_x, const scalar_t* p_b, const scalar_t* p_ref, 20 | int act, int grad, scalar_t alpha, scalar_t scale, int loop_x, int size_x, int step_b, int size_b, int use_bias, int use_ref) { 21 | int xi = blockIdx.x * loop_x * blockDim.x + threadIdx.x; 22 | 23 | scalar_t zero = 0.0; 24 | 25 | for (int loop_idx = 0; loop_idx < loop_x && xi < size_x; loop_idx++, xi += blockDim.x) { 26 | scalar_t x = p_x[xi]; 27 | 28 | if (use_bias) { 29 | x += p_b[(xi / step_b) % size_b]; 30 | } 31 | 32 | scalar_t ref = use_ref ? p_ref[xi] : zero; 33 | 34 | scalar_t y; 35 | 36 | switch (act * 10 + grad) { 37 | default: 38 | case 10: y = x; break; 39 | case 11: y = x; break; 40 | case 12: y = 0.0; break; 41 | 42 | case 30: y = (x > 0.0) ? x : x * alpha; break; 43 | case 31: y = (ref > 0.0) ? x : x * alpha; break; 44 | case 32: y = 0.0; break; 45 | } 46 | 47 | out[xi] = y * scale; 48 | } 49 | } 50 | 51 | 52 | torch::Tensor fused_bias_act_op(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer, 53 | int act, int grad, float alpha, float scale) { 54 | int curDevice = -1; 55 | cudaGetDevice(&curDevice); 56 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(curDevice); 57 | 58 | auto x = input.contiguous(); 59 | auto b = bias.contiguous(); 60 | auto ref = refer.contiguous(); 61 | 62 | int use_bias = b.numel() ? 1 : 0; 63 | int use_ref = ref.numel() ? 1 : 0; 64 | 65 | int size_x = x.numel(); 66 | int size_b = b.numel(); 67 | int step_b = 1; 68 | 69 | for (int i = 1 + 1; i < x.dim(); i++) { 70 | step_b *= x.size(i); 71 | } 72 | 73 | int loop_x = 4; 74 | int block_size = 4 * 32; 75 | int grid_size = (size_x - 1) / (loop_x * block_size) + 1; 76 | 77 | auto y = torch::empty_like(x); 78 | 79 | AT_DISPATCH_FLOATING_TYPES_AND_HALF(x.scalar_type(), "fused_bias_act_kernel", [&] { 80 | fused_bias_act_kernel<<>>( 81 | y.data_ptr(), 82 | x.data_ptr(), 83 | b.data_ptr(), 84 | ref.data_ptr(), 85 | act, 86 | grad, 87 | alpha, 88 | scale, 89 | loop_x, 90 | size_x, 91 | step_b, 92 | size_b, 93 | use_bias, 94 | use_ref 95 | ); 96 | }); 97 | 98 | return y; 99 | } -------------------------------------------------------------------------------- /face_model/op/upfirdn2d.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | 4 | torch::Tensor upfirdn2d_op(const torch::Tensor& input, const torch::Tensor& kernel, 5 | int up_x, int up_y, int down_x, int down_y, 6 | int pad_x0, int pad_x1, int pad_y0, int pad_y1); 7 | 8 | #define CHECK_CUDA(x) TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor") 9 | #define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous") 10 | #define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) 11 | 12 | torch::Tensor upfirdn2d(const torch::Tensor& input, const torch::Tensor& kernel, 13 | int up_x, int up_y, int down_x, int down_y, 14 | int pad_x0, int pad_x1, int pad_y0, int pad_y1) { 15 | CHECK_CUDA(input); 16 | CHECK_CUDA(kernel); 17 | 18 | return upfirdn2d_op(input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1); 19 | } 20 | 21 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 22 | m.def("upfirdn2d", &upfirdn2d, "upfirdn2d (CUDA)"); 23 | } -------------------------------------------------------------------------------- /face_model/op/upfirdn2d.py: -------------------------------------------------------------------------------- 1 | import os 2 | import platform 3 | 4 | import torch 5 | import torch.nn.functional as F 6 | from torch.autograd import Function 7 | from torch.utils.cpp_extension import load, _import_module_from_library 8 | 9 | # if running GPEN without cuda, please comment line 10-18 10 | if platform.system() == 'Linux' and torch.cuda.is_available(): 11 | module_path = os.path.dirname(__file__) 12 | upfirdn2d_op = load( 13 | 'upfirdn2d', 14 | sources=[ 15 | os.path.join(module_path, 'upfirdn2d.cpp'), 16 | os.path.join(module_path, 'upfirdn2d_kernel.cu'), 17 | ], 18 | ) 19 | 20 | 21 | #upfirdn2d_op = _import_module_from_library('upfirdn2d', '/tmp/torch_extensions/upfirdn2d', True) 22 | 23 | class UpFirDn2dBackward(Function): 24 | @staticmethod 25 | def forward( 26 | ctx, grad_output, kernel, grad_kernel, up, down, pad, g_pad, in_size, out_size 27 | ): 28 | 29 | up_x, up_y = up 30 | down_x, down_y = down 31 | g_pad_x0, g_pad_x1, g_pad_y0, g_pad_y1 = g_pad 32 | 33 | grad_output = grad_output.reshape(-1, out_size[0], out_size[1], 1) 34 | 35 | grad_input = upfirdn2d_op.upfirdn2d( 36 | grad_output, 37 | grad_kernel, 38 | down_x, 39 | down_y, 40 | up_x, 41 | up_y, 42 | g_pad_x0, 43 | g_pad_x1, 44 | g_pad_y0, 45 | g_pad_y1, 46 | ) 47 | grad_input = grad_input.view(in_size[0], in_size[1], in_size[2], in_size[3]) 48 | 49 | ctx.save_for_backward(kernel) 50 | 51 | pad_x0, pad_x1, pad_y0, pad_y1 = pad 52 | 53 | ctx.up_x = up_x 54 | ctx.up_y = up_y 55 | ctx.down_x = down_x 56 | ctx.down_y = down_y 57 | ctx.pad_x0 = pad_x0 58 | ctx.pad_x1 = pad_x1 59 | ctx.pad_y0 = pad_y0 60 | ctx.pad_y1 = pad_y1 61 | ctx.in_size = in_size 62 | ctx.out_size = out_size 63 | 64 | return grad_input 65 | 66 | @staticmethod 67 | def backward(ctx, gradgrad_input): 68 | kernel, = ctx.saved_tensors 69 | 70 | gradgrad_input = gradgrad_input.reshape(-1, ctx.in_size[2], ctx.in_size[3], 1) 71 | 72 | gradgrad_out = upfirdn2d_op.upfirdn2d( 73 | gradgrad_input, 74 | kernel, 75 | ctx.up_x, 76 | ctx.up_y, 77 | ctx.down_x, 78 | ctx.down_y, 79 | ctx.pad_x0, 80 | ctx.pad_x1, 81 | ctx.pad_y0, 82 | ctx.pad_y1, 83 | ) 84 | # gradgrad_out = gradgrad_out.view(ctx.in_size[0], ctx.out_size[0], ctx.out_size[1], ctx.in_size[3]) 85 | gradgrad_out = gradgrad_out.view( 86 | ctx.in_size[0], ctx.in_size[1], ctx.out_size[0], ctx.out_size[1] 87 | ) 88 | 89 | return gradgrad_out, None, None, None, None, None, None, None, None 90 | 91 | 92 | class UpFirDn2d(Function): 93 | @staticmethod 94 | def forward(ctx, input, kernel, up, down, pad): 95 | up_x, up_y = up 96 | down_x, down_y = down 97 | pad_x0, pad_x1, pad_y0, pad_y1 = pad 98 | 99 | kernel_h, kernel_w = kernel.shape 100 | batch, channel, in_h, in_w = input.shape 101 | ctx.in_size = input.shape 102 | 103 | input = input.reshape(-1, in_h, in_w, 1) 104 | 105 | ctx.save_for_backward(kernel, torch.flip(kernel, [0, 1])) 106 | 107 | out_h = (in_h * up_y + pad_y0 + pad_y1 - kernel_h) // down_y + 1 108 | out_w = (in_w * up_x + pad_x0 + pad_x1 - kernel_w) // down_x + 1 109 | ctx.out_size = (out_h, out_w) 110 | 111 | ctx.up = (up_x, up_y) 112 | ctx.down = (down_x, down_y) 113 | ctx.pad = (pad_x0, pad_x1, pad_y0, pad_y1) 114 | 115 | g_pad_x0 = kernel_w - pad_x0 - 1 116 | g_pad_y0 = kernel_h - pad_y0 - 1 117 | g_pad_x1 = in_w * up_x - out_w * down_x + pad_x0 - up_x + 1 118 | g_pad_y1 = in_h * up_y - out_h * down_y + pad_y0 - up_y + 1 119 | 120 | ctx.g_pad = (g_pad_x0, g_pad_x1, g_pad_y0, g_pad_y1) 121 | 122 | out = upfirdn2d_op.upfirdn2d( 123 | input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1 124 | ) 125 | # out = out.view(major, out_h, out_w, minor) 126 | out = out.view(-1, channel, out_h, out_w) 127 | 128 | return out 129 | 130 | @staticmethod 131 | def backward(ctx, grad_output): 132 | kernel, grad_kernel = ctx.saved_tensors 133 | 134 | grad_input = UpFirDn2dBackward.apply( 135 | grad_output, 136 | kernel, 137 | grad_kernel, 138 | ctx.up, 139 | ctx.down, 140 | ctx.pad, 141 | ctx.g_pad, 142 | ctx.in_size, 143 | ctx.out_size, 144 | ) 145 | 146 | return grad_input, None, None, None, None 147 | 148 | 149 | def upfirdn2d(input, kernel, up=1, down=1, pad=(0, 0), device='cpu'): 150 | if platform.system() == 'Linux' and torch.cuda.is_available() and device != 'cpu': 151 | out = UpFirDn2d.apply( 152 | input, kernel, (up, up), (down, down), (pad[0], pad[1], pad[0], pad[1]) 153 | ) 154 | else: 155 | out = upfirdn2d_native(input, kernel, up, up, down, down, pad[0], pad[1], pad[0], pad[1]) 156 | 157 | return out 158 | 159 | 160 | def upfirdn2d_native( 161 | input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1 162 | ): 163 | input = input.permute(0, 2, 3, 1) 164 | _, in_h, in_w, minor = input.shape 165 | kernel_h, kernel_w = kernel.shape 166 | out = input.view(-1, in_h, 1, in_w, 1, minor) 167 | out = F.pad(out, [0, 0, 0, up_x - 1, 0, 0, 0, up_y - 1]) 168 | out = out.view(-1, in_h * up_y, in_w * up_x, minor) 169 | 170 | out = F.pad( 171 | out, [0, 0, max(pad_x0, 0), max(pad_x1, 0), max(pad_y0, 0), max(pad_y1, 0)] 172 | ) 173 | out = out[ 174 | :, 175 | max(-pad_y0, 0) : out.shape[1] - max(-pad_y1, 0), 176 | max(-pad_x0, 0) : out.shape[2] - max(-pad_x1, 0), 177 | :, 178 | ] 179 | 180 | out = out.permute(0, 3, 1, 2) 181 | out = out.reshape( 182 | [-1, 1, in_h * up_y + pad_y0 + pad_y1, in_w * up_x + pad_x0 + pad_x1] 183 | ) 184 | w = torch.flip(kernel, [0, 1]).view(1, 1, kernel_h, kernel_w) 185 | out = F.conv2d(out, w) 186 | out = out.reshape( 187 | -1, 188 | minor, 189 | in_h * up_y + pad_y0 + pad_y1 - kernel_h + 1, 190 | in_w * up_x + pad_x0 + pad_x1 - kernel_w + 1, 191 | ) 192 | # out = out.permute(0, 2, 3, 1) 193 | return out[:, :, ::down_y, ::down_x] 194 | 195 | -------------------------------------------------------------------------------- /face_parse/blocks.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | import torch 3 | import torch.nn as nn 4 | from torch.nn.parameter import Parameter 5 | from torch.nn import functional as F 6 | import numpy as np 7 | 8 | class NormLayer(nn.Module): 9 | """Normalization Layers. 10 | ------------ 11 | # Arguments 12 | - channels: input channels, for batch norm and instance norm. 13 | - input_size: input shape without batch size, for layer norm. 14 | """ 15 | def __init__(self, channels, normalize_shape=None, norm_type='bn', ref_channels=None): 16 | super(NormLayer, self).__init__() 17 | norm_type = norm_type.lower() 18 | self.norm_type = norm_type 19 | if norm_type == 'bn': 20 | self.norm = nn.BatchNorm2d(channels, affine=True) 21 | elif norm_type == 'in': 22 | self.norm = nn.InstanceNorm2d(channels, affine=False) 23 | elif norm_type == 'gn': 24 | self.norm = nn.GroupNorm(32, channels, affine=True) 25 | elif norm_type == 'pixel': 26 | self.norm = lambda x: F.normalize(x, p=2, dim=1) 27 | elif norm_type == 'layer': 28 | self.norm = nn.LayerNorm(normalize_shape) 29 | elif norm_type == 'none': 30 | self.norm = lambda x: x*1.0 31 | else: 32 | assert 1==0, 'Norm type {} not support.'.format(norm_type) 33 | 34 | def forward(self, x, ref=None): 35 | if self.norm_type == 'spade': 36 | return self.norm(x, ref) 37 | else: 38 | return self.norm(x) 39 | 40 | 41 | class ReluLayer(nn.Module): 42 | """Relu Layer. 43 | ------------ 44 | # Arguments 45 | - relu type: type of relu layer, candidates are 46 | - ReLU 47 | - LeakyReLU: default relu slope 0.2 48 | - PRelu 49 | - SELU 50 | - none: direct pass 51 | """ 52 | def __init__(self, channels, relu_type='relu'): 53 | super(ReluLayer, self).__init__() 54 | relu_type = relu_type.lower() 55 | if relu_type == 'relu': 56 | self.func = nn.ReLU(True) 57 | elif relu_type == 'leakyrelu': 58 | self.func = nn.LeakyReLU(0.2, inplace=True) 59 | elif relu_type == 'prelu': 60 | self.func = nn.PReLU(channels) 61 | elif relu_type == 'selu': 62 | self.func = nn.SELU(True) 63 | elif relu_type == 'none': 64 | self.func = lambda x: x*1.0 65 | else: 66 | assert 1==0, 'Relu type {} not support.'.format(relu_type) 67 | 68 | def forward(self, x): 69 | return self.func(x) 70 | 71 | 72 | class ConvLayer(nn.Module): 73 | def __init__(self, in_channels, out_channels, kernel_size=3, scale='none', norm_type='none', relu_type='none', use_pad=True, bias=True): 74 | super(ConvLayer, self).__init__() 75 | self.use_pad = use_pad 76 | self.norm_type = norm_type 77 | if norm_type in ['bn']: 78 | bias = False 79 | 80 | stride = 2 if scale == 'down' else 1 81 | 82 | self.scale_func = lambda x: x 83 | if scale == 'up': 84 | self.scale_func = lambda x: nn.functional.interpolate(x, scale_factor=2, mode='nearest') 85 | 86 | self.reflection_pad = nn.ReflectionPad2d(int(np.ceil((kernel_size - 1.)/2))) 87 | self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size, stride, bias=bias) 88 | 89 | self.relu = ReluLayer(out_channels, relu_type) 90 | self.norm = NormLayer(out_channels, norm_type=norm_type) 91 | 92 | def forward(self, x): 93 | out = self.scale_func(x) 94 | if self.use_pad: 95 | out = self.reflection_pad(out) 96 | out = self.conv2d(out) 97 | out = self.norm(out) 98 | out = self.relu(out) 99 | return out 100 | 101 | 102 | class ResidualBlock(nn.Module): 103 | """ 104 | Residual block recommended in: http://torch.ch/blog/2016/02/04/resnets.html 105 | """ 106 | def __init__(self, c_in, c_out, relu_type='prelu', norm_type='bn', scale='none'): 107 | super(ResidualBlock, self).__init__() 108 | 109 | if scale == 'none' and c_in == c_out: 110 | self.shortcut_func = lambda x: x 111 | else: 112 | self.shortcut_func = ConvLayer(c_in, c_out, 3, scale) 113 | 114 | scale_config_dict = {'down': ['none', 'down'], 'up': ['up', 'none'], 'none': ['none', 'none']} 115 | scale_conf = scale_config_dict[scale] 116 | 117 | self.conv1 = ConvLayer(c_in, c_out, 3, scale_conf[0], norm_type=norm_type, relu_type=relu_type) 118 | self.conv2 = ConvLayer(c_out, c_out, 3, scale_conf[1], norm_type=norm_type, relu_type='none') 119 | 120 | def forward(self, x): 121 | identity = self.shortcut_func(x) 122 | 123 | res = self.conv1(x) 124 | res = self.conv2(res) 125 | return identity + res 126 | 127 | 128 | -------------------------------------------------------------------------------- /face_parse/face_parsing.py: -------------------------------------------------------------------------------- 1 | ''' 2 | @paper: GAN Prior Embedded Network for Blind Face Restoration in the Wild (CVPR2021) 3 | @author: yangxy (yangtao9009@gmail.com) 4 | ''' 5 | import os 6 | import cv2 7 | import torch 8 | import numpy as np 9 | from parse_model import ParseNet 10 | import torch.nn.functional as F 11 | 12 | class FaceParse(object): 13 | def __init__(self, base_dir='./', model='ParseNet-latest', device='cuda'): 14 | self.mfile = os.path.join(base_dir, 'weights', model+'.pth') 15 | self.size = 512 16 | self.device = device 17 | 18 | ''' 19 | 0: 'background' 1: 'skin' 2: 'nose' 20 | 3: 'eye_g' 4: 'l_eye' 5: 'r_eye' 21 | 6: 'l_brow' 7: 'r_brow' 8: 'l_ear' 22 | 9: 'r_ear' 10: 'mouth' 11: 'u_lip' 23 | 12: 'l_lip' 13: 'hair' 14: 'hat' 24 | 15: 'ear_r' 16: 'neck_l' 17: 'neck' 25 | 18: 'cloth' 26 | ''' 27 | #self.MASK_COLORMAP = [[0, 0, 0], [204, 0, 0], [76, 153, 0], [204, 204, 0], [51, 51, 255], [204, 0, 204], [0, 255, 255], [255, 204, 204], [102, 51, 0], [255, 0, 0], [102, 204, 0], [255, 255, 0], [0, 0, 153], [0, 0, 204], [255, 51, 153], [0, 204, 204], [0, 51, 0], [255, 153, 51], [0, 204, 0]] 28 | #self.#MASK_COLORMAP = [[0, 0, 0], [204, 0, 0], [76, 153, 0], [204, 204, 0], [51, 51, 255], [204, 0, 204], [0, 255, 255], [255, 204, 204], [102, 51, 0], [255, 0, 0], [102, 204, 0], [255, 255, 0], [0, 0, 153], [0, 0, 204], [255, 51, 153], [0, 204, 204], [0, 51, 0], [255, 153, 51], [0, 204, 0]] = [[0, 0, 0], [204, 0, 0], [76, 153, 0], [204, 204, 0], [51, 51, 255], [204, 0, 204], [0, 255, 255], [255, 204, 204], [102, 51, 0], [255, 0, 0], [102, 204, 0], [255, 255, 0], [0, 0, 153], [0, 0, 204], [255, 51, 153], [0, 204, 204], [0, 51, 0], [0, 0, 0], [0, 0, 0]] 29 | self.MASK_COLORMAP = [0, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 0, 255, 255, 255, 0] 30 | self.load_model() 31 | 32 | def load_model(self): 33 | self.faceparse = ParseNet(self.size, self.size, 32, 64, 19, norm_type='bn', relu_type='LeakyReLU', ch_range=[32, 256]) 34 | self.faceparse.load_state_dict(torch.load(self.mfile, map_location=torch.device('cpu'))) 35 | self.faceparse.to(self.device) 36 | self.faceparse.eval() 37 | 38 | def process(self, im): 39 | im = cv2.resize(im, (self.size, self.size)) 40 | imt = self.img2tensor(im) 41 | pred_mask, sr_img_tensor = self.faceparse(imt) 42 | mask = self.tenor2mask(pred_mask) 43 | 44 | return mask 45 | 46 | def process_tensor(self, imt): 47 | imt = F.interpolate(imt.flip(1)*2-1, (self.size, self.size)) 48 | pred_mask, sr_img_tensor = self.faceparse(imt) 49 | 50 | mask = pred_mask.argmax(dim=1) 51 | for idx, color in enumerate(self.MASK_COLORMAP): 52 | mask = torch.where(mask==idx, color, mask) 53 | #mask = mask.repeat(3, 1, 1).unsqueeze(0) #.cpu().float().numpy() 54 | mask = mask.unsqueeze(0) 55 | 56 | return mask 57 | 58 | def img2tensor(self, img): 59 | img = img[..., ::-1] 60 | img = img / 255. * 2 - 1 61 | img_tensor = torch.from_numpy(img.transpose(2, 0, 1)).unsqueeze(0).to(self.device) 62 | return img_tensor.float() 63 | 64 | def tenor2mask(self, tensor): 65 | if len(tensor.shape) < 4: 66 | tensor = tensor.unsqueeze(0) 67 | if tensor.shape[1] > 1: 68 | tensor = tensor.argmax(dim=1) 69 | 70 | tensor = tensor.squeeze(1).data.cpu().numpy() 71 | color_maps = [] 72 | for t in tensor: 73 | #tmp_img = np.zeros(tensor.shape[1:] + (3,)) 74 | tmp_img = np.zeros(tensor.shape[1:]) 75 | for idx, color in enumerate(self.MASK_COLORMAP): 76 | tmp_img[t == idx] = color 77 | color_maps.append(tmp_img.astype(np.uint8)) 78 | return color_maps -------------------------------------------------------------------------------- /face_parse/mask.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/face_parse/mask.png -------------------------------------------------------------------------------- /face_parse/parse_model.py: -------------------------------------------------------------------------------- 1 | ''' 2 | @Created by chaofengc (chaofenghust@gmail.com) 3 | 4 | @Modified by yangxy (yangtao9009@gmail.com) 5 | ''' 6 | 7 | from blocks import * 8 | import torch 9 | from torch import nn 10 | import numpy as np 11 | 12 | def define_P(in_size=512, out_size=512, min_feat_size=32, relu_type='LeakyReLU', isTrain=False, weight_path=None): 13 | net = ParseNet(in_size, out_size, min_feat_size, 64, 19, norm_type='bn', relu_type=relu_type, ch_range=[32, 256]) 14 | if not isTrain: 15 | net.eval() 16 | if weight_path is not None: 17 | net.load_state_dict(torch.load(weight_path)) 18 | return net 19 | 20 | 21 | class ParseNet(nn.Module): 22 | def __init__(self, 23 | in_size=128, 24 | out_size=128, 25 | min_feat_size=32, 26 | base_ch=64, 27 | parsing_ch=19, 28 | res_depth=10, 29 | relu_type='prelu', 30 | norm_type='bn', 31 | ch_range=[32, 512], 32 | ): 33 | super().__init__() 34 | self.res_depth = res_depth 35 | act_args = {'norm_type': norm_type, 'relu_type': relu_type} 36 | min_ch, max_ch = ch_range 37 | 38 | ch_clip = lambda x: max(min_ch, min(x, max_ch)) 39 | min_feat_size = min(in_size, min_feat_size) 40 | 41 | down_steps = int(np.log2(in_size//min_feat_size)) 42 | up_steps = int(np.log2(out_size//min_feat_size)) 43 | 44 | # =============== define encoder-body-decoder ==================== 45 | self.encoder = [] 46 | self.encoder.append(ConvLayer(3, base_ch, 3, 1)) 47 | head_ch = base_ch 48 | for i in range(down_steps): 49 | cin, cout = ch_clip(head_ch), ch_clip(head_ch * 2) 50 | self.encoder.append(ResidualBlock(cin, cout, scale='down', **act_args)) 51 | head_ch = head_ch * 2 52 | 53 | self.body = [] 54 | for i in range(res_depth): 55 | self.body.append(ResidualBlock(ch_clip(head_ch), ch_clip(head_ch), **act_args)) 56 | 57 | self.decoder = [] 58 | for i in range(up_steps): 59 | cin, cout = ch_clip(head_ch), ch_clip(head_ch // 2) 60 | self.decoder.append(ResidualBlock(cin, cout, scale='up', **act_args)) 61 | head_ch = head_ch // 2 62 | 63 | self.encoder = nn.Sequential(*self.encoder) 64 | self.body = nn.Sequential(*self.body) 65 | self.decoder = nn.Sequential(*self.decoder) 66 | self.out_img_conv = ConvLayer(ch_clip(head_ch), 3) 67 | self.out_mask_conv = ConvLayer(ch_clip(head_ch), parsing_ch) 68 | 69 | def forward(self, x): 70 | feat = self.encoder(x) 71 | x = feat + self.body(feat) 72 | x = self.decoder(x) 73 | out_img = self.out_img_conv(x) 74 | out_mask = self.out_mask_conv(x) 75 | return out_mask, out_img 76 | 77 | 78 | -------------------------------------------------------------------------------- /face_parse/test.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/face_parse/test.png -------------------------------------------------------------------------------- /figs/Solvay_conference_1927_comp.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/figs/Solvay_conference_1927_comp.jpg -------------------------------------------------------------------------------- /figs/architecture.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/figs/architecture.png -------------------------------------------------------------------------------- /figs/colorization_00.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/figs/colorization_00.jpg -------------------------------------------------------------------------------- /figs/colorization_01.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/figs/colorization_01.jpg -------------------------------------------------------------------------------- /figs/inpainting_00.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/figs/inpainting_00.jpg -------------------------------------------------------------------------------- /figs/inpainting_01.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/figs/inpainting_01.jpg -------------------------------------------------------------------------------- /figs/real_00.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/figs/real_00.png -------------------------------------------------------------------------------- /figs/real_01.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/figs/real_01.png -------------------------------------------------------------------------------- /figs/real_02.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/figs/real_02.png -------------------------------------------------------------------------------- /figs/real_03.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/figs/real_03.png -------------------------------------------------------------------------------- /figs/seg2face_00.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/figs/seg2face_00.jpg -------------------------------------------------------------------------------- /figs/seg2face_01.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/figs/seg2face_01.jpg -------------------------------------------------------------------------------- /figs/selfie_00.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/figs/selfie_00.jpg -------------------------------------------------------------------------------- /figs/selfie_01.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/figs/selfie_01.jpg -------------------------------------------------------------------------------- /misc/cog.yaml: -------------------------------------------------------------------------------- 1 | build: 2 | gpu: true 3 | python_version: "3.8" 4 | system_packages: 5 | - "libgl1-mesa-glx" 6 | - "libglib2.0-0" 7 | - "ninja-build" 8 | python_packages: 9 | - "torch==1.7.1" 10 | - "torchvision==0.8.2" 11 | - "numpy==1.20.1" 12 | - "ipython==7.21.0" 13 | - "Pillow==8.3.1" 14 | - "scikit-image==0.18.3" 15 | - "opencv-python==4.5.3.56" 16 | 17 | predict: "predict.py:Predictor" 18 | -------------------------------------------------------------------------------- /misc/onnx_export.py: -------------------------------------------------------------------------------- 1 | import math 2 | import os 3 | import sys 4 | 5 | import torch 6 | import torch.nn as nn 7 | from tqdm import tqdm 8 | from torch.autograd import Variable 9 | 10 | import __init_paths 11 | #from face_model import model 12 | from face_model.gpen_model import FullGenerator 13 | 14 | def model_load(model, path): 15 | """Load model.""" 16 | 17 | if not os.path.exists(path): 18 | print("Model '{}' does not exist.".format(path)) 19 | return 20 | 21 | state_dict = torch.load(path, map_location=lambda storage, loc: storage) 22 | 23 | model.load_state_dict(state_dict) 24 | 25 | 26 | def export_onnx(model, path, force_cpu): 27 | """Export onnx model.""" 28 | 29 | import onnx 30 | import onnxruntime 31 | #from onnx import optimizer 32 | import numpy as np 33 | 34 | onnx_file_name = os.path.join(path, model+".onnx") 35 | model_weight_file = os.path.join(path, model+".pth") 36 | dummy_input = Variable(torch.randn(1, 3, 1024, 1024)) 37 | 38 | # 1. Create and load model. 39 | model_setenv(force_cpu) 40 | torch_model = get_model(model_weight_file) 41 | torch_model.eval() 42 | 43 | # 2. Model export 44 | print("Export model ...") 45 | 46 | input_names = ["input"] 47 | output_names = ["output"] 48 | device = model_device() 49 | # torch.onnx.export(torch_model, dummy_input.to(device), onnx_file_name, 50 | # input_names=input_names, 51 | # output_names=output_names, 52 | # verbose=False, 53 | # opset_version=12, 54 | # keep_initializers_as_inputs=False, 55 | # export_params=True, 56 | # operator_export_type=torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK) 57 | torch.onnx.export(torch_model, dummy_input.to(device), onnx_file_name, 58 | input_names=input_names, 59 | output_names=output_names, 60 | verbose=False, 61 | opset_version=10, 62 | operator_export_type=torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK) 63 | 64 | # 3. Optimize model 65 | print('Checking model ...') 66 | onnx_model = onnx.load(onnx_file_name) 67 | onnx.checker.check_model(onnx_model) 68 | # https://github.com/onnx/optimizer 69 | print('Done checking model ...') 70 | # 4. Visual model 71 | # python -c "import netron; netron.start('output/image_zoom.onnx')" 72 | 73 | def verify_onnx(model, path, force_cpu): 74 | """Verify onnx model.""" 75 | 76 | import onnxruntime 77 | import numpy as np 78 | 79 | 80 | model_weight_file = os.path.join(path, model+".pth") 81 | 82 | model_weight_file = "./weights/GPEN-512.pth" 83 | 84 | model_setenv(force_cpu) 85 | torch_model = get_model(model_weight_file) 86 | torch_model.eval() 87 | 88 | onnx_file_name = os.path.join(path, model+".onnx") 89 | onnxruntime_engine = onnxruntime.InferenceSession(onnx_file_name) 90 | 91 | def to_numpy(tensor): 92 | return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy() 93 | 94 | dummy_input = Variable(torch.randn(1, 3, 512, 512)) 95 | with torch.no_grad(): 96 | torch_output, _ = torch_model(dummy_input) 97 | onnxruntime_inputs = {onnxruntime_engine.get_inputs()[0].name: to_numpy(dummy_input)} 98 | onnxruntime_outputs = onnxruntime_engine.run(None, onnxruntime_inputs) 99 | np.testing.assert_allclose(to_numpy(torch_output), onnxruntime_outputs[0], rtol=1e-02, atol=1e-02) 100 | print("Example: Onnx model has been tested with ONNXRuntime, the result looks good !") 101 | 102 | def get_model(checkpoint): 103 | """Create encoder model.""" 104 | 105 | #model_setenv() 106 | model = FullGenerator(1024, 512, 8, 2, narrow=1) #TODO 107 | model_load(model, checkpoint) 108 | device = model_device() 109 | model.to(device) 110 | return model 111 | 112 | 113 | def model_device(): 114 | """Please call after model_setenv. """ 115 | 116 | return torch.device(os.environ["DEVICE"]) 117 | 118 | 119 | def model_setenv(cpu_only): 120 | """Setup environ ...""" 121 | 122 | # random init ... 123 | import random 124 | random.seed(42) 125 | torch.manual_seed(42) 126 | 127 | # Set default device to avoid exceptions 128 | if cpu_only: 129 | os.environ["DEVICE"] = 'cpu' 130 | else: 131 | if os.environ.get("DEVICE") != "cuda" and os.environ.get("DEVICE") != "cpu": 132 | os.environ["DEVICE"] = 'cuda' if torch.cuda.is_available() else 'cpu' 133 | if os.environ["DEVICE"] == 'cuda': 134 | torch.backends.cudnn.enabled = True 135 | torch.backends.cudnn.benchmark = True 136 | 137 | print("Running Environment:") 138 | print("----------------------------------------------") 139 | #print(" PWD: ", os.environ["PWD"]) 140 | print(" DEVICE: ", os.environ["DEVICE"]) 141 | 142 | # def export_torch(model, path): 143 | # """Export torch model.""" 144 | 145 | # script_file = os.path.join(path, model+".pt") 146 | # weight_file = os.path.join(path, model+".onnx") 147 | 148 | # # 1. Load model 149 | # print("Loading model ...") 150 | # model = get_model(weight_file) 151 | # model.eval() 152 | 153 | # # 2. Model export 154 | # print("Export model ...") 155 | # dummy_input = Variable(torch.randn(1, 3, 512, 512)) 156 | # device = model_device() 157 | # traced_script_module = torch.jit.trace(model, dummy_input.to(device), _force_outplace=True) 158 | # traced_script_module.save(script_file) 159 | 160 | 161 | if __name__ == '__main__': 162 | """Test model ...""" 163 | import argparse 164 | 165 | parser = argparse.ArgumentParser() 166 | parser.add_argument('--model', type=str, required=True) 167 | parser.add_argument('--path', type=str, default='./') 168 | parser.add_argument('--export', help="Export onnx model", action='store_true') 169 | parser.add_argument('--verify', help="Verify onnx model", action='store_true') 170 | parser.add_argument('--force-cpu', dest='force_cpu', help="Verify onnx model", action='store_true') 171 | 172 | args = parser.parse_args() 173 | 174 | # export_torch() 175 | 176 | 177 | 178 | if args.export: 179 | export_onnx(model = args.model, path = args.path, force_cpu=args.force_cpu) 180 | 181 | if args.verify: 182 | verify_onnx(model = args.model, path = args.path, force_cpu=args.force_cpu) 183 | -------------------------------------------------------------------------------- /misc/predict.py: -------------------------------------------------------------------------------- 1 | import os 2 | import tempfile 3 | from pathlib import Path 4 | import glob 5 | import numpy as np 6 | import cv2 7 | from zipfile import ZipFile 8 | from PIL import Image 9 | import shutil 10 | import cog 11 | from face_enhancement import FaceEnhancement 12 | from face_colorization import FaceColorization 13 | from face_inpainting import FaceInpainting, brush_stroke_mask 14 | 15 | 16 | class Predictor(cog.Predictor): 17 | def setup(self): 18 | faceenhancer_model = {'name': 'GPEN-BFR-256', 'size': 256, 'channel_multiplier': 1, 'narrow': 0.5} 19 | self.faceenhancer = FaceEnhancement(size=faceenhancer_model['size'], model=faceenhancer_model['name'], 20 | channel_multiplier=faceenhancer_model['channel_multiplier'], 21 | narrow=faceenhancer_model['narrow']) 22 | faceinpainter_model = {'name': 'GPEN-Inpainting-1024', 'size': 1024} 23 | self.faceinpainter = FaceInpainting(size=faceinpainter_model['size'], model=faceinpainter_model['name'], 24 | channel_multiplier=2) 25 | facecolorizer_model = {'name': 'GPEN-Colorization-1024', 'size': 1024} 26 | self.facecolorizer = FaceColorization(size=facecolorizer_model['size'], model=facecolorizer_model['name'], 27 | channel_multiplier=2) 28 | 29 | @cog.input( 30 | "image", 31 | type=Path, 32 | help="input image", 33 | ) 34 | @cog.input( 35 | "task", 36 | type=str, 37 | options=['Face Restoration', 'Face Colorization', 'Face Inpainting'], 38 | default='Face Restoration', 39 | help="choose task type" 40 | ) 41 | @cog.input( 42 | "output_individual", 43 | type=bool, 44 | default=False, 45 | help="whether outputs individual enhanced faces, valid for Face Restoration. When set to true, a zip folder of " 46 | "all the enhanced faces in the input will be generated for download." 47 | ) 48 | @cog.input( 49 | "broken_image", 50 | type=bool, 51 | default=True, 52 | help="whether the input image is broken, valid for Face Inpainting. When set to True, the output will be the " 53 | "'fixed' image. When set to False, the image will randomly add brush strokes to simulate a broken image, " 54 | "and the output will be broken + fixed image" 55 | ) 56 | def predict(self, image, task='Face Restoration', output_individual=False, broken_image=True): 57 | out_path = Path(tempfile.mkdtemp()) / "out.png" 58 | if task == 'Face Restoration': 59 | im = cv2.imread(str(image), cv2.IMREAD_COLOR) # BGR 60 | assert isinstance(im, np.ndarray), 'input filename error' 61 | im = cv2.resize(im, (0, 0), fx=2, fy=2) 62 | img, orig_faces, enhanced_faces = self.faceenhancer.process(im) 63 | cv2.imwrite(str(out_path), img) 64 | if output_individual: 65 | zip_folder = 'out_zip' 66 | os.makedirs(zip_folder, exist_ok=True) 67 | out_path = Path(tempfile.mkdtemp()) / "out.zip" 68 | try: 69 | cv2.imwrite(os.path.join(zip_folder, 'whole_image.jpg'), img) 70 | for m, ef in enumerate(enhanced_faces): 71 | cv2.imwrite(os.path.join(zip_folder, f'face_{m}.jpg'), ef) 72 | img_list = sorted(glob.glob(os.path.join(zip_folder, '*'))) 73 | with ZipFile(str(out_path), 'w') as zipfile: 74 | for img in img_list: 75 | zipfile.write(img) 76 | finally: 77 | clean_folder(zip_folder) 78 | elif task == 'Face Colorization': 79 | grayf = cv2.imread(str(image), cv2.IMREAD_GRAYSCALE) 80 | grayf = cv2.cvtColor(grayf, cv2.COLOR_GRAY2BGR) # channel: 1->3 81 | colorf = self.facecolorizer.process(grayf) 82 | cv2.imwrite(str(out_path), colorf) 83 | else: 84 | originf = cv2.imread(str(image), cv2.IMREAD_COLOR) 85 | brokenf = originf 86 | if not broken_image: 87 | brokenf = np.asarray(brush_stroke_mask(Image.fromarray(originf))) 88 | completef = self.faceinpainter.process(brokenf) 89 | brokenf = cv2.resize(brokenf, completef.shape[:2]) 90 | out_img = completef if broken_image else np.hstack((brokenf, completef)) 91 | cv2.imwrite(str(out_path), out_img) 92 | 93 | return out_path 94 | 95 | 96 | def clean_folder(folder): 97 | for filename in os.listdir(folder): 98 | file_path = os.path.join(folder, filename) 99 | try: 100 | if os.path.isfile(file_path) or os.path.islink(file_path): 101 | os.unlink(file_path) 102 | elif os.path.isdir(file_path): 103 | shutil.rmtree(file_path) 104 | except Exception as e: 105 | print('Failed to delete %s. Reason: %s' % (file_path, e)) -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | ninja 2 | torch 3 | torchvision 4 | opencv-python 5 | numpy 6 | scikit-image 7 | scipy 8 | pillow 9 | tqdm -------------------------------------------------------------------------------- /run.sh: -------------------------------------------------------------------------------- 1 | ############## TEST ################ 2 | #python demo.py --task FaceEnhancement --model GPEN-BFR-512 --in_size 512 --channel_multiplier 2 --narrow 1 --use_sr --sr_scale 2 --use_cuda --save_face --indir examples/imgs --outdir examples/outs-bfr 3 | #python demo.py --task FaceEnhancement --model GPEN-BFR-256 --in_size 256 --channel_multiplier 1 --narrow 0.5 --use_sr --sr_scale 4 --use_cuda --save_face --indir examples/imgs --outdir examples/outs-bfr 4 | python demo.py --task FaceEnhancement --model GPEN-BFR-2048 --in_size 2048 --channel_multiplier 2 --narrow 1 --alpha 0.8 --use_sr --sr_scale 2 --use_cuda --tile_size 400 --save_face --indir examples/selfie --outdir examples/outs-selfie 5 | 6 | #python demo.py --task FaceColorization --model GPEN-Colorization-1024 --in_size 1024 --use_cuda --indir examples/tmp --outdir examples/outs-colorization 7 | 8 | #python demo.py --task FaceInpainting --model GPEN-Inpainting-1024 --in_size 1024 --use_cuda --indir examples/ffhq-10 --outdir examples/outs-inpainting 9 | 10 | #python demo.py --task Segmentation2Face --model GPEN-Seg2face-512 --in_size 512 --use_cuda --indir examples/segs --outdir examples/outs-seg2face 11 | 12 | ############## TRAIN ################ 13 | #CUDA_VISIBLE_DEVICES='0,1,' python -m torch.distributed.launch --nproc_per_node=2 --master_port=4321 train_simple.py --size 512 --channel_multiplier 2 --narrow 1 --ckpt ckpt-512 --sample sample-512 --batch 1 --path your_path_of_aligned_faces 14 | -------------------------------------------------------------------------------- /segmentation2face.py: -------------------------------------------------------------------------------- 1 | ''' 2 | @paper: GAN Prior Embedded Network for Blind Face Restoration in the Wild (CVPR2021) 3 | @author: yangxy (yangtao9009@gmail.com) 4 | ''' 5 | from face_model.face_gan import FaceGAN 6 | 7 | class Segmentation2Face(object): 8 | def __init__(self, base_dir='./', in_size=1024, out_size=None, model=None, channel_multiplier=2, narrow=1, key=None, is_norm=True, device='cuda'): 9 | self.facegan = FaceGAN(base_dir, in_size, out_size, model, channel_multiplier, narrow, key, is_norm, device=device) 10 | 11 | # make sure the face image is well aligned. Please refer to face_enhancement.py 12 | def process(self, segf, aligned=True): 13 | # from segmentations to faces 14 | out = self.facegan.process(segf) 15 | 16 | return out, [segf], [out] 17 | 18 | 19 | -------------------------------------------------------------------------------- /sr_model/arch_util.py: -------------------------------------------------------------------------------- 1 | import math 2 | import torch 3 | from torch import nn as nn 4 | from torch.nn import functional as F 5 | from torch.nn import init as init 6 | from torch.nn.modules.batchnorm import _BatchNorm 7 | 8 | @torch.no_grad() 9 | def default_init_weights(module_list, scale=1, bias_fill=0, **kwargs): 10 | """Initialize network weights. 11 | 12 | Args: 13 | module_list (list[nn.Module] | nn.Module): Modules to be initialized. 14 | scale (float): Scale initialized weights, especially for residual 15 | blocks. Default: 1. 16 | bias_fill (float): The value to fill bias. Default: 0 17 | kwargs (dict): Other arguments for initialization function. 18 | """ 19 | if not isinstance(module_list, list): 20 | module_list = [module_list] 21 | for module in module_list: 22 | for m in module.modules(): 23 | if isinstance(m, nn.Conv2d): 24 | init.kaiming_normal_(m.weight, **kwargs) 25 | m.weight.data *= scale 26 | if m.bias is not None: 27 | m.bias.data.fill_(bias_fill) 28 | elif isinstance(m, nn.Linear): 29 | init.kaiming_normal_(m.weight, **kwargs) 30 | m.weight.data *= scale 31 | if m.bias is not None: 32 | m.bias.data.fill_(bias_fill) 33 | elif isinstance(m, _BatchNorm): 34 | init.constant_(m.weight, 1) 35 | if m.bias is not None: 36 | m.bias.data.fill_(bias_fill) 37 | 38 | 39 | def make_layer(basic_block, num_basic_block, **kwarg): 40 | """Make layers by stacking the same blocks. 41 | 42 | Args: 43 | basic_block (nn.module): nn.module class for basic block. 44 | num_basic_block (int): number of blocks. 45 | 46 | Returns: 47 | nn.Sequential: Stacked blocks in nn.Sequential. 48 | """ 49 | layers = [] 50 | for _ in range(num_basic_block): 51 | layers.append(basic_block(**kwarg)) 52 | return nn.Sequential(*layers) 53 | 54 | 55 | class ResidualBlockNoBN(nn.Module): 56 | """Residual block without BN. 57 | 58 | It has a style of: 59 | ---Conv-ReLU-Conv-+- 60 | |________________| 61 | 62 | Args: 63 | num_feat (int): Channel number of intermediate features. 64 | Default: 64. 65 | res_scale (float): Residual scale. Default: 1. 66 | pytorch_init (bool): If set to True, use pytorch default init, 67 | otherwise, use default_init_weights. Default: False. 68 | """ 69 | 70 | def __init__(self, num_feat=64, res_scale=1, pytorch_init=False): 71 | super(ResidualBlockNoBN, self).__init__() 72 | self.res_scale = res_scale 73 | self.conv1 = nn.Conv2d(num_feat, num_feat, 3, 1, 1, bias=True) 74 | self.conv2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1, bias=True) 75 | self.relu = nn.ReLU(inplace=True) 76 | 77 | if not pytorch_init: 78 | default_init_weights([self.conv1, self.conv2], 0.1) 79 | 80 | def forward(self, x): 81 | identity = x 82 | out = self.conv2(self.relu(self.conv1(x))) 83 | return identity + out * self.res_scale 84 | 85 | 86 | class Upsample(nn.Sequential): 87 | """Upsample module. 88 | 89 | Args: 90 | scale (int): Scale factor. Supported scales: 2^n and 3. 91 | num_feat (int): Channel number of intermediate features. 92 | """ 93 | 94 | def __init__(self, scale, num_feat): 95 | m = [] 96 | if (scale & (scale - 1)) == 0: # scale = 2^n 97 | for _ in range(int(math.log(scale, 2))): 98 | m.append(nn.Conv2d(num_feat, 4 * num_feat, 3, 1, 1)) 99 | m.append(nn.PixelShuffle(2)) 100 | elif scale == 3: 101 | m.append(nn.Conv2d(num_feat, 9 * num_feat, 3, 1, 1)) 102 | m.append(nn.PixelShuffle(3)) 103 | else: 104 | raise ValueError(f'scale {scale} is not supported. ' 105 | 'Supported scales: 2^n and 3.') 106 | super(Upsample, self).__init__(*m) 107 | 108 | # TODO: may write a cpp file 109 | def pixel_unshuffle(x, scale): 110 | """ Pixel unshuffle. 111 | 112 | Args: 113 | x (Tensor): Input feature with shape (b, c, hh, hw). 114 | scale (int): Downsample ratio. 115 | 116 | Returns: 117 | Tensor: the pixel unshuffled feature. 118 | """ 119 | b, c, hh, hw = x.size() 120 | out_channel = c * (scale**2) 121 | assert hh % scale == 0 and hw % scale == 0 122 | h = hh // scale 123 | w = hw // scale 124 | x_view = x.view(b, c, h, scale, w, scale) 125 | return x_view.permute(0, 1, 3, 5, 2, 4).reshape(b, out_channel, h, w) -------------------------------------------------------------------------------- /sr_model/real_esrnet.py: -------------------------------------------------------------------------------- 1 | import os 2 | import math 3 | import torch 4 | import numpy as np 5 | from rrdbnet_arch import RRDBNet 6 | from torch.nn import functional as F 7 | 8 | class RealESRNet(object): 9 | def __init__(self, base_dir='./', model=None, scale=2, tile_size=0, tile_pad=10, device='cuda'): 10 | self.base_dir = base_dir 11 | self.scale = scale 12 | self.tile_size = tile_size 13 | self.tile_pad = tile_pad 14 | self.device = device 15 | self.load_srmodel(base_dir, model) 16 | 17 | def load_srmodel(self, base_dir, model): 18 | self.srmodel = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=32, num_block=23, num_grow_ch=32, scale=self.scale) 19 | if model is None: 20 | loadnet = torch.load(os.path.join(self.base_dir, 'weights', 'realesrnet_x%d.pth'%self.scale)) 21 | else: 22 | loadnet = torch.load(os.path.join(self.base_dir, 'weights', model+'_x%d.pth'%self.scale)) 23 | #print(loadnet['params_ema'].keys) 24 | self.srmodel.load_state_dict(loadnet['params_ema'], strict=True) 25 | self.srmodel.eval() 26 | self.srmodel = self.srmodel.to(self.device) 27 | 28 | def tile_process(self, img): 29 | """It will first crop input images to tiles, and then process each tile. 30 | Finally, all the processed tiles are merged into one images. 31 | 32 | Modified from: https://github.com/ata4/esrgan-launcher 33 | """ 34 | batch, channel, height, width = img.shape 35 | output_height = height * self.scale 36 | output_width = width * self.scale 37 | output_shape = (batch, channel, output_height, output_width) 38 | 39 | # start with black image 40 | output = img.new_zeros(output_shape) 41 | tiles_x = math.ceil(width / self.tile_size) 42 | tiles_y = math.ceil(height / self.tile_size) 43 | 44 | # loop over all tiles 45 | for y in range(tiles_y): 46 | for x in range(tiles_x): 47 | # extract tile from input image 48 | ofs_x = x * self.tile_size 49 | ofs_y = y * self.tile_size 50 | # input tile area on total image 51 | input_start_x = ofs_x 52 | input_end_x = min(ofs_x + self.tile_size, width) 53 | input_start_y = ofs_y 54 | input_end_y = min(ofs_y + self.tile_size, height) 55 | 56 | # input tile area on total image with padding 57 | input_start_x_pad = max(input_start_x - self.tile_pad, 0) 58 | input_end_x_pad = min(input_end_x + self.tile_pad, width) 59 | input_start_y_pad = max(input_start_y - self.tile_pad, 0) 60 | input_end_y_pad = min(input_end_y + self.tile_pad, height) 61 | 62 | # input tile dimensions 63 | input_tile_width = input_end_x - input_start_x 64 | input_tile_height = input_end_y - input_start_y 65 | tile_idx = y * tiles_x + x + 1 66 | input_tile = img[:, :, input_start_y_pad:input_end_y_pad, input_start_x_pad:input_end_x_pad] 67 | 68 | # upscale tile 69 | try: 70 | with torch.no_grad(): 71 | output_tile = self.srmodel(input_tile) 72 | except RuntimeError as error: 73 | print('Error', error) 74 | return None 75 | if tile_idx%10==0: print(f'\tTile {tile_idx}/{tiles_x * tiles_y}') 76 | 77 | # output tile area on total image 78 | output_start_x = input_start_x * self.scale 79 | output_end_x = input_end_x * self.scale 80 | output_start_y = input_start_y * self.scale 81 | output_end_y = input_end_y * self.scale 82 | 83 | # output tile area without padding 84 | output_start_x_tile = (input_start_x - input_start_x_pad) * self.scale 85 | output_end_x_tile = output_start_x_tile + input_tile_width * self.scale 86 | output_start_y_tile = (input_start_y - input_start_y_pad) * self.scale 87 | output_end_y_tile = output_start_y_tile + input_tile_height * self.scale 88 | 89 | # put tile into output image 90 | output[:, :, output_start_y:output_end_y, 91 | output_start_x:output_end_x] = output_tile[:, :, output_start_y_tile:output_end_y_tile, 92 | output_start_x_tile:output_end_x_tile] 93 | return output 94 | 95 | def process(self, img): 96 | img = img.astype(np.float32) / 255. 97 | img = torch.from_numpy(np.transpose(img[:, :, [2, 1, 0]], (2, 0, 1))).float() 98 | img = img.unsqueeze(0).to(self.device) 99 | 100 | if self.scale == 2: 101 | mod_scale = 2 102 | elif self.scale == 1: 103 | mod_scale = 4 104 | else: 105 | mod_scale = None 106 | if mod_scale is not None: 107 | h_pad, w_pad = 0, 0 108 | _, _, h, w = img.size() 109 | if (h % mod_scale != 0): 110 | h_pad = (mod_scale - h % mod_scale) 111 | if (w % mod_scale != 0): 112 | w_pad = (mod_scale - w % mod_scale) 113 | img = F.pad(img, (0, w_pad, 0, h_pad), 'reflect') 114 | 115 | try: 116 | with torch.no_grad(): 117 | if self.tile_size > 0: 118 | output = self.tile_process(img) 119 | else: 120 | output = self.srmodel(img) 121 | del img 122 | # remove extra pad 123 | if mod_scale is not None: 124 | _, _, h, w = output.size() 125 | output = output[:, :, 0:h - h_pad, 0:w - w_pad] 126 | output = output.data.squeeze().float().cpu().clamp_(0, 1).numpy() 127 | output = np.transpose(output[[2, 1, 0], :, :], (1, 2, 0)) 128 | output = (output * 255.0).round().astype(np.uint8) 129 | 130 | return output 131 | except Exception as e: 132 | print('sr failed:', e) 133 | return None -------------------------------------------------------------------------------- /sr_model/rrdbnet_arch.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn as nn 3 | from torch.nn import functional as F 4 | 5 | from arch_util import default_init_weights, make_layer, pixel_unshuffle 6 | 7 | 8 | class ResidualDenseBlock(nn.Module): 9 | """Residual Dense Block. 10 | 11 | Used in RRDB block in ESRGAN. 12 | 13 | Args: 14 | num_feat (int): Channel number of intermediate features. 15 | num_grow_ch (int): Channels for each growth. 16 | """ 17 | 18 | def __init__(self, num_feat=64, num_grow_ch=32): 19 | super(ResidualDenseBlock, self).__init__() 20 | self.conv1 = nn.Conv2d(num_feat, num_grow_ch, 3, 1, 1) 21 | self.conv2 = nn.Conv2d(num_feat + num_grow_ch, num_grow_ch, 3, 1, 1) 22 | self.conv3 = nn.Conv2d(num_feat + 2 * num_grow_ch, num_grow_ch, 3, 1, 1) 23 | self.conv4 = nn.Conv2d(num_feat + 3 * num_grow_ch, num_grow_ch, 3, 1, 1) 24 | self.conv5 = nn.Conv2d(num_feat + 4 * num_grow_ch, num_feat, 3, 1, 1) 25 | 26 | self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True) 27 | 28 | # initialization 29 | default_init_weights([self.conv1, self.conv2, self.conv3, self.conv4, self.conv5], 0.1) 30 | 31 | def forward(self, x): 32 | x1 = self.lrelu(self.conv1(x)) 33 | x2 = self.lrelu(self.conv2(torch.cat((x, x1), 1))) 34 | x3 = self.lrelu(self.conv3(torch.cat((x, x1, x2), 1))) 35 | x4 = self.lrelu(self.conv4(torch.cat((x, x1, x2, x3), 1))) 36 | x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1)) 37 | # Emperically, we use 0.2 to scale the residual for better performance 38 | return x5 * 0.2 + x 39 | 40 | 41 | class RRDB(nn.Module): 42 | """Residual in Residual Dense Block. 43 | 44 | Used in RRDB-Net in ESRGAN. 45 | 46 | Args: 47 | num_feat (int): Channel number of intermediate features. 48 | num_grow_ch (int): Channels for each growth. 49 | """ 50 | 51 | def __init__(self, num_feat, num_grow_ch=32): 52 | super(RRDB, self).__init__() 53 | self.rdb1 = ResidualDenseBlock(num_feat, num_grow_ch) 54 | self.rdb2 = ResidualDenseBlock(num_feat, num_grow_ch) 55 | self.rdb3 = ResidualDenseBlock(num_feat, num_grow_ch) 56 | 57 | def forward(self, x): 58 | out = self.rdb1(x) 59 | out = self.rdb2(out) 60 | out = self.rdb3(out) 61 | # Emperically, we use 0.2 to scale the residual for better performance 62 | return out * 0.2 + x 63 | 64 | class RRDBNet(nn.Module): 65 | """Networks consisting of Residual in Residual Dense Block, which is used 66 | in ESRGAN. 67 | 68 | ESRGAN: Enhanced Super-Resolution Generative Adversarial Networks. 69 | 70 | We extend ESRGAN for scale x2 and scale x1. 71 | Note: This is one option for scale 1, scale 2 in RRDBNet. 72 | We first employ the pixel-unshuffle (an inverse operation of pixelshuffle to reduce the spatial size 73 | and enlarge the channel size before feeding inputs into the main ESRGAN architecture. 74 | 75 | Args: 76 | num_in_ch (int): Channel number of inputs. 77 | num_out_ch (int): Channel number of outputs. 78 | num_feat (int): Channel number of intermediate features. 79 | Default: 64 80 | num_block (int): Block number in the trunk network. Defaults: 23 81 | num_grow_ch (int): Channels for each growth. Default: 32. 82 | """ 83 | 84 | def __init__(self, num_in_ch, num_out_ch, scale=4, num_feat=64, num_block=23, num_grow_ch=32): 85 | super(RRDBNet, self).__init__() 86 | self.scale = scale 87 | if scale == 2: 88 | num_in_ch = num_in_ch * 4 89 | elif scale == 1: 90 | num_in_ch = num_in_ch * 16 91 | self.conv_first = nn.Conv2d(num_in_ch, num_feat, 3, 1, 1) 92 | self.body = make_layer(RRDB, num_block, num_feat=num_feat, num_grow_ch=num_grow_ch) 93 | self.conv_body = nn.Conv2d(num_feat, num_feat, 3, 1, 1) 94 | # upsample 95 | self.conv_up1 = nn.Conv2d(num_feat, num_feat, 3, 1, 1) 96 | self.conv_up2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1) 97 | self.conv_hr = nn.Conv2d(num_feat, num_feat, 3, 1, 1) 98 | self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1) 99 | 100 | self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True) 101 | 102 | def forward(self, x): 103 | if self.scale == 2: 104 | feat = pixel_unshuffle(x, scale=2) 105 | elif self.scale == 1: 106 | feat = pixel_unshuffle(x, scale=4) 107 | else: 108 | feat = x 109 | feat = self.conv_first(feat) 110 | body_feat = self.conv_body(self.body(feat)) 111 | feat = feat + body_feat 112 | # upsample 113 | feat = self.lrelu(self.conv_up1(F.interpolate(feat, scale_factor=2, mode='nearest'))) 114 | feat = self.lrelu(self.conv_up2(F.interpolate(feat, scale_factor=2, mode='nearest'))) 115 | out = self.conv_last(self.lrelu(self.conv_hr(feat))) 116 | return out 117 | -------------------------------------------------------------------------------- /training/data_loader/dataset_face.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import cv2 3 | import os 4 | import glob 5 | import math 6 | import random 7 | import torch 8 | import torch.nn.functional as F 9 | from torch.utils.data import Dataset 10 | 11 | import degradations 12 | 13 | 14 | class GFPGAN_degradation(object): 15 | def __init__(self): 16 | self.kernel_list = ['iso', 'aniso'] 17 | self.kernel_prob = [0.5, 0.5] 18 | self.blur_kernel_size = 41 19 | self.blur_sigma = [0.1, 10] 20 | self.downsample_range = [0.8, 8] 21 | self.noise_range = [0, 20] 22 | self.jpeg_range = [60, 100] 23 | self.gray_prob = 0.2 24 | self.color_jitter_prob = 0.0 25 | self.color_jitter_pt_prob = 0.0 26 | self.shift = 20/255. 27 | 28 | def degrade_process(self, img_gt): 29 | if random.random() > 0.5: 30 | img_gt = cv2.flip(img_gt, 1) 31 | 32 | h, w = img_gt.shape[:2] 33 | 34 | # random color jitter 35 | if np.random.uniform() < self.color_jitter_prob: 36 | jitter_val = np.random.uniform(-self.shift, self.shift, 3).astype(np.float32) 37 | img_gt = img_gt + jitter_val 38 | img_gt = np.clip(img_gt, 0, 1) 39 | 40 | # random grayscale 41 | if np.random.uniform() < self.gray_prob: 42 | img_gt = cv2.cvtColor(img_gt, cv2.COLOR_BGR2GRAY) 43 | img_gt = np.tile(img_gt[:, :, None], [1, 1, 3]) 44 | 45 | # ------------------------ generate lq image ------------------------ # 46 | # blur 47 | kernel = degradations.random_mixed_kernels( 48 | self.kernel_list, 49 | self.kernel_prob, 50 | self.blur_kernel_size, 51 | self.blur_sigma, 52 | self.blur_sigma, [-math.pi, math.pi], 53 | noise_range=None) 54 | img_lq = cv2.filter2D(img_gt, -1, kernel) 55 | # downsample 56 | scale = np.random.uniform(self.downsample_range[0], self.downsample_range[1]) 57 | img_lq = cv2.resize(img_lq, (int(w // scale), int(h // scale)), interpolation=cv2.INTER_LINEAR) 58 | 59 | # noise 60 | if self.noise_range is not None: 61 | img_lq = degradations.random_add_gaussian_noise(img_lq, self.noise_range) 62 | # jpeg compression 63 | if self.jpeg_range is not None: 64 | img_lq = degradations.random_add_jpg_compression(img_lq, self.jpeg_range) 65 | 66 | # round and clip 67 | img_lq = np.clip((img_lq * 255.0).round(), 0, 255) / 255. 68 | 69 | # resize to original size 70 | img_lq = cv2.resize(img_lq, (w, h), interpolation=cv2.INTER_LINEAR) 71 | 72 | return img_gt, img_lq 73 | 74 | class FaceDataset(Dataset): 75 | def __init__(self, path, resolution=512): 76 | self.resolution = resolution 77 | 78 | self.HQ_imgs = glob.glob(os.path.join(path, '*.*')) 79 | self.length = len(self.HQ_imgs) 80 | 81 | self.degrader = GFPGAN_degradation() 82 | 83 | def __len__(self): 84 | return self.length 85 | 86 | def __getitem__(self, index): 87 | img_gt = cv2.imread(self.HQ_imgs[index], cv2.IMREAD_COLOR) 88 | img_gt = cv2.resize(img_gt, (self.resolution, self.resolution), interpolation=cv2.INTER_AREA) 89 | 90 | # BFR degradation 91 | # We adopt the degradation of GFPGAN for simplicity, which however differs from our implementation in the paper. 92 | # Data degradation plays a key role in BFR. Please replace it with your own methods. 93 | img_gt = img_gt.astype(np.float32)/255. 94 | img_gt, img_lq = self.degrader.degrade_process(img_gt) 95 | 96 | img_gt = (torch.from_numpy(img_gt) - 0.5) / 0.5 97 | img_lq = (torch.from_numpy(img_lq) - 0.5) / 0.5 98 | 99 | img_gt = img_gt.permute(2, 0, 1).flip(0) # BGR->RGB 100 | img_lq = img_lq.permute(2, 0, 1).flip(0) # BGR->RGB 101 | 102 | return img_lq, img_gt 103 | 104 | -------------------------------------------------------------------------------- /training/loss/helpers.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | from torch.nn import Conv2d, BatchNorm2d, PReLU, ReLU, Sigmoid, MaxPool2d, AdaptiveAvgPool2d, Sequential, Module 4 | 5 | """ 6 | ArcFace implementation from [TreB1eN](https://github.com/TreB1eN/InsightFace_Pytorch) 7 | """ 8 | 9 | 10 | class Flatten(Module): 11 | def forward(self, input): 12 | return input.view(input.size(0), -1) 13 | 14 | 15 | def l2_norm(input, axis=1): 16 | norm = torch.norm(input, 2, axis, True) 17 | output = torch.div(input, norm) 18 | return output 19 | 20 | 21 | class Bottleneck(namedtuple('Block', ['in_channel', 'depth', 'stride'])): 22 | """ A named tuple describing a ResNet block. """ 23 | 24 | 25 | def get_block(in_channel, depth, num_units, stride=2): 26 | return [Bottleneck(in_channel, depth, stride)] + [Bottleneck(depth, depth, 1) for i in range(num_units - 1)] 27 | 28 | 29 | def get_blocks(num_layers): 30 | if num_layers == 50: 31 | blocks = [ 32 | get_block(in_channel=64, depth=64, num_units=3), 33 | get_block(in_channel=64, depth=128, num_units=4), 34 | get_block(in_channel=128, depth=256, num_units=14), 35 | get_block(in_channel=256, depth=512, num_units=3) 36 | ] 37 | elif num_layers == 100: 38 | blocks = [ 39 | get_block(in_channel=64, depth=64, num_units=3), 40 | get_block(in_channel=64, depth=128, num_units=13), 41 | get_block(in_channel=128, depth=256, num_units=30), 42 | get_block(in_channel=256, depth=512, num_units=3) 43 | ] 44 | elif num_layers == 152: 45 | blocks = [ 46 | get_block(in_channel=64, depth=64, num_units=3), 47 | get_block(in_channel=64, depth=128, num_units=8), 48 | get_block(in_channel=128, depth=256, num_units=36), 49 | get_block(in_channel=256, depth=512, num_units=3) 50 | ] 51 | else: 52 | raise ValueError("Invalid number of layers: {}. Must be one of [50, 100, 152]".format(num_layers)) 53 | return blocks 54 | 55 | 56 | class SEModule(Module): 57 | def __init__(self, channels, reduction): 58 | super(SEModule, self).__init__() 59 | self.avg_pool = AdaptiveAvgPool2d(1) 60 | self.fc1 = Conv2d(channels, channels // reduction, kernel_size=1, padding=0, bias=False) 61 | self.relu = ReLU(inplace=True) 62 | self.fc2 = Conv2d(channels // reduction, channels, kernel_size=1, padding=0, bias=False) 63 | self.sigmoid = Sigmoid() 64 | 65 | def forward(self, x): 66 | module_input = x 67 | x = self.avg_pool(x) 68 | x = self.fc1(x) 69 | x = self.relu(x) 70 | x = self.fc2(x) 71 | x = self.sigmoid(x) 72 | return module_input * x 73 | 74 | 75 | class bottleneck_IR(Module): 76 | def __init__(self, in_channel, depth, stride): 77 | super(bottleneck_IR, self).__init__() 78 | if in_channel == depth: 79 | self.shortcut_layer = MaxPool2d(1, stride) 80 | else: 81 | self.shortcut_layer = Sequential( 82 | Conv2d(in_channel, depth, (1, 1), stride, bias=False), 83 | BatchNorm2d(depth) 84 | ) 85 | self.res_layer = Sequential( 86 | BatchNorm2d(in_channel), 87 | Conv2d(in_channel, depth, (3, 3), (1, 1), 1, bias=False), PReLU(depth), 88 | Conv2d(depth, depth, (3, 3), stride, 1, bias=False), BatchNorm2d(depth) 89 | ) 90 | 91 | def forward(self, x): 92 | shortcut = self.shortcut_layer(x) 93 | res = self.res_layer(x) 94 | return res + shortcut 95 | 96 | 97 | class bottleneck_IR_SE(Module): 98 | def __init__(self, in_channel, depth, stride): 99 | super(bottleneck_IR_SE, self).__init__() 100 | if in_channel == depth: 101 | self.shortcut_layer = MaxPool2d(1, stride) 102 | else: 103 | self.shortcut_layer = Sequential( 104 | Conv2d(in_channel, depth, (1, 1), stride, bias=False), 105 | BatchNorm2d(depth) 106 | ) 107 | self.res_layer = Sequential( 108 | BatchNorm2d(in_channel), 109 | Conv2d(in_channel, depth, (3, 3), (1, 1), 1, bias=False), 110 | PReLU(depth), 111 | Conv2d(depth, depth, (3, 3), stride, 1, bias=False), 112 | BatchNorm2d(depth), 113 | SEModule(depth, 16) 114 | ) 115 | 116 | def forward(self, x): 117 | shortcut = self.shortcut_layer(x) 118 | res = self.res_layer(x) 119 | return res + shortcut 120 | -------------------------------------------------------------------------------- /training/loss/id_loss.py: -------------------------------------------------------------------------------- 1 | import os 2 | import torch 3 | from torch import nn 4 | from model_irse import Backbone 5 | 6 | class IDLoss(nn.Module): 7 | def __init__(self, base_dir='./', device='cuda', ckpt_dict=None): 8 | super(IDLoss, self).__init__() 9 | print('Loading ResNet ArcFace') 10 | self.facenet = Backbone(input_size=112, num_layers=50, drop_ratio=0.6, mode='ir_se').to(device) 11 | if ckpt_dict is None: 12 | self.facenet.load_state_dict(torch.load(os.path.join(base_dir, 'weights', 'model_ir_se50.pth'), map_location=torch.device('cpu'))) 13 | else: 14 | self.facenet.load_state_dict(ckpt_dict) 15 | self.face_pool = torch.nn.AdaptiveAvgPool2d((112, 112)) 16 | self.facenet.eval() 17 | 18 | def extract_feats(self, x): 19 | _, _, h, w = x.shape 20 | assert h==w 21 | ss = h//256 22 | x = x[:, :, 35*ss:-33*ss, 32*ss:-36*ss] # Crop interesting region 23 | x = self.face_pool(x) 24 | x_feats = self.facenet(x) 25 | return x_feats 26 | 27 | def forward(self, y_hat, y, x): 28 | n_samples = x.shape[0] 29 | x_feats = self.extract_feats(x) 30 | y_feats = self.extract_feats(y) # Otherwise use the feature from there 31 | y_hat_feats = self.extract_feats(y_hat) 32 | y_feats = y_feats.detach() 33 | loss = 0 34 | sim_improvement = 0 35 | id_logs = [] 36 | count = 0 37 | for i in range(n_samples): 38 | diff_target = y_hat_feats[i].dot(y_feats[i]) 39 | diff_input = y_hat_feats[i].dot(x_feats[i]) 40 | diff_views = y_feats[i].dot(x_feats[i]) 41 | id_logs.append({'diff_target': float(diff_target), 42 | 'diff_input': float(diff_input), 43 | 'diff_views': float(diff_views)}) 44 | loss += 1 - diff_target 45 | id_diff = float(diff_target) - float(diff_views) 46 | sim_improvement += id_diff 47 | count += 1 48 | 49 | return loss / count, sim_improvement / count, id_logs 50 | 51 | -------------------------------------------------------------------------------- /training/loss/model_irse.py: -------------------------------------------------------------------------------- 1 | from torch.nn import Linear, Conv2d, BatchNorm1d, BatchNorm2d, PReLU, Dropout, Sequential, Module 2 | #from models.encoders.helpers import get_blocks, Flatten, bottleneck_IR, bottleneck_IR_SE, l2_norm 3 | from helpers import get_blocks, Flatten, bottleneck_IR, bottleneck_IR_SE, l2_norm 4 | 5 | """ 6 | Modified Backbone implementation from [TreB1eN](https://github.com/TreB1eN/InsightFace_Pytorch) 7 | """ 8 | 9 | 10 | class Backbone(Module): 11 | def __init__(self, input_size, num_layers, mode='ir', drop_ratio=0.4, affine=True): 12 | super(Backbone, self).__init__() 13 | assert input_size in [112, 224], "input_size should be 112 or 224" 14 | assert num_layers in [50, 100, 152], "num_layers should be 50, 100 or 152" 15 | assert mode in ['ir', 'ir_se'], "mode should be ir or ir_se" 16 | blocks = get_blocks(num_layers) 17 | if mode == 'ir': 18 | unit_module = bottleneck_IR 19 | elif mode == 'ir_se': 20 | unit_module = bottleneck_IR_SE 21 | self.input_layer = Sequential(Conv2d(3, 64, (3, 3), 1, 1, bias=False), 22 | BatchNorm2d(64), 23 | PReLU(64)) 24 | if input_size == 112: 25 | self.output_layer = Sequential(BatchNorm2d(512), 26 | Dropout(drop_ratio), 27 | Flatten(), 28 | Linear(512 * 7 * 7, 512), 29 | BatchNorm1d(512, affine=affine)) 30 | else: 31 | self.output_layer = Sequential(BatchNorm2d(512), 32 | Dropout(drop_ratio), 33 | Flatten(), 34 | Linear(512 * 14 * 14, 512), 35 | BatchNorm1d(512, affine=affine)) 36 | 37 | modules = [] 38 | for block in blocks: 39 | for bottleneck in block: 40 | modules.append(unit_module(bottleneck.in_channel, 41 | bottleneck.depth, 42 | bottleneck.stride)) 43 | self.body = Sequential(*modules) 44 | 45 | def forward(self, x): 46 | x = self.input_layer(x) 47 | x = self.body(x) 48 | x = self.output_layer(x) 49 | return l2_norm(x) 50 | 51 | 52 | def IR_50(input_size): 53 | """Constructs a ir-50 model.""" 54 | model = Backbone(input_size, num_layers=50, mode='ir', drop_ratio=0.4, affine=False) 55 | return model 56 | 57 | 58 | def IR_101(input_size): 59 | """Constructs a ir-101 model.""" 60 | model = Backbone(input_size, num_layers=100, mode='ir', drop_ratio=0.4, affine=False) 61 | return model 62 | 63 | 64 | def IR_152(input_size): 65 | """Constructs a ir-152 model.""" 66 | model = Backbone(input_size, num_layers=152, mode='ir', drop_ratio=0.4, affine=False) 67 | return model 68 | 69 | 70 | def IR_SE_50(input_size): 71 | """Constructs a ir_se-50 model.""" 72 | model = Backbone(input_size, num_layers=50, mode='ir_se', drop_ratio=0.4, affine=False) 73 | return model 74 | 75 | 76 | def IR_SE_101(input_size): 77 | """Constructs a ir_se-101 model.""" 78 | model = Backbone(input_size, num_layers=100, mode='ir_se', drop_ratio=0.4, affine=False) 79 | return model 80 | 81 | 82 | def IR_SE_152(input_size): 83 | """Constructs a ir_se-152 model.""" 84 | model = Backbone(input_size, num_layers=152, mode='ir_se', drop_ratio=0.4, affine=False) 85 | return model 86 | -------------------------------------------------------------------------------- /training/lpips/__init__.py: -------------------------------------------------------------------------------- 1 | 2 | from __future__ import absolute_import 3 | from __future__ import division 4 | from __future__ import print_function 5 | 6 | import numpy as np 7 | import torch 8 | # from torch.autograd import Variable 9 | 10 | from lpips.trainer import * 11 | from lpips.lpips import * 12 | 13 | # class PerceptualLoss(torch.nn.Module): 14 | # def __init__(self, model='lpips', net='alex', spatial=False, use_gpu=False, gpu_ids=[0], version='0.1'): # VGG using our perceptually-learned weights (LPIPS metric) 15 | # # def __init__(self, model='net', net='vgg', use_gpu=True): # "default" way of using VGG as a perceptual loss 16 | # super(PerceptualLoss, self).__init__() 17 | # print('Setting up Perceptual loss...') 18 | # self.use_gpu = use_gpu 19 | # self.spatial = spatial 20 | # self.gpu_ids = gpu_ids 21 | # self.model = dist_model.DistModel() 22 | # self.model.initialize(model=model, net=net, use_gpu=use_gpu, spatial=self.spatial, gpu_ids=gpu_ids, version=version) 23 | # print('...[%s] initialized'%self.model.name()) 24 | # print('...Done') 25 | 26 | # def forward(self, pred, target, normalize=False): 27 | # """ 28 | # Pred and target are Variables. 29 | # If normalize is True, assumes the images are between [0,1] and then scales them between [-1,+1] 30 | # If normalize is False, assumes the images are already between [-1,+1] 31 | 32 | # Inputs pred and target are Nx3xHxW 33 | # Output pytorch Variable N long 34 | # """ 35 | 36 | # if normalize: 37 | # target = 2 * target - 1 38 | # pred = 2 * pred - 1 39 | 40 | # return self.model.forward(target, pred) 41 | 42 | def normalize_tensor(in_feat,eps=1e-10): 43 | norm_factor = torch.sqrt(torch.sum(in_feat**2,dim=1,keepdim=True)) 44 | return in_feat/(norm_factor+eps) 45 | 46 | def l2(p0, p1, range=255.): 47 | return .5*np.mean((p0 / range - p1 / range)**2) 48 | 49 | def psnr(p0, p1, peak=255.): 50 | return 10*np.log10(peak**2/np.mean((1.*p0-1.*p1)**2)) 51 | 52 | def dssim(p0, p1, range=255.): 53 | from skimage.measure import compare_ssim 54 | return (1 - compare_ssim(p0, p1, data_range=range, multichannel=True)) / 2. 55 | 56 | def rgb2lab(in_img,mean_cent=False): 57 | from skimage import color 58 | img_lab = color.rgb2lab(in_img) 59 | if(mean_cent): 60 | img_lab[:,:,0] = img_lab[:,:,0]-50 61 | return img_lab 62 | 63 | def tensor2np(tensor_obj): 64 | # change dimension of a tensor object into a numpy array 65 | return tensor_obj[0].cpu().float().numpy().transpose((1,2,0)) 66 | 67 | def np2tensor(np_obj): 68 | # change dimenion of np array into tensor array 69 | return torch.Tensor(np_obj[:, :, :, np.newaxis].transpose((3, 2, 0, 1))) 70 | 71 | def tensor2tensorlab(image_tensor,to_norm=True,mc_only=False): 72 | # image tensor to lab tensor 73 | from skimage import color 74 | 75 | img = tensor2im(image_tensor) 76 | img_lab = color.rgb2lab(img) 77 | if(mc_only): 78 | img_lab[:,:,0] = img_lab[:,:,0]-50 79 | if(to_norm and not mc_only): 80 | img_lab[:,:,0] = img_lab[:,:,0]-50 81 | img_lab = img_lab/100. 82 | 83 | return np2tensor(img_lab) 84 | 85 | def tensorlab2tensor(lab_tensor,return_inbnd=False): 86 | from skimage import color 87 | import warnings 88 | warnings.filterwarnings("ignore") 89 | 90 | lab = tensor2np(lab_tensor)*100. 91 | lab[:,:,0] = lab[:,:,0]+50 92 | 93 | rgb_back = 255.*np.clip(color.lab2rgb(lab.astype('float')),0,1) 94 | if(return_inbnd): 95 | # convert back to lab, see if we match 96 | lab_back = color.rgb2lab(rgb_back.astype('uint8')) 97 | mask = 1.*np.isclose(lab_back,lab,atol=2.) 98 | mask = np2tensor(np.prod(mask,axis=2)[:,:,np.newaxis]) 99 | return (im2tensor(rgb_back),mask) 100 | else: 101 | return im2tensor(rgb_back) 102 | 103 | def load_image(path): 104 | if(path[-3:] == 'dng'): 105 | import rawpy 106 | with rawpy.imread(path) as raw: 107 | img = raw.postprocess() 108 | elif(path[-3:]=='bmp' or path[-3:]=='jpg' or path[-3:]=='png'): 109 | import cv2 110 | return cv2.imread(path)[:,:,::-1] 111 | else: 112 | img = (255*plt.imread(path)[:,:,:3]).astype('uint8') 113 | 114 | return img 115 | 116 | def rgb2lab(input): 117 | from skimage import color 118 | return color.rgb2lab(input / 255.) 119 | 120 | def tensor2im(image_tensor, imtype=np.uint8, cent=1., factor=255./2.): 121 | image_numpy = image_tensor[0].cpu().float().numpy() 122 | image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + cent) * factor 123 | return image_numpy.astype(imtype) 124 | 125 | def im2tensor(image, imtype=np.uint8, cent=1., factor=255./2.): 126 | return torch.Tensor((image / factor - cent) 127 | [:, :, :, np.newaxis].transpose((3, 2, 0, 1))) 128 | 129 | def tensor2vec(vector_tensor): 130 | return vector_tensor.data.cpu().numpy()[:, :, 0, 0] 131 | 132 | 133 | def tensor2im(image_tensor, imtype=np.uint8, cent=1., factor=255./2.): 134 | # def tensor2im(image_tensor, imtype=np.uint8, cent=1., factor=1.): 135 | image_numpy = image_tensor[0].cpu().float().numpy() 136 | image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + cent) * factor 137 | return image_numpy.astype(imtype) 138 | 139 | def im2tensor(image, imtype=np.uint8, cent=1., factor=255./2.): 140 | # def im2tensor(image, imtype=np.uint8, cent=1., factor=1.): 141 | return torch.Tensor((image / factor - cent) 142 | [:, :, :, np.newaxis].transpose((3, 2, 0, 1))) 143 | 144 | 145 | 146 | def voc_ap(rec, prec, use_07_metric=False): 147 | """ ap = voc_ap(rec, prec, [use_07_metric]) 148 | Compute VOC AP given precision and recall. 149 | If use_07_metric is true, uses the 150 | VOC 07 11 point method (default:False). 151 | """ 152 | if use_07_metric: 153 | # 11 point metric 154 | ap = 0. 155 | for t in np.arange(0., 1.1, 0.1): 156 | if np.sum(rec >= t) == 0: 157 | p = 0 158 | else: 159 | p = np.max(prec[rec >= t]) 160 | ap = ap + p / 11. 161 | else: 162 | # correct AP calculation 163 | # first append sentinel values at the end 164 | mrec = np.concatenate(([0.], rec, [1.])) 165 | mpre = np.concatenate(([0.], prec, [0.])) 166 | 167 | # compute the precision envelope 168 | for i in range(mpre.size - 1, 0, -1): 169 | mpre[i - 1] = np.maximum(mpre[i - 1], mpre[i]) 170 | 171 | # to calculate area under PR curve, look for points 172 | # where X axis (recall) changes value 173 | i = np.where(mrec[1:] != mrec[:-1])[0] 174 | 175 | # and sum (\Delta recall) * prec 176 | ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1]) 177 | return ap 178 | 179 | -------------------------------------------------------------------------------- /training/lpips/pretrained_networks.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import torch 3 | from torchvision import models as tv 4 | 5 | class squeezenet(torch.nn.Module): 6 | def __init__(self, requires_grad=False, pretrained=True): 7 | super(squeezenet, self).__init__() 8 | pretrained_features = tv.squeezenet1_1(pretrained=pretrained).features 9 | self.slice1 = torch.nn.Sequential() 10 | self.slice2 = torch.nn.Sequential() 11 | self.slice3 = torch.nn.Sequential() 12 | self.slice4 = torch.nn.Sequential() 13 | self.slice5 = torch.nn.Sequential() 14 | self.slice6 = torch.nn.Sequential() 15 | self.slice7 = torch.nn.Sequential() 16 | self.N_slices = 7 17 | for x in range(2): 18 | self.slice1.add_module(str(x), pretrained_features[x]) 19 | for x in range(2,5): 20 | self.slice2.add_module(str(x), pretrained_features[x]) 21 | for x in range(5, 8): 22 | self.slice3.add_module(str(x), pretrained_features[x]) 23 | for x in range(8, 10): 24 | self.slice4.add_module(str(x), pretrained_features[x]) 25 | for x in range(10, 11): 26 | self.slice5.add_module(str(x), pretrained_features[x]) 27 | for x in range(11, 12): 28 | self.slice6.add_module(str(x), pretrained_features[x]) 29 | for x in range(12, 13): 30 | self.slice7.add_module(str(x), pretrained_features[x]) 31 | if not requires_grad: 32 | for param in self.parameters(): 33 | param.requires_grad = False 34 | 35 | def forward(self, X): 36 | h = self.slice1(X) 37 | h_relu1 = h 38 | h = self.slice2(h) 39 | h_relu2 = h 40 | h = self.slice3(h) 41 | h_relu3 = h 42 | h = self.slice4(h) 43 | h_relu4 = h 44 | h = self.slice5(h) 45 | h_relu5 = h 46 | h = self.slice6(h) 47 | h_relu6 = h 48 | h = self.slice7(h) 49 | h_relu7 = h 50 | vgg_outputs = namedtuple("SqueezeOutputs", ['relu1','relu2','relu3','relu4','relu5','relu6','relu7']) 51 | out = vgg_outputs(h_relu1,h_relu2,h_relu3,h_relu4,h_relu5,h_relu6,h_relu7) 52 | 53 | return out 54 | 55 | 56 | class alexnet(torch.nn.Module): 57 | def __init__(self, requires_grad=False, pretrained=True): 58 | super(alexnet, self).__init__() 59 | alexnet_pretrained_features = tv.alexnet(pretrained=pretrained).features 60 | self.slice1 = torch.nn.Sequential() 61 | self.slice2 = torch.nn.Sequential() 62 | self.slice3 = torch.nn.Sequential() 63 | self.slice4 = torch.nn.Sequential() 64 | self.slice5 = torch.nn.Sequential() 65 | self.N_slices = 5 66 | for x in range(2): 67 | self.slice1.add_module(str(x), alexnet_pretrained_features[x]) 68 | for x in range(2, 5): 69 | self.slice2.add_module(str(x), alexnet_pretrained_features[x]) 70 | for x in range(5, 8): 71 | self.slice3.add_module(str(x), alexnet_pretrained_features[x]) 72 | for x in range(8, 10): 73 | self.slice4.add_module(str(x), alexnet_pretrained_features[x]) 74 | for x in range(10, 12): 75 | self.slice5.add_module(str(x), alexnet_pretrained_features[x]) 76 | if not requires_grad: 77 | for param in self.parameters(): 78 | param.requires_grad = False 79 | 80 | def forward(self, X): 81 | h = self.slice1(X) 82 | h_relu1 = h 83 | h = self.slice2(h) 84 | h_relu2 = h 85 | h = self.slice3(h) 86 | h_relu3 = h 87 | h = self.slice4(h) 88 | h_relu4 = h 89 | h = self.slice5(h) 90 | h_relu5 = h 91 | alexnet_outputs = namedtuple("AlexnetOutputs", ['relu1', 'relu2', 'relu3', 'relu4', 'relu5']) 92 | out = alexnet_outputs(h_relu1, h_relu2, h_relu3, h_relu4, h_relu5) 93 | 94 | return out 95 | 96 | class vgg16(torch.nn.Module): 97 | def __init__(self, requires_grad=False, pretrained=True): 98 | super(vgg16, self).__init__() 99 | vgg_pretrained_features = tv.vgg16(pretrained=pretrained).features 100 | self.slice1 = torch.nn.Sequential() 101 | self.slice2 = torch.nn.Sequential() 102 | self.slice3 = torch.nn.Sequential() 103 | self.slice4 = torch.nn.Sequential() 104 | self.slice5 = torch.nn.Sequential() 105 | self.N_slices = 5 106 | for x in range(4): 107 | self.slice1.add_module(str(x), vgg_pretrained_features[x]) 108 | for x in range(4, 9): 109 | self.slice2.add_module(str(x), vgg_pretrained_features[x]) 110 | for x in range(9, 16): 111 | self.slice3.add_module(str(x), vgg_pretrained_features[x]) 112 | for x in range(16, 23): 113 | self.slice4.add_module(str(x), vgg_pretrained_features[x]) 114 | for x in range(23, 30): 115 | self.slice5.add_module(str(x), vgg_pretrained_features[x]) 116 | if not requires_grad: 117 | for param in self.parameters(): 118 | param.requires_grad = False 119 | 120 | def forward(self, X): 121 | h = self.slice1(X) 122 | h_relu1_2 = h 123 | h = self.slice2(h) 124 | h_relu2_2 = h 125 | h = self.slice3(h) 126 | h_relu3_3 = h 127 | h = self.slice4(h) 128 | h_relu4_3 = h 129 | h = self.slice5(h) 130 | h_relu5_3 = h 131 | vgg_outputs = namedtuple("VggOutputs", ['relu1_2', 'relu2_2', 'relu3_3', 'relu4_3', 'relu5_3']) 132 | out = vgg_outputs(h_relu1_2, h_relu2_2, h_relu3_3, h_relu4_3, h_relu5_3) 133 | 134 | return out 135 | 136 | 137 | 138 | class resnet(torch.nn.Module): 139 | def __init__(self, requires_grad=False, pretrained=True, num=18): 140 | super(resnet, self).__init__() 141 | if(num==18): 142 | self.net = tv.resnet18(pretrained=pretrained) 143 | elif(num==34): 144 | self.net = tv.resnet34(pretrained=pretrained) 145 | elif(num==50): 146 | self.net = tv.resnet50(pretrained=pretrained) 147 | elif(num==101): 148 | self.net = tv.resnet101(pretrained=pretrained) 149 | elif(num==152): 150 | self.net = tv.resnet152(pretrained=pretrained) 151 | self.N_slices = 5 152 | 153 | self.conv1 = self.net.conv1 154 | self.bn1 = self.net.bn1 155 | self.relu = self.net.relu 156 | self.maxpool = self.net.maxpool 157 | self.layer1 = self.net.layer1 158 | self.layer2 = self.net.layer2 159 | self.layer3 = self.net.layer3 160 | self.layer4 = self.net.layer4 161 | 162 | def forward(self, X): 163 | h = self.conv1(X) 164 | h = self.bn1(h) 165 | h = self.relu(h) 166 | h_relu1 = h 167 | h = self.maxpool(h) 168 | h = self.layer1(h) 169 | h_conv2 = h 170 | h = self.layer2(h) 171 | h_conv3 = h 172 | h = self.layer3(h) 173 | h_conv4 = h 174 | h = self.layer4(h) 175 | h_conv5 = h 176 | 177 | outputs = namedtuple("Outputs", ['relu1','conv2','conv3','conv4','conv5']) 178 | out = outputs(h_relu1, h_conv2, h_conv3, h_conv4, h_conv5) 179 | 180 | return out 181 | -------------------------------------------------------------------------------- /training/lpips/weights/v0.0/alex.pth: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/training/lpips/weights/v0.0/alex.pth -------------------------------------------------------------------------------- /training/lpips/weights/v0.0/squeeze.pth: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/training/lpips/weights/v0.0/squeeze.pth -------------------------------------------------------------------------------- /training/lpips/weights/v0.0/vgg.pth: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/training/lpips/weights/v0.0/vgg.pth -------------------------------------------------------------------------------- /training/lpips/weights/v0.1/alex.pth: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/training/lpips/weights/v0.1/alex.pth -------------------------------------------------------------------------------- /training/lpips/weights/v0.1/squeeze.pth: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/training/lpips/weights/v0.1/squeeze.pth -------------------------------------------------------------------------------- /training/lpips/weights/v0.1/vgg.pth: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/training/lpips/weights/v0.1/vgg.pth -------------------------------------------------------------------------------- /val/hq/ffhq_00000.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/val/hq/ffhq_00000.png -------------------------------------------------------------------------------- /val/hq/ffhq_00001.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/val/hq/ffhq_00001.png -------------------------------------------------------------------------------- /val/hq/ffhq_00002.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/val/hq/ffhq_00002.png -------------------------------------------------------------------------------- /val/hq/ffhq_00003.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/val/hq/ffhq_00003.png -------------------------------------------------------------------------------- /val/lq/ffhq_00000.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/val/lq/ffhq_00000.png -------------------------------------------------------------------------------- /val/lq/ffhq_00001.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/val/lq/ffhq_00001.png -------------------------------------------------------------------------------- /val/lq/ffhq_00002.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/val/lq/ffhq_00002.png -------------------------------------------------------------------------------- /val/lq/ffhq_00003.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yangxy/GPEN/b611a9f2d05fdc5b82715219ed0f252e42e4d8be/val/lq/ffhq_00003.png -------------------------------------------------------------------------------- /weights/README.md: -------------------------------------------------------------------------------- 1 | ## Pre-trained Model 2 | 3 | Download RetinaFace model and our pre-trained model and put them here. 4 | 5 | [RetinaFace-R50](https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/RetinaFace-R50.pth) | [ParseNet-latest](https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/ParseNet-latest.pth) | [model_ir_se50](https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/model_ir_se50.pth) | [GPEN-BFR-512](https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/GPEN-BFR-512.pth) | [GPEN-BFR-512-D](https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/GPEN-BFR-512-D.pth) | [GPEN-BFR-256](https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/GPEN-BFR-256.pth) | [GPEN-BFR-256-D](https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/GPEN-BFR-256-D.pth) | [GPEN-Colorization-1024](https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/GPEN-Colorization-1024.pth) | [GPEN-Inpainting-1024](https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/GPEN-Inpainting-1024.pth) | [GPEN-Seg2face-512](https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/GPEN-Seg2face-512.pth) | [realesrnet_x1](https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/realesrnet_x1.pth) | [realesrnet_x2](https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/realesrnet_x2.pth) | [realesrnet_x4](https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/realesrnet_x4.pth) --------------------------------------------------------------------------------