paulo061/codeformer

1import cv22import numpy as np3import torch4 5 6def compute_increased_bbox(bbox, increase_area, preserve_aspect=True):7    left, top, right, bot = bbox8    width = right - left9    height = bot - top10 11    if preserve_aspect:12        width_increase = max(increase_area, ((1 + 2 * increase_area) * height - width) / (2 * width))13        height_increase = max(increase_area, ((1 + 2 * increase_area) * width - height) / (2 * height))14    else:15        width_increase = height_increase = increase_area16    left = int(left - width_increase * width)17    top = int(top - height_increase * height)18    right = int(right + width_increase * width)19    bot = int(bot + height_increase * height)20    return (left, top, right, bot)21 22 23def get_valid_bboxes(bboxes, h, w):24    left = max(bboxes[0], 0)25    top = max(bboxes[1], 0)26    right = min(bboxes[2], w)27    bottom = min(bboxes[3], h)28    return (left, top, right, bottom)29 30 31def align_crop_face_landmarks(img,32                              landmarks,33                              output_size,34                              transform_size=None,35                              enable_padding=True,36                              return_inverse_affine=False,37                              shrink_ratio=(1, 1)):38    """Align and crop face with landmarks.39 40    The output_size and transform_size are based on width. The height is41    adjusted based on shrink_ratio_h/shring_ration_w.42 43    Modified from:44    https://github.com/NVlabs/ffhq-dataset/blob/master/download_ffhq.py45 46    Args:47        img (Numpy array): Input image.48        landmarks (Numpy array): 5 or 68 or 98 landmarks.49        output_size (int): Output face size.50        transform_size (ing): Transform size. Usually the four time of51            output_size.52        enable_padding (float): Default: True.53        shrink_ratio (float | tuple[float] | list[float]): Shring the whole54            face for height and width (crop larger area). Default: (1, 1).55 56    Returns:57        (Numpy array): Cropped face.58    """59    lm_type = 'retinaface_5'  # Options: dlib_5, retinaface_560 61    if isinstance(shrink_ratio, (float, int)):62        shrink_ratio = (shrink_ratio, shrink_ratio)63    if transform_size is None:64        transform_size = output_size * 465 66    # Parse landmarks67    lm = np.array(landmarks)68    if lm.shape[0] == 5 and lm_type == 'retinaface_5':69        eye_left = lm[0]70        eye_right = lm[1]71        mouth_avg = (lm[3] + lm[4]) * 0.572    elif lm.shape[0] == 5 and lm_type == 'dlib_5':73        lm_eye_left = lm[2:4]74        lm_eye_right = lm[0:2]75        eye_left = np.mean(lm_eye_left, axis=0)76        eye_right = np.mean(lm_eye_right, axis=0)77        mouth_avg = lm[4]78    elif lm.shape[0] == 68:79        lm_eye_left = lm[36:42]80        lm_eye_right = lm[42:48]81        eye_left = np.mean(lm_eye_left, axis=0)82        eye_right = np.mean(lm_eye_right, axis=0)83        mouth_avg = (lm[48] + lm[54]) * 0.584    elif lm.shape[0] == 98:85        lm_eye_left = lm[60:68]86        lm_eye_right = lm[68:76]87        eye_left = np.mean(lm_eye_left, axis=0)88        eye_right = np.mean(lm_eye_right, axis=0)89        mouth_avg = (lm[76] + lm[82]) * 0.590 91    eye_avg = (eye_left + eye_right) * 0.592    eye_to_eye = eye_right - eye_left93    eye_to_mouth = mouth_avg - eye_avg94 95    # Get the oriented crop rectangle96    # x: half width of the oriented crop rectangle97    x = eye_to_eye - np.flipud(eye_to_mouth) * [-1, 1]98    #  - np.flipud(eye_to_mouth) * [-1, 1]: rotate 90 clockwise99    # norm with the hypotenuse: get the direction100    x /= np.hypot(*x)  # get the hypotenuse of a right triangle101    rect_scale = 1  # TODO: you can edit it to get larger rect102    x *= max(np.hypot(*eye_to_eye) * 2.0 * rect_scale, np.hypot(*eye_to_mouth) * 1.8 * rect_scale)103    # y: half height of the oriented crop rectangle104    y = np.flipud(x) * [-1, 1]105 106    x *= shrink_ratio[1]  # width107    y *= shrink_ratio[0]  # height108 109    # c: center110    c = eye_avg + eye_to_mouth * 0.1111    # quad: (left_top, left_bottom, right_bottom, right_top)112    quad = np.stack([c - x - y, c - x + y, c + x + y, c + x - y])113    # qsize: side length of the square114    qsize = np.hypot(*x) * 2115 116    quad_ori = np.copy(quad)117    # Shrink, for large face118    # TODO: do we really need shrink119    shrink = int(np.floor(qsize / output_size * 0.5))120    if shrink > 1:121        h, w = img.shape[0:2]122        rsize = (int(np.rint(float(w) / shrink)), int(np.rint(float(h) / shrink)))123        img = cv2.resize(img, rsize, interpolation=cv2.INTER_AREA)124        quad /= shrink125        qsize /= shrink126 127    # Crop128    h, w = img.shape[0:2]129    border = max(int(np.rint(qsize * 0.1)), 3)130    crop = (int(np.floor(min(quad[:, 0]))), int(np.floor(min(quad[:, 1]))), int(np.ceil(max(quad[:, 0]))),131            int(np.ceil(max(quad[:, 1]))))132    crop = (max(crop[0] - border, 0), max(crop[1] - border, 0), min(crop[2] + border, w), min(crop[3] + border, h))133    if crop[2] - crop[0] < w or crop[3] - crop[1] < h:134        img = img[crop[1]:crop[3], crop[0]:crop[2], :]135        quad -= crop[0:2]136 137    # Pad138    # pad: (width_left, height_top, width_right, height_bottom)139    h, w = img.shape[0:2]140    pad = (int(np.floor(min(quad[:, 0]))), int(np.floor(min(quad[:, 1]))), int(np.ceil(max(quad[:, 0]))),141           int(np.ceil(max(quad[:, 1]))))142    pad = (max(-pad[0] + border, 0), max(-pad[1] + border, 0), max(pad[2] - w + border, 0), max(pad[3] - h + border, 0))143    if enable_padding and max(pad) > border - 4:144        pad = np.maximum(pad, int(np.rint(qsize * 0.3)))145        img = np.pad(img, ((pad[1], pad[3]), (pad[0], pad[2]), (0, 0)), 'reflect')146        h, w = img.shape[0:2]147        y, x, _ = np.ogrid[:h, :w, :1]148        mask = np.maximum(1.0 - np.minimum(np.float32(x) / pad[0],149                                           np.float32(w - 1 - x) / pad[2]),150                          1.0 - np.minimum(np.float32(y) / pad[1],151                                           np.float32(h - 1 - y) / pad[3]))152        blur = int(qsize * 0.02)153        if blur % 2 == 0:154            blur += 1155        blur_img = cv2.boxFilter(img, 0, ksize=(blur, blur))156 157        img = img.astype('float32')158        img += (blur_img - img) * np.clip(mask * 3.0 + 1.0, 0.0, 1.0)159        img += (np.median(img, axis=(0, 1)) - img) * np.clip(mask, 0.0, 1.0)160        img = np.clip(img, 0, 255)  # float32, [0, 255]161        quad += pad[:2]162 163    # Transform use cv2164    h_ratio = shrink_ratio[0] / shrink_ratio[1]165    dst_h, dst_w = int(transform_size * h_ratio), transform_size166    template = np.array([[0, 0], [0, dst_h], [dst_w, dst_h], [dst_w, 0]])167    # use cv2.LMEDS method for the equivalence to skimage transform168    # ref: https://blog.csdn.net/yichxi/article/details/115827338169    affine_matrix = cv2.estimateAffinePartial2D(quad, template, method=cv2.LMEDS)[0]170    cropped_face = cv2.warpAffine(171        img, affine_matrix, (dst_w, dst_h), borderMode=cv2.BORDER_CONSTANT, borderValue=(135, 133, 132))  # gray172 173    if output_size < transform_size:174        cropped_face = cv2.resize(175            cropped_face, (output_size, int(output_size * h_ratio)), interpolation=cv2.INTER_LINEAR)176 177    if return_inverse_affine:178        dst_h, dst_w = int(output_size * h_ratio), output_size179        template = np.array([[0, 0], [0, dst_h], [dst_w, dst_h], [dst_w, 0]])180        # use cv2.LMEDS method for the equivalence to skimage transform181        # ref: https://blog.csdn.net/yichxi/article/details/115827338182        affine_matrix = cv2.estimateAffinePartial2D(183            quad_ori, np.array([[0, 0], [0, output_size], [dst_w, dst_h], [dst_w, 0]]), method=cv2.LMEDS)[0]184        inverse_affine = cv2.invertAffineTransform(affine_matrix)185    else:186        inverse_affine = None187    return cropped_face, inverse_affine188 189 190def paste_face_back(img, face, inverse_affine):191    h, w = img.shape[0:2]192    face_h, face_w = face.shape[0:2]193    inv_restored = cv2.warpAffine(face, inverse_affine, (w, h))194    mask = np.ones((face_h, face_w, 3), dtype=np.float32)195    inv_mask = cv2.warpAffine(mask, inverse_affine, (w, h))196    # remove the black borders197    inv_mask_erosion = cv2.erode(inv_mask, np.ones((2, 2), np.uint8))198    inv_restored_remove_border = inv_mask_erosion * inv_restored199    total_face_area = np.sum(inv_mask_erosion) // 3200    # compute the fusion edge based on the area of face201    w_edge = int(total_face_area**0.5) // 20202    erosion_radius = w_edge * 2203    inv_mask_center = cv2.erode(inv_mask_erosion, np.ones((erosion_radius, erosion_radius), np.uint8))204    blur_size = w_edge * 2205    inv_soft_mask = cv2.GaussianBlur(inv_mask_center, (blur_size + 1, blur_size + 1), 0)206    img = inv_soft_mask * inv_restored_remove_border + (1 - inv_soft_mask) * img207    # float32, [0, 255]208    return img209 210 211if __name__ == '__main__':212    import os213 214    from facelib.detection import init_detection_model215    from facelib.utils.face_restoration_helper import get_largest_face216 217    img_path = '/home/wxt/datasets/ffhq/ffhq_wild/00009.png'218    img_name = os.splitext(os.path.basename(img_path))[0]219 220    # initialize model221    det_net = init_detection_model('retinaface_resnet50', half=False)222    img_ori = cv2.imread(img_path)223    h, w = img_ori.shape[0:2]224    # if larger than 800, scale it225    scale = max(h / 800, w / 800)226    if scale > 1:227        img = cv2.resize(img_ori, (int(w / scale), int(h / scale)), interpolation=cv2.INTER_LINEAR)228 229    with torch.no_grad():230        bboxes = det_net.detect_faces(img, 0.97)231    if scale > 1:232        bboxes *= scale  # the score is incorrect233    bboxes = get_largest_face(bboxes, h, w)[0]234 235    landmarks = np.array([[bboxes[i], bboxes[i + 1]] for i in range(5, 15, 2)])236 237    cropped_face, inverse_affine = align_crop_face_landmarks(238        img_ori,239        landmarks,240        output_size=512,241        transform_size=None,242        enable_padding=True,243        return_inverse_affine=True,244        shrink_ratio=(1, 1))245 246    cv2.imwrite(f'tmp/{img_name}_cropeed_face.png', cropped_face)247    img = paste_face_back(img_ori, cropped_face, inverse_affine)248    cv2.imwrite(f'tmp/{img_name}_back.png', img)249