add paste back and update README

Author: xinntao

README.md

@@ -8,13 +8,13 @@ GFPGAN is a blind face restoration algorithm towards real-world face images.
 [Colab Demo](https://colab.research.google.com/drive/1sVsoBd9AjckIXThgtZhGrHRfFI6UUYOo)
 
 ### :book: GFP-GAN: Towards Real-World Blind Face Restoration with Generative Facial Prior
-> [[Paper](https://arxiv.org/abs/2101.04061)] &emsp; [[Project Page](https://xinntao.github.io/projects/gfpgan)] &emsp; [[Demo]()] <br>
+> [[Paper](https://arxiv.org/abs/2101.04061)] &emsp; [[Project Page](https://xinntao.github.io/projects/gfpgan)] &emsp; [Demo] <br>
 > [Xintao Wang](https://xinntao.github.io/), [Yu Li](https://yu-li.github.io/), [Honglun Zhang](https://scholar.google.com/citations?hl=en&user=KjQLROoAAAAJ), [Ying Shan](https://scholar.google.com/citations?user=4oXBp9UAAAAJ&hl=en) <br>
 > Applied Research Center (ARC), Tencent PCG
 
 #### Abstract
 
-Blind face restoration usually relies on facial priors, such as facial geometry prior or reference prior, to restore realistic and faithful details. However, very low-quality inputs cannot offer accurate geometric prior while high-quality references are inaccessible, limiting the applicability in real-world scenarios. In this work, we propose GFP-GAN that leverages rich and diverse priors encapsulated in a pretrained face GAN for blind face restoration. This Generative Facial Prior (GFP) is incorporated into the face restoration process via novel channel-split spatial feature transform layers, which allow our method to achieve a good balance of realness and fidelity. Thanks to the powerful generative facial prior and delicate designs, our GFP-GAN could jointly restore facial details and enhance colors with just a single forward pass, while GAN inversion methods require expensive image-specific optimization at inference. Extensive experiments show that our method achieves superior performance to prior art on both synthetic and real-world datasets.
+Blind face restoration usually relies on facial priors, such as facial geometry prior or reference prior, to restore realistic and faithful details. However, very low-quality inputs cannot offer accurate geometric prior while high-quality references are inaccessible, limiting the applicability in real-world scenarios. In this work, we propose GFP-GAN that leverages **rich and diverse priors encapsulated in a pretrained face GAN** for blind face restoration. This Generative Facial Prior (GFP) is incorporated into the face restoration process via novel channel-split spatial feature transform layers, which allow our method to achieve a good balance of realness and fidelity. Thanks to the powerful generative facial prior and delicate designs, our GFP-GAN could jointly restore facial details and enhance colors with just a single forward pass, while GAN inversion methods require expensive image-specific optimization at inference. Extensive experiments show that our method achieves superior performance to prior art on both synthetic and real-world datasets.
 
 #### BibTeX
 
@@ -43,18 +43,39 @@ Blind face restoration usually relies on facial priors, such as facial geometry
 
     ```bash
     git clone https://github.com/xinntao/GFPGAN.git
+    cd GFPGAN
     ```
 
 1. Install dependent packages
 
     ```bash
-    cd GFPGAN
+    # Install basicsr - https://github.com/xinntao/BasicSR
+    # We use BasicSR for both training and inference
+    # Set BASICSR_EXT=True to compile the cuda extensions in the BasicSR - It may take several minutes to compile, please be patient
+    BASICSR_EXT=True pip install basicsr
+
+    # Install facexlib - https://github.com/xinntao/facexlib
+    # We use face detection and face restoration helper in the facexlib package
+    pip install facexlib
+
     pip install -r requirements.txt
     ```
 
 ## :zap: Quick Inference
 
-> python inference_gfpgan_full.py --model_path experiments/pretrained_models/GFPGANv1.pth --test_path inputs
+Download pre-trained models: [GFPGANv1.pth](https://github.com/TencentARC/GFPGAN/releases/download/v0.1.0/GFPGANv1.pth)
+
+> wget https://github.com/TencentARC/GFPGAN/releases/download/v0.1.0/GFPGANv1.pth -P experiments/pretrained_models
+
+> python inference_gfpgan_full.py --model_path experiments/pretrained_models/GFPGANv1.pth --test_path inputs/whole_imgs
+> python inference_gfpgan_full.py --model_path experiments/pretrained_models/GFPGANv1.pth --test_path inputs/cropped_faces --aligned
+
+## :computer: Training
+
+We provide a complete training for GFPGAN. <br>
+You could improve it according to your own needs.
+
+> python -m torch.distributed.launch --nproc_per_node=4 --master_port=22021 train.py -opt train_gfpgan_v1.yml --launcher pytorch
 
 ## :scroll: License and Acknowledgement
 

README_CN.md

@@ -8,13 +8,13 @@ GFPGAN is a blind face restoration algorithm towards real-world face images.
 [Colab Demo](https://colab.research.google.com/drive/1sVsoBd9AjckIXThgtZhGrHRfFI6UUYOo)
 
 ### :book: GFP-GAN: Towards Real-World Blind Face Restoration with Generative Facial Prior
-> [[Paper](https://arxiv.org/abs/2101.04061)] &emsp; [[Project Page](https://xinntao.github.io/projects/gfpgan)] &emsp; [[Demo]()] <br>
+> [[Paper](https://arxiv.org/abs/2101.04061)] &emsp; [[Project Page](https://xinntao.github.io/projects/gfpgan)] &emsp; [Demo] <br>
 > [Xintao Wang](https://xinntao.github.io/), [Yu Li](https://yu-li.github.io/), [Honglun Zhang](https://scholar.google.com/citations?hl=en&user=KjQLROoAAAAJ), [Ying Shan](https://scholar.google.com/citations?user=4oXBp9UAAAAJ&hl=en) <br>
 > Applied Research Center (ARC), Tencent PCG
 
 #### Abstract
 
-Blind face restoration usually relies on facial priors, such as facial geometry prior or reference prior, to restore realistic and faithful details. However, very low-quality inputs cannot offer accurate geometric prior while high-quality references are inaccessible, limiting the applicability in real-world scenarios. In this work, we propose GFP-GAN that leverages rich and diverse priors encapsulated in a pretrained face GAN for blind face restoration. This Generative Facial Prior (GFP) is incorporated into the face restoration process via novel channel-split spatial feature transform layers, which allow our method to achieve a good balance of realness and fidelity. Thanks to the powerful generative facial prior and delicate designs, our GFP-GAN could jointly restore facial details and enhance colors with just a single forward pass, while GAN inversion methods require expensive image-specific optimization at inference. Extensive experiments show that our method achieves superior performance to prior art on both synthetic and real-world datasets.
+Blind face restoration usually relies on facial priors, such as facial geometry prior or reference prior, to restore realistic and faithful details. However, very low-quality inputs cannot offer accurate geometric prior while high-quality references are inaccessible, limiting the applicability in real-world scenarios. In this work, we propose GFP-GAN that leverages **rich and diverse priors encapsulated in a pretrained face GAN** for blind face restoration. This Generative Facial Prior (GFP) is incorporated into the face restoration process via novel channel-split spatial feature transform layers, which allow our method to achieve a good balance of realness and fidelity. Thanks to the powerful generative facial prior and delicate designs, our GFP-GAN could jointly restore facial details and enhance colors with just a single forward pass, while GAN inversion methods require expensive image-specific optimization at inference. Extensive experiments show that our method achieves superior performance to prior art on both synthetic and real-world datasets.
 
 #### BibTeX
 
@@ -43,18 +43,39 @@ Blind face restoration usually relies on facial priors, such as facial geometry
 
     ```bash
     git clone https://github.com/xinntao/GFPGAN.git
+    cd GFPGAN
     ```
 
 1. Install dependent packages
 
     ```bash
-    cd GFPGAN
+    # Install basicsr - https://github.com/xinntao/BasicSR
+    # We use BasicSR for both training and inference
+    # Set BASICSR_EXT=True to compile the cuda extensions in the BasicSR - It may take several minutes to compile, please be patient
+    BASICSR_EXT=True pip install basicsr
+
+    # Install facexlib - https://github.com/xinntao/facexlib
+    # We use face detection and face restoration helper in the facexlib package
+    pip install facexlib
+
     pip install -r requirements.txt
     ```
 
 ## :zap: Quick Inference
 
-> python inference_gfpgan_full.py --model_path experiments/pretrained_models/GFPGANv1.pth --test_path inputs
+Download pre-trained models: [GFPGANv1.pth](https://github.com/TencentARC/GFPGAN/releases/download/v0.1.0/GFPGANv1.pth)
+
+> wget https://github.com/TencentARC/GFPGAN/releases/download/v0.1.0/GFPGANv1.pth -P experiments/pretrained_models
+
+> python inference_gfpgan_full.py --model_path experiments/pretrained_models/GFPGANv1.pth --test_path inputs/whole_imgs
+> python inference_gfpgan_full.py --model_path experiments/pretrained_models/GFPGANv1.pth --test_path inputs/cropped_faces --aligned
+
+## :computer: Training
+
+We provide a complete training for GFPGAN. <br>
+You could improve it according to your own needs.
+
+> python -m torch.distributed.launch --nproc_per_node=4 --master_port=22021 train.py -opt train_gfpgan_v1.yml --launcher pytorch
 
 ## :scroll: License and Acknowledgement
 

inference_gfpgan_full.py

@@ -11,7 +11,14 @@
 from basicsr.utils import img2tensor, imwrite, tensor2img
 
 
-def restoration(gfpgan, face_helper, img_path, save_root, has_aligned=False, only_center_face=True, suffix=None):
+def restoration(gfpgan,
+                face_helper,
+                img_path,
+                save_root,
+                has_aligned=False,
+                only_center_face=True,
+                suffix=None,
+                paste_back=False):
     # read image
     img_name = os.path.basename(img_path)
     print(f'Processing {img_name} ...')
@@ -60,6 +67,12 @@ def restoration(gfpgan, face_helper, img_path, save_root, has_aligned=False, onl
         cmp_img = np.concatenate((cropped_face, restored_face), axis=1)
         imwrite(cmp_img, os.path.join(save_root, 'cmp', f'{basename}_{idx:02d}.png'))
 
+    if not has_aligned and paste_back:
+        face_helper.get_inverse_affine(None)
+        save_restore_path = os.path.join(save_root, 'restored_imgs', img_name)
+        # paste each restored face to the input image
+        face_helper.paste_faces_to_input_image(save_restore_path)
+
 
 if __name__ == '__main__':
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
@@ -71,6 +84,7 @@ def restoration(gfpgan, face_helper, img_path, save_root, has_aligned=False, onl
     parser.add_argument('--suffix', type=str, default=None, help='Suffix of the restored faces')
     parser.add_argument('--only_center_face', action='store_true')
     parser.add_argument('--aligned', action='store_true')
+    parser.add_argument('--paste_back', action='store_true')
 
     args = parser.parse_args()
     if args.test_path.endswith('/'):
@@ -110,6 +124,7 @@ def restoration(gfpgan, face_helper, img_path, save_root, has_aligned=False, onl
             save_root,
             has_aligned=args.aligned,
             only_center_face=args.only_center_face,
-            suffix=args.suffix)
+            suffix=args.suffix,
+            paste_back=args.paste_back)
 
     print('Results are in the <results> folder.')