update

2022-04-24 15:44:47 +08:00
parent 99ed65aaa3
commit 29d8914c0a
138 changed files with 24864 additions and 353 deletions
@@ -0,0 +1,164 @@
+import argparse
+import math
+import torch
+
+from gfpgan.archs.gfpganv1_clean_arch import GFPGANv1Clean
+
+
+def modify_checkpoint(checkpoint_bilinear, checkpoint_clean):
+    for ori_k, ori_v in checkpoint_bilinear.items():
+        if 'stylegan_decoder' in ori_k:
+            if 'style_mlp' in ori_k:  # style_mlp_layers
+                lr_mul = 0.01
+                prefix, name, idx, var = ori_k.split('.')
+                idx = (int(idx) * 2) - 1
+                crt_k = f'{prefix}.{name}.{idx}.{var}'
+                if var == 'weight':
+                    _, c_in = ori_v.size()
+                    scale = (1 / math.sqrt(c_in)) * lr_mul
+                    crt_v = ori_v * scale * 2**0.5
+                else:
+                    crt_v = ori_v * lr_mul * 2**0.5
+                checkpoint_clean[crt_k] = crt_v
+            elif 'modulation' in ori_k:  # modulation in StyleConv
+                lr_mul = 1
+                crt_k = ori_k
+                var = ori_k.split('.')[-1]
+                if var == 'weight':
+                    _, c_in = ori_v.size()
+                    scale = (1 / math.sqrt(c_in)) * lr_mul
+                    crt_v = ori_v * scale
+                else:
+                    crt_v = ori_v * lr_mul
+                checkpoint_clean[crt_k] = crt_v
+            elif 'style_conv' in ori_k:
+                # StyleConv in style_conv1 and style_convs
+                if 'activate' in ori_k:  # FusedLeakyReLU
+                    # eg. style_conv1.activate.bias
+                    # eg. style_convs.13.activate.bias
+                    split_rlt = ori_k.split('.')
+                    if len(split_rlt) == 4:
+                        prefix, name, _, var = split_rlt
+                        crt_k = f'{prefix}.{name}.{var}'
+                    elif len(split_rlt) == 5:
+                        prefix, name, idx, _, var = split_rlt
+                        crt_k = f'{prefix}.{name}.{idx}.{var}'
+                    crt_v = ori_v * 2**0.5  # 2**0.5 used in FusedLeakyReLU
+                    c = crt_v.size(0)
+                    checkpoint_clean[crt_k] = crt_v.view(1, c, 1, 1)
+                elif 'modulated_conv' in ori_k:
+                    # eg. style_conv1.modulated_conv.weight
+                    # eg. style_convs.13.modulated_conv.weight
+                    _, c_out, c_in, k1, k2 = ori_v.size()
+                    scale = 1 / math.sqrt(c_in * k1 * k2)
+                    crt_k = ori_k
+                    checkpoint_clean[crt_k] = ori_v * scale
+                elif 'weight' in ori_k:
+                    crt_k = ori_k
+                    checkpoint_clean[crt_k] = ori_v * 2**0.5
+            elif 'to_rgb' in ori_k:  # StyleConv in to_rgb1 and to_rgbs
+                if 'modulated_conv' in ori_k:
+                    # eg. to_rgb1.modulated_conv.weight
+                    # eg. to_rgbs.5.modulated_conv.weight
+                    _, c_out, c_in, k1, k2 = ori_v.size()
+                    scale = 1 / math.sqrt(c_in * k1 * k2)
+                    crt_k = ori_k
+                    checkpoint_clean[crt_k] = ori_v * scale
+                else:
+                    crt_k = ori_k
+                    checkpoint_clean[crt_k] = ori_v
+            else:
+                crt_k = ori_k
+                checkpoint_clean[crt_k] = ori_v
+            # end of 'stylegan_decoder'
+        elif 'conv_body_first' in ori_k or 'final_conv' in ori_k:
+            # key name
+            name, _, var = ori_k.split('.')
+            crt_k = f'{name}.{var}'
+            # weight and bias
+            if var == 'weight':
+                c_out, c_in, k1, k2 = ori_v.size()
+                scale = 1 / math.sqrt(c_in * k1 * k2)
+                checkpoint_clean[crt_k] = ori_v * scale * 2**0.5
+            else:
+                checkpoint_clean[crt_k] = ori_v * 2**0.5
+        elif 'conv_body' in ori_k:
+            if 'conv_body_up' in ori_k:
+                ori_k = ori_k.replace('conv2.weight', 'conv2.1.weight')
+                ori_k = ori_k.replace('skip.weight', 'skip.1.weight')
+            name1, idx1, name2, _, var = ori_k.split('.')
+            crt_k = f'{name1}.{idx1}.{name2}.{var}'
+            if name2 == 'skip':
+                c_out, c_in, k1, k2 = ori_v.size()
+                scale = 1 / math.sqrt(c_in * k1 * k2)
+                checkpoint_clean[crt_k] = ori_v * scale / 2**0.5
+            else:
+                if var == 'weight':
+                    c_out, c_in, k1, k2 = ori_v.size()
+                    scale = 1 / math.sqrt(c_in * k1 * k2)
+                    checkpoint_clean[crt_k] = ori_v * scale
+                else:
+                    checkpoint_clean[crt_k] = ori_v
+                if 'conv1' in ori_k:
+                    checkpoint_clean[crt_k] *= 2**0.5
+        elif 'toRGB' in ori_k:
+            crt_k = ori_k
+            if 'weight' in ori_k:
+                c_out, c_in, k1, k2 = ori_v.size()
+                scale = 1 / math.sqrt(c_in * k1 * k2)
+                checkpoint_clean[crt_k] = ori_v * scale
+            else:
+                checkpoint_clean[crt_k] = ori_v
+        elif 'final_linear' in ori_k:
+            crt_k = ori_k
+            if 'weight' in ori_k:
+                _, c_in = ori_v.size()
+                scale = 1 / math.sqrt(c_in)
+                checkpoint_clean[crt_k] = ori_v * scale
+            else:
+                checkpoint_clean[crt_k] = ori_v
+        elif 'condition' in ori_k:
+            crt_k = ori_k
+            if '0.weight' in ori_k:
+                c_out, c_in, k1, k2 = ori_v.size()
+                scale = 1 / math.sqrt(c_in * k1 * k2)
+                checkpoint_clean[crt_k] = ori_v * scale * 2**0.5
+            elif '0.bias' in ori_k:
+                checkpoint_clean[crt_k] = ori_v * 2**0.5
+            elif '2.weight' in ori_k:
+                c_out, c_in, k1, k2 = ori_v.size()
+                scale = 1 / math.sqrt(c_in * k1 * k2)
+                checkpoint_clean[crt_k] = ori_v * scale
+            elif '2.bias' in ori_k:
+                checkpoint_clean[crt_k] = ori_v
+
+    return checkpoint_clean
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--ori_path', type=str, help='Path to the original model')
+    parser.add_argument('--narrow', type=float, default=1)
+    parser.add_argument('--channel_multiplier', type=float, default=2)
+    parser.add_argument('--save_path', type=str)
+    args = parser.parse_args()
+
+    ori_ckpt = torch.load(args.ori_path)['params_ema']
+
+    net = GFPGANv1Clean(
+        512,
+        num_style_feat=512,
+        channel_multiplier=args.channel_multiplier,
+        decoder_load_path=None,
+        fix_decoder=False,
+        # for stylegan decoder
+        num_mlp=8,
+        input_is_latent=True,
+        different_w=True,
+        narrow=args.narrow,
+        sft_half=True)
+    crt_ckpt = net.state_dict()
+
+    crt_ckpt = modify_checkpoint(ori_ckpt, crt_ckpt)
+    print(f'Save to {args.save_path}.')
+    torch.save(dict(params_ema=crt_ckpt), args.save_path, _use_new_zipfile_serialization=False)
@@ -0,0 +1,85 @@
+import cv2
+import json
+import numpy as np
+import os
+import torch
+from basicsr.utils import FileClient, imfrombytes
+from collections import OrderedDict
+
+# ---------------------------- This script is used to parse facial landmarks ------------------------------------- #
+# Configurations
+save_img = False
+scale = 0.5  # 0.5 for official FFHQ (512x512), 1 for others
+enlarge_ratio = 1.4  # only for eyes
+json_path = 'ffhq-dataset-v2.json'
+face_path = 'datasets/ffhq/ffhq_512.lmdb'
+save_path = './FFHQ_eye_mouth_landmarks_512.pth'
+
+print('Load JSON metadata...')
+# use the official json file in FFHQ dataset
+with open(json_path, 'rb') as f:
+    json_data = json.load(f, object_pairs_hook=OrderedDict)
+
+print('Open LMDB file...')
+# read ffhq images
+file_client = FileClient('lmdb', db_paths=face_path)
+with open(os.path.join(face_path, 'meta_info.txt')) as fin:
+    paths = [line.split('.')[0] for line in fin]
+
+save_dict = {}
+
+for item_idx, item in enumerate(json_data.values()):
+    print(f'\r{item_idx} / {len(json_data)}, {item["image"]["file_path"]} ', end='', flush=True)
+
+    # parse landmarks
+    lm = np.array(item['image']['face_landmarks'])
+    lm = lm * scale
+
+    item_dict = {}
+    # get image
+    if save_img:
+        img_bytes = file_client.get(paths[item_idx])
+        img = imfrombytes(img_bytes, float32=True)
+
+    # get landmarks for each component
+    map_left_eye = list(range(36, 42))
+    map_right_eye = list(range(42, 48))
+    map_mouth = list(range(48, 68))
+
+    # eye_left
+    mean_left_eye = np.mean(lm[map_left_eye], 0)  # (x, y)
+    half_len_left_eye = np.max((np.max(np.max(lm[map_left_eye], 0) - np.min(lm[map_left_eye], 0)) / 2, 16))
+    item_dict['left_eye'] = [mean_left_eye[0], mean_left_eye[1], half_len_left_eye]
+    # mean_left_eye[0] = 512 - mean_left_eye[0]  # for testing flip
+    half_len_left_eye *= enlarge_ratio
+    loc_left_eye = np.hstack((mean_left_eye - half_len_left_eye + 1, mean_left_eye + half_len_left_eye)).astype(int)
+    if save_img:
+        eye_left_img = img[loc_left_eye[1]:loc_left_eye[3], loc_left_eye[0]:loc_left_eye[2], :]
+        cv2.imwrite(f'tmp/{item_idx:08d}_eye_left.png', eye_left_img * 255)
+
+    # eye_right
+    mean_right_eye = np.mean(lm[map_right_eye], 0)
+    half_len_right_eye = np.max((np.max(np.max(lm[map_right_eye], 0) - np.min(lm[map_right_eye], 0)) / 2, 16))
+    item_dict['right_eye'] = [mean_right_eye[0], mean_right_eye[1], half_len_right_eye]
+    # mean_right_eye[0] = 512 - mean_right_eye[0]  # # for testing flip
+    half_len_right_eye *= enlarge_ratio
+    loc_right_eye = np.hstack(
+        (mean_right_eye - half_len_right_eye + 1, mean_right_eye + half_len_right_eye)).astype(int)
+    if save_img:
+        eye_right_img = img[loc_right_eye[1]:loc_right_eye[3], loc_right_eye[0]:loc_right_eye[2], :]
+        cv2.imwrite(f'tmp/{item_idx:08d}_eye_right.png', eye_right_img * 255)
+
+    # mouth
+    mean_mouth = np.mean(lm[map_mouth], 0)
+    half_len_mouth = np.max((np.max(np.max(lm[map_mouth], 0) - np.min(lm[map_mouth], 0)) / 2, 16))
+    item_dict['mouth'] = [mean_mouth[0], mean_mouth[1], half_len_mouth]
+    # mean_mouth[0] = 512 - mean_mouth[0]  # for testing flip
+    loc_mouth = np.hstack((mean_mouth - half_len_mouth + 1, mean_mouth + half_len_mouth)).astype(int)
+    if save_img:
+        mouth_img = img[loc_mouth[1]:loc_mouth[3], loc_mouth[0]:loc_mouth[2], :]
+        cv2.imwrite(f'tmp/{item_idx:08d}_mouth.png', mouth_img * 255)
+
+    save_dict[f'{item_idx:08d}'] = item_dict
+
+print('Save...')
+torch.save(save_dict, save_path)