| |
|
|
| import numpy as np |
| import cv2 |
|
|
|
|
| def transformation_from_points(points1, points0, smooth=True, p_bias=None): |
| points2 = np.array(points0) |
| points2 = points2.astype(np.float64) |
| points1 = points1.astype(np.float64) |
| c1 = np.mean(points1, axis=0) |
| c2 = np.mean(points2, axis=0) |
| points1 -= c1 |
| points2 -= c2 |
| s1 = np.std(points1) |
| s2 = np.std(points2) |
| points1 /= s1 |
| points2 /= s2 |
| U, S, Vt = np.linalg.svd(np.matmul(points1.T, points2)) |
| R = (np.matmul(U, Vt)).T |
| sR = (s2 / s1) * R |
| T = c2.reshape(2, 1) - (s2 / s1) * np.matmul(R, c1.reshape(2, 1)) |
| M = np.concatenate((sR, T), axis=1) |
| if smooth: |
| bias = points2[2] - points1[2] |
| if p_bias is None: |
| p_bias = bias |
| else: |
| bias = p_bias * 0.2 + bias * 0.8 |
| p_bias = bias |
| M[:, 2] = M[:, 2] + bias |
| return M, p_bias |
|
|
|
|
| class AlignRestore(object): |
| def __init__(self, align_points=3): |
| if align_points == 3: |
| self.upscale_factor = 1 |
| self.crop_ratio = (2.8, 2.8) |
| self.face_template = np.array([[19 - 2, 30 - 10], [56 + 2, 30 - 10], [37.5, 45 - 5]]) |
| self.face_template = self.face_template * 2.8 |
| |
| self.face_size = (int(75 * self.crop_ratio[0]), int(100 * self.crop_ratio[1])) |
| self.p_bias = None |
|
|
| def process(self, img, lmk_align=None, smooth=True, align_points=3): |
| aligned_face, affine_matrix = self.align_warp_face(img, lmk_align, smooth) |
| restored_img = self.restore_img(img, aligned_face, affine_matrix) |
| cv2.imwrite("restored.jpg", restored_img) |
| cv2.imwrite("aligned.jpg", aligned_face) |
| return aligned_face, restored_img |
|
|
| def align_warp_face(self, img, lmks3, smooth=True, border_mode="constant"): |
| affine_matrix, self.p_bias = transformation_from_points(lmks3, self.face_template, smooth, self.p_bias) |
| if border_mode == "constant": |
| border_mode = cv2.BORDER_CONSTANT |
| elif border_mode == "reflect101": |
| border_mode = cv2.BORDER_REFLECT101 |
| elif border_mode == "reflect": |
| border_mode = cv2.BORDER_REFLECT |
| cropped_face = cv2.warpAffine( |
| img, affine_matrix, self.face_size, borderMode=border_mode, borderValue=[127, 127, 127] |
| ) |
| return cropped_face, affine_matrix |
|
|
| def align_warp_face2(self, img, landmark, border_mode="constant"): |
| affine_matrix = cv2.estimateAffinePartial2D(landmark, self.face_template)[0] |
| if border_mode == "constant": |
| border_mode = cv2.BORDER_CONSTANT |
| elif border_mode == "reflect101": |
| border_mode = cv2.BORDER_REFLECT101 |
| elif border_mode == "reflect": |
| border_mode = cv2.BORDER_REFLECT |
| cropped_face = cv2.warpAffine( |
| img, affine_matrix, self.face_size, borderMode=border_mode, borderValue=(135, 133, 132) |
| ) |
| return cropped_face, affine_matrix |
|
|
| def restore_img(self, input_img, face, affine_matrix): |
| h, w, _ = input_img.shape |
| h_up, w_up = int(h * self.upscale_factor), int(w * self.upscale_factor) |
| upsample_img = cv2.resize(input_img, (w_up, h_up), interpolation=cv2.INTER_LANCZOS4) |
| inverse_affine = cv2.invertAffineTransform(affine_matrix) |
| inverse_affine *= self.upscale_factor |
| if self.upscale_factor > 1: |
| extra_offset = 0.5 * self.upscale_factor |
| else: |
| extra_offset = 0 |
| inverse_affine[:, 2] += extra_offset |
| inv_restored = cv2.warpAffine(face, inverse_affine, (w_up, h_up)) |
| mask = np.ones((self.face_size[1], self.face_size[0]), dtype=np.float32) |
| inv_mask = cv2.warpAffine(mask, inverse_affine, (w_up, h_up)) |
| inv_mask_erosion = cv2.erode( |
| inv_mask, np.ones((int(2 * self.upscale_factor), int(2 * self.upscale_factor)), np.uint8) |
| ) |
| pasted_face = inv_mask_erosion[:, :, None] * inv_restored |
| total_face_area = np.sum(inv_mask_erosion) |
| 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) |
| inv_soft_mask = inv_soft_mask[:, :, None] |
| upsample_img = inv_soft_mask * pasted_face + (1 - inv_soft_mask) * upsample_img |
| if np.max(upsample_img) > 256: |
| upsample_img = upsample_img.astype(np.uint16) |
| else: |
| upsample_img = upsample_img.astype(np.uint8) |
| return upsample_img |
|
|
|
|
| class laplacianSmooth: |
| def __init__(self, smoothAlpha=0.3): |
| self.smoothAlpha = smoothAlpha |
| self.pts_last = None |
|
|
| def smooth(self, pts_cur): |
| if self.pts_last is None: |
| self.pts_last = pts_cur.copy() |
| return pts_cur.copy() |
| x1 = min(pts_cur[:, 0]) |
| x2 = max(pts_cur[:, 0]) |
| y1 = min(pts_cur[:, 1]) |
| y2 = max(pts_cur[:, 1]) |
| width = x2 - x1 |
| pts_update = [] |
| for i in range(len(pts_cur)): |
| x_new, y_new = pts_cur[i] |
| x_old, y_old = self.pts_last[i] |
| tmp = (x_new - x_old) ** 2 + (y_new - y_old) ** 2 |
| w = np.exp(-tmp / (width * self.smoothAlpha)) |
| x = x_old * w + x_new * (1 - w) |
| y = y_old * w + y_new * (1 - w) |
| pts_update.append([x, y]) |
| pts_update = np.array(pts_update) |
| self.pts_last = pts_update.copy() |
|
|
| return pts_update |
|
|
