Upload 5 files

scripts/crop_align_face.py

@@ -0,0 +1,205 @@
+"""
+brief: face alignment with FFHQ method (https://github.com/NVlabs/ffhq-dataset)
+author: lzhbrian (https://lzhbrian.me)
+link: https://gist.github.com/lzhbrian/bde87ab23b499dd02ba4f588258f57d5
+date: 2020.1.5
+note: code is heavily borrowed from
+    https://github.com/NVlabs/ffhq-dataset
+    http://dlib.net/face_landmark_detection.py.html
+requirements:
+    conda install Pillow numpy scipy
+    conda install -c conda-forge dlib
+    # download face landmark model from:
+    # http://dlib.net/files/shape_predictor_68_face_landmarks.dat.bz2
+"""
+
+import os
+import glob
+import numpy as np
+import PIL
+import PIL.Image
+import scipy
+import scipy.ndimage
+import argparse
+from basicsr.utils.download_util import load_file_from_url
+
+try:
+    import dlib
+except ImportError:
+    print('Please install dlib by running:' 'conda install -c conda-forge dlib')
+
+# download model from: http://dlib.net/files/shape_predictor_68_face_landmarks.dat.bz2
+shape_predictor_url = 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/shape_predictor_68_face_landmarks-fbdc2cb8.dat'
+ckpt_path = load_file_from_url(url=shape_predictor_url, 
+                                    model_dir='weights/dlib', progress=True, file_name=None)
+predictor = dlib.shape_predictor('weights/dlib/shape_predictor_68_face_landmarks-fbdc2cb8.dat')
+
+
+def get_landmark(filepath, only_keep_largest=True):
+    """get landmark with dlib
+    :return: np.array shape=(68, 2)
+    """
+    detector = dlib.get_frontal_face_detector()
+
+    img = dlib.load_rgb_image(filepath)
+    dets = detector(img, 1)
+
+    # Shangchen modified
+    print("\tNumber of faces detected: {}".format(len(dets)))
+    if only_keep_largest:
+        print('\tOnly keep the largest.')
+        face_areas = []
+        for k, d in enumerate(dets):
+            face_area = (d.right() - d.left()) * (d.bottom() - d.top())
+            face_areas.append(face_area)
+
+        largest_idx = face_areas.index(max(face_areas))
+        d = dets[largest_idx]
+        shape = predictor(img, d)
+        # print("Part 0: {}, Part 1: {} ...".format(
+        #     shape.part(0), shape.part(1)))
+    else:
+        for k, d in enumerate(dets):
+            # print("Detection {}: Left: {} Top: {} Right: {} Bottom: {}".format(
+            #     k, d.left(), d.top(), d.right(), d.bottom()))
+            # Get the landmarks/parts for the face in box d.
+            shape = predictor(img, d)
+            # print("Part 0: {}, Part 1: {} ...".format(
+            #     shape.part(0), shape.part(1)))
+
+    t = list(shape.parts())
+    a = []
+    for tt in t:
+        a.append([tt.x, tt.y])
+    lm = np.array(a)
+    # lm is a shape=(68,2) np.array
+    return lm
+
+def align_face(filepath, out_path):
+    """
+    :param filepath: str
+    :return: PIL Image
+    """
+    try:
+        lm = get_landmark(filepath)
+    except:
+        print('No landmark ...')
+        return
+
+    lm_chin = lm[0:17]  # left-right
+    lm_eyebrow_left = lm[17:22]  # left-right
+    lm_eyebrow_right = lm[22:27]  # left-right
+    lm_nose = lm[27:31]  # top-down
+    lm_nostrils = lm[31:36]  # top-down
+    lm_eye_left = lm[36:42]  # left-clockwise
+    lm_eye_right = lm[42:48]  # left-clockwise
+    lm_mouth_outer = lm[48:60]  # left-clockwise
+    lm_mouth_inner = lm[60:68]  # left-clockwise
+
+    # Calculate auxiliary vectors.
+    eye_left = np.mean(lm_eye_left, axis=0)
+    eye_right = np.mean(lm_eye_right, axis=0)
+    eye_avg = (eye_left + eye_right) * 0.5
+    eye_to_eye = eye_right - eye_left
+    mouth_left = lm_mouth_outer[0]
+    mouth_right = lm_mouth_outer[6]
+    mouth_avg = (mouth_left + mouth_right) * 0.5
+    eye_to_mouth = mouth_avg - eye_avg
+
+    # Choose oriented crop rectangle.
+    x = eye_to_eye - np.flipud(eye_to_mouth) * [-1, 1]
+    x /= np.hypot(*x)
+    x *= max(np.hypot(*eye_to_eye) * 2.0, np.hypot(*eye_to_mouth) * 1.8)
+    y = np.flipud(x) * [-1, 1]
+    c = eye_avg + eye_to_mouth * 0.1
+    quad = np.stack([c - x - y, c - x + y, c + x + y, c + x - y])
+    qsize = np.hypot(*x) * 2
+
+    # read image
+    img = PIL.Image.open(filepath)
+
+    output_size = 512
+    transform_size = 4096
+    enable_padding = False
+
+    # Shrink.
+    shrink = int(np.floor(qsize / output_size * 0.5))
+    if shrink > 1:
+        rsize = (int(np.rint(float(img.size[0]) / shrink)),
+                 int(np.rint(float(img.size[1]) / shrink)))
+        img = img.resize(rsize, PIL.Image.ANTIALIAS)
+        quad /= shrink
+        qsize /= shrink
+ 
+    # Crop.
+    border = max(int(np.rint(qsize * 0.1)), 3)
+    crop = (int(np.floor(min(quad[:, 0]))), int(np.floor(min(quad[:, 1]))),
+            int(np.ceil(max(quad[:, 0]))), int(np.ceil(max(quad[:, 1]))))
+    crop = (max(crop[0] - border, 0), max(crop[1] - border, 0),
+            min(crop[2] + border,
+                img.size[0]), min(crop[3] + border, img.size[1]))
+    if crop[2] - crop[0] < img.size[0] or crop[3] - crop[1] < img.size[1]:
+        img = img.crop(crop)
+        quad -= crop[0:2]
+
+    # Pad.
+    pad = (int(np.floor(min(quad[:, 0]))), int(np.floor(min(quad[:, 1]))),
+           int(np.ceil(max(quad[:, 0]))), int(np.ceil(max(quad[:, 1]))))
+    pad = (max(-pad[0] + border,
+               0), max(-pad[1] + border,
+                       0), max(pad[2] - img.size[0] + border,
+                               0), max(pad[3] - img.size[1] + border, 0))
+    if enable_padding and max(pad) > border - 4:
+        pad = np.maximum(pad, int(np.rint(qsize * 0.3)))
+        img = np.pad(
+            np.float32(img), ((pad[1], pad[3]), (pad[0], pad[2]), (0, 0)),
+            'reflect')
+        h, w, _ = img.shape
+        y, x, _ = np.ogrid[:h, :w, :1]
+        mask = np.maximum(
+            1.0 -
+            np.minimum(np.float32(x) / pad[0],
+                       np.float32(w - 1 - x) / pad[2]), 1.0 -
+            np.minimum(np.float32(y) / pad[1],
+                       np.float32(h - 1 - y) / pad[3]))
+        blur = qsize * 0.02
+        img += (scipy.ndimage.gaussian_filter(img, [blur, blur, 0]) -
+                img) * np.clip(mask * 3.0 + 1.0, 0.0, 1.0)
+        img += (np.median(img, axis=(0, 1)) - img) * np.clip(mask, 0.0, 1.0)
+        img = PIL.Image.fromarray(
+            np.uint8(np.clip(np.rint(img), 0, 255)), 'RGB')
+        quad += pad[:2]
+
+    img = img.transform((transform_size, transform_size), PIL.Image.QUAD,
+                        (quad + 0.5).flatten(), PIL.Image.BILINEAR)
+
+    if output_size < transform_size:
+        img = img.resize((output_size, output_size), PIL.Image.ANTIALIAS)
+
+    # Save aligned image.
+    # print('saveing: ', out_path)
+    img.save(out_path)
+
+    return img, np.max(quad[:, 0]) - np.min(quad[:, 0])
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-i', '--in_dir', type=str, default='./inputs/whole_imgs')
+    parser.add_argument('-o', '--out_dir', type=str, default='./inputs/cropped_faces')
+    args = parser.parse_args()
+
+    if args.out_dir.endswith('/'):  # solve when path ends with /
+        args.out_dir = args.out_dir[:-1]
+    dir_name = os.path.abspath(args.out_dir)
+    os.makedirs(dir_name, exist_ok=True)
+
+    img_list = sorted(glob.glob(os.path.join(args.in_dir, '*.[jpJP][pnPN]*[gG]')))
+    test_img_num = len(img_list)
+
+    for i, in_path in enumerate(img_list):
+        img_name = os.path.basename(in_path)
+        print(f'[{i+1}/{test_img_num}] Processing: {img_name}')
+        out_path = os.path.join(args.out_dir, in_path.split("/")[-1])        
+        out_path = out_path.replace('.jpg', '.png')
+        size_ = align_face(in_path, out_path)

scripts/download_pretrained_models.py

@@ -0,0 +1,44 @@
+import argparse
+import os
+from os import path as osp
+
+from basicsr.utils.download_util import load_file_from_url
+
+
+def download_pretrained_models(method, file_urls):
+    save_path_root = f'./weights/{method}'
+    os.makedirs(save_path_root, exist_ok=True)
+
+    for file_name, file_url in file_urls.items():
+        save_path = load_file_from_url(url=file_url, model_dir=save_path_root, progress=True, file_name=file_name)
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        'method',
+        type=str,
+        help=("Options: 'CodeFormer' 'facelib' 'dlib'. Set to 'all' to download all the models."))
+    args = parser.parse_args()
+
+    file_urls = {
+        'CodeFormer': {
+            'codeformer.pth': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth'
+        },
+        'facelib': {
+            # 'yolov5l-face.pth': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/yolov5l-face.pth',
+            'detection_Resnet50_Final.pth': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/detection_Resnet50_Final.pth',
+            'parsing_parsenet.pth': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/parsing_parsenet.pth'
+        },
+        'dlib': {
+            'mmod_human_face_detector-4cb19393.dat': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/mmod_human_face_detector-4cb19393.dat',
+            'shape_predictor_5_face_landmarks-c4b1e980.dat': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/shape_predictor_5_face_landmarks-c4b1e980.dat'
+        }
+    }
+
+    if args.method == 'all':
+        for method in file_urls.keys():
+            download_pretrained_models(method, file_urls[method])
+    else:
+        download_pretrained_models(args.method, file_urls[args.method])

scripts/download_pretrained_models_from_gdrive.py

@@ -0,0 +1,60 @@
+import argparse
+import os
+from os import path as osp
+
+# from basicsr.utils.download_util import download_file_from_google_drive
+import gdown
+
+
+def download_pretrained_models(method, file_ids):
+    save_path_root = f'./weights/{method}'
+    os.makedirs(save_path_root, exist_ok=True)
+
+    for file_name, file_id in file_ids.items():
+        file_url = 'https://drive.google.com/uc?id='+file_id
+        save_path = osp.abspath(osp.join(save_path_root, file_name))
+        if osp.exists(save_path):
+            user_response = input(f'{file_name} already exist. Do you want to cover it? Y/N\n')
+            if user_response.lower() == 'y':
+                print(f'Covering {file_name} to {save_path}')
+                gdown.download(file_url, save_path, quiet=False)
+                # download_file_from_google_drive(file_id, save_path)
+            elif user_response.lower() == 'n':
+                print(f'Skipping {file_name}')
+            else:
+                raise ValueError('Wrong input. Only accepts Y/N.')
+        else:
+            print(f'Downloading {file_name} to {save_path}')
+            gdown.download(file_url, save_path, quiet=False)
+            # download_file_from_google_drive(file_id, save_path)
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        'method',
+        type=str,
+        help=("Options: 'CodeFormer' 'facelib'. Set to 'all' to download all the models."))
+    args = parser.parse_args()
+
+    # file name: file id
+    # 'dlib': {
+    #     'mmod_human_face_detector-4cb19393.dat': '1qD-OqY8M6j4PWUP_FtqfwUPFPRMu6ubX',
+    #     'shape_predictor_5_face_landmarks-c4b1e980.dat': '1vF3WBUApw4662v9Pw6wke3uk1qxnmLdg',
+    #     'shape_predictor_68_face_landmarks-fbdc2cb8.dat': '1tJyIVdCHaU6IDMDx86BZCxLGZfsWB8yq'
+    # }
+    file_ids = {
+        'CodeFormer': {
+            'codeformer.pth': '1v_E_vZvP-dQPF55Kc5SRCjaKTQXDz-JB'
+        },
+        'facelib': {
+            'yolov5l-face.pth': '131578zMA6B2x8VQHyHfa6GEPtulMCNzV',
+            'parsing_parsenet.pth': '16pkohyZZ8ViHGBk3QtVqxLZKzdo466bK'
+        }
+    }
+
+    if args.method == 'all':
+        for method in file_ids.keys():
+            download_pretrained_models(method, file_ids[method])
+    else:
+        download_pretrained_models(args.method, file_ids[args.method])

scripts/generate_latent_gt.py

@@ -0,0 +1,67 @@
+import argparse
+import glob
+import numpy as np
+import os
+import cv2
+import torch
+from torchvision.transforms.functional import normalize
+from basicsr.utils import imwrite, img2tensor, tensor2img
+
+from basicsr.utils.registry import ARCH_REGISTRY
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-i', '--test_path', type=str, default='datasets/ffhq/ffhq_512')
+    parser.add_argument('-o', '--save_root', type=str, default='./experiments/pretrained_models/vqgan')
+    parser.add_argument('--codebook_size', type=int, default=1024)
+    parser.add_argument('--ckpt_path', type=str, default='./experiments/pretrained_models/vqgan/net_g.pth')
+    args = parser.parse_args()
+
+    if args.save_root.endswith('/'):  # solve when path ends with /
+        args.save_root = args.save_root[:-1]
+    dir_name = os.path.abspath(args.save_root)
+    os.makedirs(dir_name, exist_ok=True)
+
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    test_path = args.test_path
+    save_root = args.save_root
+    ckpt_path = args.ckpt_path
+    codebook_size = args.codebook_size
+
+    vqgan = ARCH_REGISTRY.get('VQAutoEncoder')(512, 64, [1, 2, 2, 4, 4, 8], 'nearest',
+                                                codebook_size=codebook_size).to(device)
+    checkpoint = torch.load(ckpt_path)['params_ema']
+
+    vqgan.load_state_dict(checkpoint)
+    vqgan.eval()
+
+    sum_latent = np.zeros((codebook_size)).astype('float64')
+    size_latent = 16
+    latent = {}
+    latent['orig'] = {}
+    latent['hflip'] = {}
+    for i in ['orig', 'hflip']:
+    # for i in ['hflip']:
+        for img_path in sorted(glob.glob(os.path.join(test_path, '*.[jp][pn]g'))):
+            img_name = os.path.basename(img_path)
+            img = cv2.imread(img_path)
+            if i == 'hflip':
+                cv2.flip(img, 1, img)
+            img = img2tensor(img / 255., bgr2rgb=True, float32=True)
+            normalize(img, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
+            img = img.unsqueeze(0).to(device)
+            with torch.no_grad():
+                # output = net(img)[0]
+                x, feat_dict = vqgan.encoder(img, True)
+                x, _, log = vqgan.quantize(x)
+            # del output
+            torch.cuda.empty_cache()
+
+            min_encoding_indices = log['min_encoding_indices']
+            min_encoding_indices = min_encoding_indices.view(size_latent,size_latent)
+            latent[i][img_name[:-4]] = min_encoding_indices.cpu().numpy()
+            print(img_name, latent[i][img_name[:-4]].shape)
+
+    latent_save_path = os.path.join(save_root, f'latent_gt_code{codebook_size}.pth')
+    torch.save(latent, latent_save_path)
+    print(f'\nLatent GT code are saved in {save_root}')

scripts/inference_vqgan.py

@@ -0,0 +1,59 @@
+import argparse
+import glob
+import numpy as np
+import os
+import cv2
+import torch
+from torchvision.transforms.functional import normalize
+from basicsr.utils import imwrite, img2tensor, tensor2img
+
+from basicsr.utils.registry import ARCH_REGISTRY
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-i', '--test_path', type=str, default='datasets/ffhq/ffhq_512')
+    parser.add_argument('-o', '--save_root', type=str, default='./results/vqgan_rec')
+    parser.add_argument('--codebook_size', type=int, default=1024)
+    parser.add_argument('--ckpt_path', type=str, default='./experiments/pretrained_models/vqgan/net_g.pth')
+    args = parser.parse_args()
+
+    if args.save_root.endswith('/'):  # solve when path ends with /
+        args.save_root = args.save_root[:-1]
+    dir_name = os.path.abspath(args.save_root)
+    os.makedirs(dir_name, exist_ok=True)
+
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    test_path = args.test_path
+    save_root = args.save_root
+    ckpt_path = args.ckpt_path
+    codebook_size = args.codebook_size
+
+    vqgan = ARCH_REGISTRY.get('VQAutoEncoder')(512, 64, [1, 2, 2, 4, 4, 8], 'nearest',
+                                                codebook_size=codebook_size).to(device)
+    checkpoint = torch.load(ckpt_path)['params_ema']
+
+    vqgan.load_state_dict(checkpoint)
+    vqgan.eval()
+
+    for img_path in sorted(glob.glob(os.path.join(test_path, '*.[jp][pn]g'))):
+        img_name = os.path.basename(img_path)
+        print(img_name)
+        img = cv2.imread(img_path)
+        img = img2tensor(img / 255., bgr2rgb=True, float32=True)
+        normalize(img, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
+        img = img.unsqueeze(0).to(device)
+        with torch.no_grad():
+            output = vqgan(img)[0]
+            output = tensor2img(output, min_max=[-1,1])
+            img = tensor2img(img, min_max=[-1,1])
+            restored_img = np.concatenate([img, output], axis=1)
+            restored_img = output
+        del output
+        torch.cuda.empty_cache()
+
+        path = os.path.splitext(os.path.join(save_root, img_name))[0]
+        save_path = f'{path}.png'
+        imwrite(restored_img, save_path)
+
+    print(f'\nAll results are saved in {save_root}')
+