Spaces:
Running
Running
| from codeformer.facelib.utils.face_restoration_helper import FaceRestoreHelper | |
| import numpy as np | |
| from codeformer.basicsr.utils.misc import get_device | |
| class FaceRestoreHelperOptimized(FaceRestoreHelper): | |
| def __init__( | |
| self, | |
| upscale_factor, | |
| face_size=512, | |
| crop_ratio=(1, 1), | |
| det_model="retinaface_resnet50", | |
| save_ext="png", | |
| template_3points=False, | |
| pad_blur=False, | |
| use_parse=False, | |
| device=None, | |
| ): | |
| self.template_3points = template_3points # improve robustness | |
| self.upscale_factor = int(upscale_factor) | |
| # the cropped face ratio based on the square face | |
| self.crop_ratio = crop_ratio # (h, w) | |
| assert self.crop_ratio[0] >= 1 and self.crop_ratio[1] >= 1, "crop ration only supports >=1" | |
| self.face_size = (int(face_size * self.crop_ratio[1]), int(face_size * self.crop_ratio[0])) | |
| self.det_model = det_model | |
| if self.det_model == "dlib": | |
| # 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) | |
| elif self.template_3points: | |
| self.face_template = np.array([[192, 240], [319, 240], [257, 371]]) | |
| else: | |
| # standard 5 landmarks for FFHQ faces with 512 x 512 | |
| # facexlib | |
| self.face_template = np.array( | |
| [ | |
| [192.98138, 239.94708], | |
| [318.90277, 240.1936], | |
| [256.63416, 314.01935], | |
| [201.26117, 371.41043], | |
| [313.08905, 371.15118], | |
| ] | |
| ) | |
| # dlib: left_eye: 36:41 right_eye: 42:47 nose: 30,32,33,34 left mouth corner: 48 right mouth corner: 54 | |
| # self.face_template = np.array([[193.65928, 242.98541], [318.32558, 243.06108], [255.67984, 328.82894], | |
| # [198.22603, 372.82502], [313.91018, 372.75659]]) | |
| self.face_template = self.face_template * (face_size / 512.0) | |
| if self.crop_ratio[0] > 1: | |
| self.face_template[:, 1] += face_size * (self.crop_ratio[0] - 1) / 2 | |
| if self.crop_ratio[1] > 1: | |
| self.face_template[:, 0] += face_size * (self.crop_ratio[1] - 1) / 2 | |
| self.save_ext = save_ext | |
| self.pad_blur = pad_blur | |
| if self.pad_blur is True: | |
| self.template_3points = False | |
| self.all_landmarks_5 = [] | |
| self.det_faces = [] | |
| self.affine_matrices = [] | |
| self.inverse_affine_matrices = [] | |
| self.cropped_faces = [] | |
| self.restored_faces = [] | |
| self.pad_input_imgs = [] | |
| if device is None: | |
| # self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
| self.device = get_device() | |
| else: | |
| self.device = device | |
| # init face detection model | |
| # if self.det_model == "dlib": | |
| # self.face_detector, self.shape_predictor_5 = self.init_dlib( | |
| # dlib_model_url["face_detector"], dlib_model_url["shape_predictor_5"] | |
| # ) | |
| # else: | |
| # self.face_detector = init_detection_model(det_model, half=False, device=self.device) | |
| # init face parsing model | |
| self.use_parse = use_parse | |
| #self.face_parse = init_parsing_model(model_name="parsenet", device=self.device) | |
| # MUST set face_detector and face_parse!!! |