File size: 1,683 Bytes
9c501d2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 |
import cv2
import numpy as np
import modules.config
faceRestoreHelper = None
def align_warp_face(self, landmark, border_mode='constant'):
affine_matrix = cv2.estimateAffinePartial2D(landmark, self.face_template, method=cv2.LMEDS)[0]
self.affine_matrices.append(affine_matrix)
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
input_img = self.input_img
cropped_face = cv2.warpAffine(input_img, affine_matrix, self.face_size,
borderMode=border_mode, borderValue=(135, 133, 132))
return cropped_face
def crop_image(img_rgb):
global faceRestoreHelper
if faceRestoreHelper is None:
from extras.facexlib.utils.face_restoration_helper import FaceRestoreHelper
faceRestoreHelper = FaceRestoreHelper(
upscale_factor=1,
model_rootpath=modules.config.path_controlnet,
device='cpu' # use cpu is safer since we are out of memory management
)
faceRestoreHelper.clean_all()
faceRestoreHelper.read_image(np.ascontiguousarray(img_rgb[:, :, ::-1].copy()))
faceRestoreHelper.get_face_landmarks_5()
landmarks = faceRestoreHelper.all_landmarks_5
# landmarks are already sorted with confidence.
if len(landmarks) == 0:
print('No face detected')
return img_rgb
else:
print(f'Detected {len(landmarks)} faces')
result = align_warp_face(faceRestoreHelper, landmarks[0])
return np.ascontiguousarray(result[:, :, ::-1].copy())
|