run_HR

Author: raywzy

Face_Detection/align_warp_back_multiple_dlib_HR.py

@@ -0,0 +1,184 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import torch
+import numpy as np
+import skimage.io as io
+
+# from FaceSDK.face_sdk import FaceDetection
+# from face_sdk import FaceDetection
+import matplotlib.pyplot as plt
+from matplotlib.patches import Rectangle
+from skimage.transform import SimilarityTransform
+from skimage.transform import warp
+from PIL import Image
+import torch.nn.functional as F
+import torchvision as tv
+import torchvision.utils as vutils
+import time
+import cv2
+import os
+from skimage import img_as_ubyte
+import json
+import argparse
+import dlib
+
+
+def _standard_face_pts():
+    pts = (
+        np.array([196.0, 226.0, 316.0, 226.0, 256.0, 286.0, 220.0, 360.4, 292.0, 360.4], np.float32) / 256.0
+        - 1.0
+    )
+
+    return np.reshape(pts, (5, 2))
+
+
+def _origin_face_pts():
+    pts = np.array([196.0, 226.0, 316.0, 226.0, 256.0, 286.0, 220.0, 360.4, 292.0, 360.4], np.float32)
+
+    return np.reshape(pts, (5, 2))
+
+
+def get_landmark(face_landmarks, id):
+    part = face_landmarks.part(id)
+    x = part.x
+    y = part.y
+
+    return (x, y)
+
+
+def search(face_landmarks):
+
+    x1, y1 = get_landmark(face_landmarks, 36)
+    x2, y2 = get_landmark(face_landmarks, 39)
+    x3, y3 = get_landmark(face_landmarks, 42)
+    x4, y4 = get_landmark(face_landmarks, 45)
+
+    x_nose, y_nose = get_landmark(face_landmarks, 30)
+
+    x_left_mouth, y_left_mouth = get_landmark(face_landmarks, 48)
+    x_right_mouth, y_right_mouth = get_landmark(face_landmarks, 54)
+
+    x_left_eye = int((x1 + x2) / 2)
+    y_left_eye = int((y1 + y2) / 2)
+    x_right_eye = int((x3 + x4) / 2)
+    y_right_eye = int((y3 + y4) / 2)
+
+    results = np.array(
+        [
+            [x_left_eye, y_left_eye],
+            [x_right_eye, y_right_eye],
+            [x_nose, y_nose],
+            [x_left_mouth, y_left_mouth],
+            [x_right_mouth, y_right_mouth],
+        ]
+    )
+
+    return results
+
+
+def compute_transformation_matrix(img, landmark, normalize, target_face_scale=1.0):
+
+    std_pts = _standard_face_pts()  # [-1,1]
+    target_pts = (std_pts * target_face_scale + 1) / 2 * 512.0
+
+    # print(target_pts)
+
+    h, w, c = img.shape
+    if normalize == True:
+        landmark[:, 0] = landmark[:, 0] / h * 2 - 1.0
+        landmark[:, 1] = landmark[:, 1] / w * 2 - 1.0
+
+    # print(landmark)
+
+    affine = SimilarityTransform()
+
+    affine.estimate(target_pts, landmark)
+
+    return affine.params
+
+
+def show_detection(image, box, landmark):
+    plt.imshow(image)
+    print(box[2] - box[0])
+    plt.gca().add_patch(
+        Rectangle(
+            (box[1], box[0]), box[2] - box[0], box[3] - box[1], linewidth=1, edgecolor="r", facecolor="none"
+        )
+    )
+    plt.scatter(landmark[0][0], landmark[0][1])
+    plt.scatter(landmark[1][0], landmark[1][1])
+    plt.scatter(landmark[2][0], landmark[2][1])
+    plt.scatter(landmark[3][0], landmark[3][1])
+    plt.scatter(landmark[4][0], landmark[4][1])
+    plt.show()
+
+
+def affine2theta(affine, input_w, input_h, target_w, target_h):
+    # param = np.linalg.inv(affine)
+    param = affine
+    theta = np.zeros([2, 3])
+    theta[0, 0] = param[0, 0] * input_h / target_h
+    theta[0, 1] = param[0, 1] * input_w / target_h
+    theta[0, 2] = (2 * param[0, 2] + param[0, 0] * input_h + param[0, 1] * input_w) / target_h - 1
+    theta[1, 0] = param[1, 0] * input_h / target_w
+    theta[1, 1] = param[1, 1] * input_w / target_w
+    theta[1, 2] = (2 * param[1, 2] + param[1, 0] * input_h + param[1, 1] * input_w) / target_w - 1
+    return theta
+
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--url", type=str, default="/home/jingliao/ziyuwan/celebrities", help="input")
+    parser.add_argument(
+        "--save_url", type=str, default="/home/jingliao/ziyuwan/celebrities_detected_face_reid", help="output"
+    )
+    opts = parser.parse_args()
+
+    url = opts.url
+    save_url = opts.save_url
+
+    ### If the origin url is None, then we don't need to reid the origin image
+
+    os.makedirs(url, exist_ok=True)
+    os.makedirs(save_url, exist_ok=True)
+
+    face_detector = dlib.get_frontal_face_detector()
+    landmark_locator = dlib.shape_predictor("shape_predictor_68_face_landmarks.dat")
+
+    count = 0
+
+    map_id = {}
+    for x in os.listdir(url):
+        img_url = os.path.join(url, x)
+        pil_img = Image.open(img_url).convert("RGB")
+
+        image = np.array(pil_img)
+
+        start = time.time()
+        faces = face_detector(image)
+        done = time.time()
+
+        if len(faces) == 0:
+            print("Warning: There is no face in %s" % (x))
+            continue
+
+        print(len(faces))
+
+        if len(faces) > 0:
+            for face_id in range(len(faces)):
+                current_face = faces[face_id]
+                face_landmarks = landmark_locator(image, current_face)
+                current_fl = search(face_landmarks)
+
+                affine = compute_transformation_matrix(image, current_fl, False, target_face_scale=1.3)
+                aligned_face = warp(image, affine, output_shape=(512, 512, 3))
+                img_name = x[:-4] + "_" + str(face_id + 1)
+                io.imsave(os.path.join(save_url, img_name + ".png"), img_as_ubyte(aligned_face))
+
+        count += 1
+
+        if count % 1000 == 0:
+            print("%d have finished ..." % (count))
+

Face_Detection/detect_all_dlib_HR.py

@@ -97,7 +97,7 @@ def compute_latent_vector_size(self, opt):
         else:
             raise ValueError("opt.num_upsampling_layers [%s] not recognized" % opt.num_upsampling_layers)
 
-        sw = opt.crop_size // (2 ** num_up_layers)
+        sw = opt.load_size // (2 ** num_up_layers)
         sh = round(sw / opt.aspect_ratio)
 
         return sw, sh

Face_Enhancement/models/networks/generator.py

@@ -97,4 +97,4 @@ def initialize(self):
         self.parser.add_argument(
             "--Scratch_and_Quality_restore", action="store_true", help="For scratched images"
         )
-
+        self.parser.add_argument("--HR", action='store_true',help='Large input size with scratches')

Global/options/test_options.py

@@ -86,6 +86,11 @@ def parameter_set(opt):
         opt.name = "mapping_scratch"
         opt.load_pretrainA = os.path.join(opt.checkpoints_dir, "VAE_A_quality")
         opt.load_pretrainB = os.path.join(opt.checkpoints_dir, "VAE_B_scratch")
+        if opt.HR:
+            opt.mapping_exp = 1
+            opt.inference_optimize = True
+            opt.mask_dilation = 3
+            opt.name = "mapping_Patch_Attention"
 
 
 if __name__ == "__main__":
@@ -135,6 +140,11 @@ def parameter_set(opt):
         if opt.NL_use_mask:
             mask_name = mask_loader[i]
             mask = Image.open(os.path.join(opt.test_mask, mask_name)).convert("RGB")
+            if opt.mask_dilation!=0:
+                kernel=np.ones((3,3),np.uint8)
+                mask=np.array(mask)
+                mask=cv2.dilate(mask,kernel,iterations=opt.mask_dilation)
+                mask=Image.fromarray(mask.astype('uint8'))
             origin = input
             input = irregular_hole_synthesize(input, mask)
             mask = mask_transform(mask)

Global/test.py

@@ -23,6 +23,10 @@ The code originates from our research project and the aim is to demonstrate the
 **We are improving the algorithm so as to process high resolution photos. It takes time and please stay tuned.**
 
 ## News
+The framework now supports the restoration of high-resolution input.
+
+<img src='imgs/HR.png'>
+
 Training code is available and welcome to have a try and learn the training details. 
 
 You can now play with our [Colab](https://colab.research.google.com/drive/1NEm6AsybIiC5TwTU_4DqDkQO0nFRB-uA?usp=sharing) and try it on your photos. 
@@ -101,6 +105,16 @@ python run.py --input_folder [test_image_folder_path] \
               --with_scratch
 ```
 
+For high-resolution images with scratches:
+
+```
+python run.py --input_folder [test_image_folder_path] \
+              --output_folder [output_path] \
+              --GPU 0 \
+              --with_scratch \
+              --HR
+```
+
 Note: Please try to use the absolute path. The final results will be saved in `./output_path/final_output/`. You could also check the produced results of different steps in `output_path`.
 
 ### 2) Scratch Detection
@@ -132,8 +146,8 @@ python test.py --Scratch_and_Quality_restore \
                --outputs_dir [output_path]
 
 python test.py --Quality_restore \
- --test_input [test_image_folder_path] \
- --outputs_dir [output_path]
+               --test_input [test_image_folder_path] \
+               --outputs_dir [output_path]
 ```
 
 <img src='imgs/global.png'>
@@ -203,14 +217,17 @@ Traing the mapping with scraches:
 python train_mapping.py --no_TTUR --NL_res --random_hole --use_SN --correlation_renormalize --training_dataset mapping --NL_use_mask --NL_fusion_method combine --non_local Setting_42 --use_v2_degradation --use_vae_which_epoch 200 --continue_train --name mapping_scratch --label_nc 0 --loadSize 256 --fineSize 256 --dataroot [your_data_folder] --no_instance --resize_or_crop crop_only --batchSize 36 --no_html --gpu_ids 0,1,2,3 --nThreads 8 --load_pretrainA [ckpt_of_domainA_SR_old_photos] --load_pretrainB [ckpt_of_domainB_old_photos] --l2_feat 60 --n_downsample_global 3 --mc 64 --k_size 4 --start_r 1 --mapping_n_block 6 --map_mc 512 --use_l1_feat --niter 150 --niter_decay 100 --outputs_dir [your_output_folder] --checkpoints_dir [your_ckpt_folder] --irregular_mask [absolute_path_of_mask_file]
 ```
 
-
+Traing the mapping with scraches (Multi-Scale Patch Attention for HR input):
+```
+python train_mapping.py --no_TTUR --NL_res --random_hole --use_SN --correlation_renormalize --training_dataset mapping --NL_use_mask --NL_fusion_method combine --non_local Setting_42 --use_v2_degradation --use_vae_which_epoch 200 --continue_train --name mapping_Pathc_Attention --label_nc 0 --loadSize 256 --fineSize 256 --dataroot [your_data_folder] --no_instance --resize_or_crop crop_only --batchSize 36 --no_html --gpu_ids 0,1,2,3 --nThreads 8 --load_pretrainA [ckpt_of_domainA_SR_old_photos] --load_pretrainB [ckpt_of_domainB_old_photos] --l2_feat 60 --n_downsample_global 3 --mc 64 --k_size 4 --start_r 1 --mapping_n_block 6 --map_mc 512 --use_l1_feat --niter 150 --niter_decay 100 --outputs_dir [your_output_folder] --checkpoints_dir [your_ckpt_folder] --irregular_mask [absolute_path_of_mask_file] --mapping_exp 1
+```
 
 ## To Do
 - [x] Clean testing code
 - [x] Release pretrained model
 - [x] Collab demo
-- [ ] Replace face detection module (dlib) with RetinaFace
 - [x] Release training code
+- [x] Processing of high-resolution input
 
 
 ## Citation