Spaces:
Running
Running
| 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 = [] | |