add the model file and (PaddlePaddle#634)

* add the model files and tipc of GPEN * add the model files and tipc of GPEN
driyca · May 18, 2022 · 9a31bdf · 9a31bdf
1 parent 1e1e2ad
commit 9a31bdf
Show file tree

Hide file tree

Showing 24 changed files with 2,324 additions and 51 deletions.
diff --git a/applications/tools/gpen.py b/applications/tools/gpen.py
@@ -0,0 +1,93 @@
+#   Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import sys
+
+sys.path.append(".")
+import argparse
+import paddle
+from ppgan.apps import GPENPredictor
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--output_path",
+                        type=str,
+                        default='output_dir',
+                        help="path to output image dir")
+
+    parser.add_argument("--weight_path",
+                        type=str,
+                        default=None,
+                        help="path to model checkpoint path")
+
+    parser.add_argument("--test_img",
+                        type=str,
+                        default='data/gpen/lite_data/15006.png',
+                        help="path of test image")
+
+    parser.add_argument("--model_type",
+                        type=str,
+                        default=None,
+                        help="type of model for loading pretrained model")
+
+    parser.add_argument("--seed",
+                        type=int,
+                        default=None,
+                        help="sample random seed for model's image generation")
+
+    parser.add_argument("--size",
+                        type=int,
+                        default=256,
+                        help="resolution of output image")
+
+    parser.add_argument("--style_dim",
+                        type=int,
+                        default=512,
+                        help="number of style dimension")
+
+    parser.add_argument("--n_mlp",
+                        type=int,
+                        default=8,
+                        help="number of mlp layer depth")
+
+    parser.add_argument("--channel_multiplier",
+                        type=int,
+                        default=1,
+                        help="number of channel multiplier")
+
+    parser.add_argument("--narrow",
+                        type=float,
+                        default=0.5,
+                        help="number of channel narrow")
+
+    parser.add_argument("--cpu",
+                        dest="cpu",
+                        action="store_true",
+                        help="cpu mode.")
+
+    args = parser.parse_args()
+
+    if args.cpu:
+        paddle.set_device('cpu')
+
+    predictor = GPENPredictor(output_path=args.output_path,
+                              weight_path=args.weight_path,
+                              model_type=args.model_type,
+                              seed=args.seed,
+                              size=args.size,
+                              style_dim=args.style_dim,
+                              n_mlp=args.n_mlp,
+                              narrow=args.narrow,
+                              channel_multiplier=args.channel_multiplier)
+    predictor.run(args.test_img)
diff --git a/configs/gpen_256_ffhq.yaml b/configs/gpen_256_ffhq.yaml
@@ -0,0 +1,77 @@
+total_iters: 200000
+output_dir: output_dir
+find_unused_parameters: True
+
+
+model:
+  name: GPENModel
+  generator:
+    name: GPEN
+    size: 256
+    style_dim: 512
+    n_mlp: 8
+    channel_multiplier: 1
+    narrow: 0.5
+  discriminator:
+    name: GPENDiscriminator
+    size: 256
+    channel_multiplier: 1
+    narrow: 0.5
+
+
+export_model:
+  - {name: 'g_ema', inputs_num: 1}
+
+dataset:
+  train:
+    name: GPENDataset
+    dataroot: data/ffhq/images256x256/
+    num_workers: 0
+    batch_size: 2 #1gpus
+    size: 256
+
+  test:
+    name: GPENDataset
+    dataroot: data/ffhq/images256x256/
+    num_workers: 0
+    batch_size: 1
+    size: 256
+    amount: 100
+
+
+lr_scheduler:
+  name: CosineAnnealingRestartLR
+  learning_rate: 0.002
+  periods: [500000, 500000, 500000, 500000]
+  restart_weights: [1, 1, 1, 1]
+  eta_min: 0.002
+
+
+optimizer:
+  optimG:
+    name: Adam
+    net_names:
+      - netG
+    beta1: 0.9
+    beta2: 0.99
+  optimD:
+    name: Adam
+    net_names:
+      - netD
+    beta1: 0.9
+    beta2: 0.99
+
+log_config:
+  interval: 100
+  visiual_interval: 500
+
+snapshot_config:
+  interval: 5000
+
+validate:
+  interval: 5000
+  save_img: false
+  metrics:
+    fid:
+        name: FID
+        batch_size: 1
diff --git a/docs/en_US/tutorials/gpen.md b/docs/en_US/tutorials/gpen.md
@@ -0,0 +1,202 @@
+English | [Chinese](../../zh_CN/tutorials/gpen.md)
+
+## GPEN Blind Face Restoration Model
+
+
+## 1、Introduction
+
+The GPEN model is a blind face restoration model. The author embeds the decoder of StyleGAN V2 proposed by the previous model as the decoder of GPEN; reconstructs a simple encoder with DNN to provide input for the decoder. In this way, while retaining the excellent performance of the StyleGAN V2 decoder, the function of the model is changed from image style conversion to blind face restoration. The overall structure of the model is shown in the following figure:
+
+![img](https://user-images.githubusercontent.com/23252220/168281766-a0972bd3-243e-4fc7-baa5-e458ef0946ce.jpg)
+
+For a more detailed introduction to the model, and refer to the repo, you can view the following AI Studio project [link]([GPEN Blind Face Repair Model Reproduction - Paddle AI Studio (baidu.com)](https://aistudio.baidu.com/ The latest version of aistudio/projectdetail/3936241?contributionType=1)).
+
+
+
+
+## 2、Ready to work
+
+### 2.1 Dataset Preparation
+
+The GPEN model training set is the classic FFHQ face data set, with a total of 70,000 high-resolution 1024 x 1024 high-resolution face pictures, and the test set is the CELEBA-HQ data set, with a total of 2,000 high-resolution face pictures. For details, please refer to **Dataset URL:** [FFHQ](https://github.com/NVlabs/ffhq-dataset), [CELEBA-HQ](https://github.com/tkarras/progressive_growing_of_gans). The specific download links are given below:
+
+**Original dataset download address:**
+
+**FFHQ ：**           https://drive.google.com/drive/folders/1tZUcXDBeOibC6jcMCtgRRz67pzrAHeHL?usp=drive_open
+
+**CELEBA-HQ：** https://drive.google.com/drive/folders/0B4qLcYyJmiz0TXY1NG02bzZVRGs?resourcekey=0-arAVTUfW9KRhN-irJchVKQ&usp=sharing
+
+
+
+Since the original FFHQ dataset is too large, you can also download the 256-resolution FFHQ dataset from the following link:
+
+https://paddlegan.bj.bcebos.com/datasets/images256x256.tar
+
+
+
+**After downloading, the file organization is as follows**
+
+```
+|-- data/GPEN
+	|-- ffhq/images256x256/
+		|-- 00000
+			|-- 00000.png
+			|-- 00001.png
+			|-- ......
+			|-- 00999.png
+		|-- 01000
+			|-- ......
+		|-- ......
+		|-- 69000
+            |-- ......
+                |-- 69999.png
+	|-- test
+		|-- 2000张png图片
+```
+
+Please modify the dataroot parameters of dataset train and test in the configs/gpen_256_ffhq.yaml configuration file to your training set and test set path.
+
+
+
+### 2.2 Model preparation
+
+**Model parameter file and training log download address:**
+
+link：https://paddlegan.bj.bcebos.com/models/gpen.zip
+
+
+Download the model parameters and test images from the link and put them in the data/ folder in the project root directory. The specific file structure is as follows:
+
+
+```
+data/gpen/weights
+    |-- model_ir_se50.pdparams
+    |-- weight_pretrain.pdparams  
+data/gpen/lite_data
+```
+
+
+
+## 3、Start using
+
+### 3.1 model training
+
+Enter the following code in the console to start training：
+
+ ```shell
+ python tools/main.py -c configs/gpen_256_ffhq.yaml
+ ```
+
+The model only supports single-card training.
+
+Model training needs to use paddle2.3 and above, and wait for paddle to implement the second-order operator related functions of elementwise_pow. The paddle2.2.2 version can run normally, but the model cannot be successfully trained because some loss functions will calculate the wrong gradient. . If an error is reported during training, training is not supported for the time being. You can skip the training part and directly use the provided model parameters for testing. Model evaluation and testing can use paddle2.2.2 and above.
+
+
+
+### 3.2 Model evaluation
+
+When evaluating the model, enter the following code in the console, using the downloaded model parameters mentioned above:
+
+ ```shell
+python tools/main.py -c configs/gpen_256_ffhq.yaml -o dataset.test.amount=2000 --load data/gpen/weights/weight_pretrain.pdparams --evaluate-only
+ ```
+
+If you want to test on your own provided model, please modify the path after --load .
+
+
+
+### 3.3 Model prediction
+
+#### 3.3.1 Export generator weights
+
+After training, you need to use ``tools/extract_weight.py`` to extract the weights of the generator from the trained model (including the generator and discriminator) for inference to `applications/tools/gpen.py` to achieve Various applications of the GPEN model. Enter the following command to extract the weights of the generator:
+
+```bash
+python tools/extract_weight.py data/gpen/weights/weight_pretrain.pdparams --net-name g_ema --output data/gpen/weights/g_ema.pdparams
+```
+
+
+
+#### 3.3.2 Process a single image
+
+After extracting the weights of the generator, enter the following command to test the images under the --test_img path. Modifying the --seed parameter can generate different degraded images to show richer effects. You can modify the path after --test_img to any image you want to test. If no weight is provided after the --weight_path parameter, the trained model weights will be automatically downloaded for testing.
+
+```bash
+python applications/tools/gpen.py --test_img data/gpen/lite_data/15006.png --seed=100 --weight_path data/gpen/weights/g_ema.pdparams --model_type gpen-ffhq-256
+```
+
+The following are the sample images and the corresponding inpainted images, from left to right, the degraded image, the generated image, and the original clear image:
+
+<p align='center'>
+<img src="https://user-images.githubusercontent.com/23252220/168281788-39c08e86-2dc3-487f-987d-93489934c14c.png" height="256px" width='768px' >
+An example output is as follows:
+
+
+```
+result saved in : output_dir/gpen_predict.png
+        FID: 92.11730631094356
+        PSNR:19.014782083825743
+```
+
+
+
+## 4. Tipc
+
+### 4.1 Export the inference model
+
+```bash
+python tools/export_model.py -c configs/gpen_256_ffhq.yaml --inputs_size=1,3,256,256 --load data/gpen/weights/weight_pretrain.pdparams
+```
+
+The above command will generate the model structure file `gpenmodel_g_ema.pdmodel` and model weight files `gpenmodel_g_ema.pdiparams` and `gpenmodel_g_ema.pdiparams.info` files required for prediction, which are stored in the `inference_model/` directory. You can also modify the parameters after --load to the model parameter file you want to test.
+
+
+
+### 4.2 Inference with a prediction engine
+
+```bash
+python tools/inference.py --model_type GPEN --seed 100 -c configs/gpen_256_ffhq.yaml -o dataset.test.dataroot="./data/gpen/lite_data/" --output_path test_tipc/output/ --model_path inference_model/gpenmodel_g_ema
+```
+
+At the end of the inference, the repaired image generated by the model will be saved in the test_tipc/output/GPEN directory by default, and the FID value obtained by the test will be output in test_tipc/output/GPEN/metric.txt.
+
+
+The default output is as follows:
+
+```
+Metric fid: 187.0158
+```
+
+Note: Since the operation of degrading high-definition pictures has a certain degree of randomness, the results of each test will be different. In order to ensure that the test results are consistent, here I fixed the random seed, so that the same degradation operation is performed on the image for each test.
+
+
+
+### 4.3 Call the script to complete the training and push test in two steps
+
+To invoke the `lite_train_lite_infer` mode of the foot test base training prediction function, run:
+
+```shell
+# Corrected format of sh file
+sed -i 's/\r//' test_tipc/prepare.sh
+sed -i 's/\r//' test_tipc/test_train_inference_python.sh
+sed -i 's/\r//' test_tipc/common_func.sh
+# prepare data
+bash test_tipc/prepare.sh ./test_tipc/configs/GPEN/train_infer_python.txt 'lite_train_lite_infer'
+# run the test
+bash test_tipc/test_train_inference_python.sh ./test_tipc/configs/GPEN/train_infer_python.txt 'lite_train_lite_infer'
+```
+
+
+
+## 5、References
+
+```
+@misc{2021GAN,
+      title={GAN Prior Embedded Network for Blind Face Restoration in the Wild},
+      author={ Yang, T.  and  Ren, P.  and  Xie, X.  and  Zhang, L. },
+      year={2021},
+      archivePrefix={CVPR},
+      primaryClass={cs.CV}
+}
+```
+