GFPGAN-V1.3-whole-image / scripts /parse_landmark.py
tjxj
Add application file
2a8871c
raw
history blame
3.52 kB
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)