Dev-Mrha / DualPriorsCorrection

Combining Generative and Geometry Priors for Wide-Angle Portrait Correction

Lan Yao¹, Chaofeng Chen², Xiaoming Li¹, Zifei Yan¹, Wangmeng Zuo^1,3

¹ Harbin Institute of Technology, ² Nanyang Technological University, ³ Pazhou Lab, Huangpu

we introduce a framework that leverages generative and geometry priors to rectify wide-angle distortions in faces and backgrounds. Experiments demonstrate that our approach outperforms previous methods by a large margin, excelling not only in quantitative measures such as line straightness and shape consistency metrics but also in terms of perceptual visual quality

Given a wide-angle image, LineCNet and FaceCNet deal with the correction in background and face region seperatly, then face fusion is done by a post-process.

TODO

• Release the source code.
• [] Release the checkpoint model.

Installation

All dependencies for defining the environment are provided in environment.yaml.

Evaluate

Because of our post process, ShapeAcc and Landmark Distance cannot directly generated by an output flow, which we calculate by using a face detection tool to get each faces keypoints and calculate according to ./evaluate.py.

Step 0: download the weights

checkpoints should be put in ./pretrained_models/

Step 1: download the dataset

public dataset released by work Practical Wide-Angle Portraits Correction with Deep Structured Models

Step 2: generate the output of test dataset

python evaluate.py --option generate --test-dir ../test/ --device cuda:0 --e4e-path ./pretrained_models/e4e_best_model.pth --linenet-path ./pretrained_models/linenet.pt --facenet-path ./pretrained_models/facenet.pt

The generated output is before post-process part, which means for each input xx.jpg, you now have a xx_out.jpg and xx_mask.jpg that you may need a image inpainting alogorithm to generate the final output. You can use lama for this part.

Step 3: Face detection to get the keypoints

Use any face detection tool to get face keypoints landmark for each faces in test images and save as file xxx_pred_landmark.json, donot mind the face sequence cause it will be sorted during calculating. You may need megvii face landmark detection to get the same 81 landmarks as test dataset(except 'left_eye_pupil' and 'left_eye_pupil' landmark when using api).

Step 4: Evaluate

python evaluate.py --option evaluate --test-dir ../test/ --device cuda:0 --e4e-path ./pretrained_models/e4e_best_model.pth --linenet-path ./pretrained_models/linenet.pt --facenet-path ./pretrained_models/facenet.pt

Training

Data for finetune e4e encoder:

• generated distorted face data using FFHQ and CelebA-HQ

Training Data for FaceCNet:

• generated distorted face data using FFHQ and CelebA-HQ
• cropped face from public datset

python train_facecnet.py

Training Data for lineCNet:

• public dataset

python train_linecnet.py

Acknowledgement

This project is built based on encoder4editing.