import cv2 import numpy as np import os import torch from skimage import transform as trans from basicsr.utils import imwrite try: import dlib except ImportError: print('Please install dlib before testing face restoration. Reference:　https://github.com/davisking/dlib') class FaceRestorationHelper(object): """Helper for the face restoration pipeline.""" def __init__(self, upscale_factor, face_size=512): self.upscale_factor = upscale_factor self.face_size = (face_size, face_size) # standard 5 landmarks for FFHQ faces with 1024 x 1024 self.face_template = np.array([[686.77227723, 488.62376238], [586.77227723, 493.59405941], [337.91089109, 488.38613861], [437.95049505, 493.51485149], [513.58415842, 678.5049505]]) self.face_template = self.face_template / (1024 // face_size) # for estimation the 2D similarity transformation self.similarity_trans = trans.SimilarityTransform() self.all_landmarks_5 = [] self.all_landmarks_68 = [] self.affine_matrices = [] self.inverse_affine_matrices = [] self.cropped_faces = [] self.restored_faces = [] self.save_png = True def init_dlib(self, detection_path, landmark5_path, landmark68_path): """Initialize the dlib detectors and predictors.""" self.face_detector = dlib.cnn_face_detection_model_v1(detection_path) self.shape_predictor_5 = dlib.shape_predictor(landmark5_path) self.shape_predictor_68 = dlib.shape_predictor(landmark68_path) def free_dlib_gpu_memory(self): del self.face_detector del self.shape_predictor_5 del self.shape_predictor_68 def read_input_image(self, img_path): # self.input_img is Numpy array, (h, w, c) with RGB order self.input_img = dlib.load_rgb_image(img_path) def detect_faces(self, img_path, upsample_num_times=1, only_keep_largest=False): """ Args: img_path (str): Image path. upsample_num_times (int): Upsamples the image before running the face detector Returns: int: Number of detected faces. """ self.read_input_image(img_path) det_faces = self.face_detector(self.input_img, upsample_num_times) if len(det_faces) == 0: print('No face detected. Try to increase upsample_num_times.') else: if only_keep_largest: print('Detect several faces and only keep the largest.') face_areas = [] for i in range(len(det_faces)): face_area = (det_faces[i].rect.right() - det_faces[i].rect.left()) * ( det_faces[i].rect.bottom() - det_faces[i].rect.top()) face_areas.append(face_area) largest_idx = face_areas.index(max(face_areas)) self.det_faces = [det_faces[largest_idx]] else: self.det_faces = det_faces return len(self.det_faces) def get_face_landmarks_5(self): for face in self.det_faces: shape = self.shape_predictor_5(self.input_img, face.rect) landmark = np.array([[part.x, part.y] for part in shape.parts()]) self.all_landmarks_5.append(landmark) return len(self.all_landmarks_5) def get_face_landmarks_68(self): """Get 68 densemarks for cropped images. Should only have one face at most in the cropped image. """ num_detected_face = 0 for idx, face in enumerate(self.cropped_faces): # face detection det_face = self.face_detector(face, 1) # TODO: can we remove it? if len(det_face) == 0: print(f'Cannot find faces in cropped image with index {idx}.') self.all_landmarks_68.append(None) else: if len(det_face) > 1: print('Detect several faces in the cropped face. Use the ' ' largest one. Note that it will also cause overlap ' 'during paste_faces_to_input_image.') face_areas = [] for i in range(len(det_face)): face_area = (det_face[i].rect.right() - det_face[i].rect.left()) * ( det_face[i].rect.bottom() - det_face[i].rect.top()) face_areas.append(face_area) largest_idx = face_areas.index(max(face_areas)) face_rect = det_face[largest_idx].rect else: face_rect = det_face[0].rect shape = self.shape_predictor_68(face, face_rect) landmark = np.array([[part.x, part.y] for part in shape.parts()]) self.all_landmarks_68.append(landmark) num_detected_face += 1 return num_detected_face def warp_crop_faces(self, save_cropped_path=None, save_inverse_affine_path=None): """Get affine matrix, warp and cropped faces. Also get inverse affine matrix for post-processing. """ for idx, landmark in enumerate(self.all_landmarks_5): # use 5 landmarks to get affine matrix self.similarity_trans.estimate(landmark, self.face_template) affine_matrix = self.similarity_trans.params[0:2, :] self.affine_matrices.append(affine_matrix) # warp and crop faces cropped_face = cv2.warpAffine(self.input_img, affine_matrix, self.face_size) self.cropped_faces.append(cropped_face) # save the cropped face if save_cropped_path is not None: path, ext = os.path.splitext(save_cropped_path) if self.save_png: save_path = f'{path}_{idx:02d}.png' else: save_path = f'{path}_{idx:02d}{ext}' imwrite(cv2.cvtColor(cropped_face, cv2.COLOR_RGB2BGR), save_path) # get inverse affine matrix self.similarity_trans.estimate(self.face_template, landmark * self.upscale_factor) inverse_affine = self.similarity_trans.params[0:2, :] self.inverse_affine_matrices.append(inverse_affine) # save inverse affine matrices if save_inverse_affine_path is not None: path, _ = os.path.splitext(save_inverse_affine_path) save_path = f'{path}_{idx:02d}.pth' torch.save(inverse_affine, save_path) def add_restored_face(self, face): self.restored_faces.append(face) def paste_faces_to_input_image(self, save_path): # operate in the BGR order input_img = cv2.cvtColor(self.input_img, cv2.COLOR_RGB2BGR) h, w, _ = input_img.shape h_up, w_up = h * self.upscale_factor, w * self.upscale_factor # simply resize the background upsample_img = cv2.resize(input_img, (w_up, h_up)) assert len(self.restored_faces) == len( self.inverse_affine_matrices), ('length of restored_faces and affine_matrices are different.') for restored_face, inverse_affine in zip(self.restored_faces, self.inverse_affine_matrices): inv_restored = cv2.warpAffine(restored_face, inverse_affine, (w_up, h_up)) mask = np.ones((*self.face_size, 3), dtype=np.float32) inv_mask = cv2.warpAffine(mask, inverse_affine, (w_up, h_up)) # remove the black borders inv_mask_erosion = cv2.erode(inv_mask, np.ones((2 * self.upscale_factor, 2 * self.upscale_factor), np.uint8)) inv_restored_remove_border = inv_mask_erosion * inv_restored total_face_area = np.sum(inv_mask_erosion) // 3 # compute the fusion edge based on the area of face w_edge = int(total_face_area**0.5) // 20 erosion_radius = w_edge * 2 inv_mask_center = cv2.erode(inv_mask_erosion, np.ones((erosion_radius, erosion_radius), np.uint8)) blur_size = w_edge * 2 inv_soft_mask = cv2.GaussianBlur(inv_mask_center, (blur_size + 1, blur_size + 1), 0) upsample_img = inv_soft_mask * inv_restored_remove_border + (1 - inv_soft_mask) * upsample_img if self.save_png: save_path = save_path.replace('.jpg', '.png').replace('.jpeg', '.png') imwrite(upsample_img.astype(np.uint8), save_path) def clean_all(self): self.all_landmarks_5 = [] self.all_landmarks_68 = [] self.restored_faces = [] self.affine_matrices = [] self.cropped_faces = [] self.inverse_affine_matrices = []