import cv2 import numpy as np ######### face enhancement from videoretalking.third_part.GPEN.face_parse.face_parsing import FaceParse from videoretalking.third_part.GPEN.face_detect.retinaface_detection import RetinaFaceDetection from videoretalking.third_part.GPEN.face_parse.face_parsing import FaceParse from videoretalking.third_part.GPEN.face_model.face_gan import FaceGAN # from sr_model.real_esrnet import RealESRNet from videoretalking.third_part.GPEN.align_faces import warp_and_crop_face, get_reference_facial_points from videoretalking.utils.inference_utils import Laplacian_Pyramid_Blending_with_mask class FaceEnhancement(object): def __init__(self, base_dir='./', size=512, model=None, use_sr=True, sr_model=None, channel_multiplier=2, narrow=1, device='cuda'): self.facedetector = RetinaFaceDetection(base_dir, device) self.facegan = FaceGAN(base_dir, size, model, channel_multiplier, narrow, device=device) # self.srmodel = RealESRNet(base_dir, sr_model, device=device) self.srmodel=None self.faceparser = FaceParse(base_dir, device=device) self.use_sr = use_sr self.size = size self.threshold = 0.9 # the mask for pasting restored faces back self.mask = np.zeros((512, 512), np.float32) cv2.rectangle(self.mask, (26, 26), (486, 486), (1, 1, 1), -1, cv2.LINE_AA) self.mask = cv2.GaussianBlur(self.mask, (101, 101), 11) self.mask = cv2.GaussianBlur(self.mask, (101, 101), 11) self.kernel = np.array(( [0.0625, 0.125, 0.0625], [0.125, 0.25, 0.125], [0.0625, 0.125, 0.0625]), dtype="float32") # get the reference 5 landmarks position in the crop settings default_square = True inner_padding_factor = 0.25 outer_padding = (0, 0) self.reference_5pts = get_reference_facial_points( (self.size, self.size), inner_padding_factor, outer_padding, default_square) def mask_postprocess(self, mask, thres=20): mask[:thres, :] = 0; mask[-thres:, :] = 0 mask[:, :thres] = 0; mask[:, -thres:] = 0 mask = cv2.GaussianBlur(mask, (101, 101), 11) mask = cv2.GaussianBlur(mask, (101, 101), 11) return mask.astype(np.float32) def process(self, img, ori_img, bbox=None, face_enhance=True, possion_blending=False): if self.use_sr: img_sr = self.srmodel.process(img) if img_sr is not None: img = cv2.resize(img, img_sr.shape[:2][::-1]) facebs, landms = self.facedetector.detect(img.copy()) orig_faces, enhanced_faces = [], [] height, width = img.shape[:2] full_mask = np.zeros((height, width), dtype=np.float32) full_img = np.zeros(ori_img.shape, dtype=np.uint8) for i, (faceb, facial5points) in enumerate(zip(facebs, landms)): if faceb[4]0)] = tmp_mask[np.where(mask>0)] full_img[np.where(mask>0)] = tmp_img[np.where(mask>0)] mask_sharp = cv2.GaussianBlur(mask_sharp, (0,0), sigmaX=1, sigmaY=1, borderType = cv2.BORDER_DEFAULT) full_mask = full_mask[:, :, np.newaxis] mask_sharp = mask_sharp[:, :, np.newaxis] if self.use_sr and img_sr is not None: img = cv2.convertScaleAbs(img_sr*(1-full_mask) + full_img*full_mask) elif possion_blending is True: if bbox is not None: y1, y2, x1, x2 = bbox mask_bbox = np.zeros_like(mask_sharp) mask_bbox[y1:y2 - 5, x1:x2] = 1 full_img, ori_img, full_mask = [cv2.resize(x,(512,512)) for x in (full_img, ori_img, np.float32(mask_sharp * mask_bbox))] else: full_img, ori_img, full_mask = [cv2.resize(x,(512,512)) for x in (full_img, ori_img, full_mask)] img = Laplacian_Pyramid_Blending_with_mask(full_img, ori_img, full_mask, 6) img = np.clip(img, 0 ,255) img = np.uint8(cv2.resize(img, (width, height))) else: img = cv2.convertScaleAbs(ori_img*(1-full_mask) + full_img*full_mask) img = cv2.convertScaleAbs(ori_img*(1-mask_sharp) + img*mask_sharp) return img, orig_faces, enhanced_faces