git地址
https://github.com/minivision-ai/photo2cartoonz
在face_detect.py里,将其更改文下方代码,加入了def __crop(image, landmarks):,(用于切割图像)
import cv2
import math
import numpy as np
import face_alignment
class FaceDetect:
def __init__(self, device, detector):
# landmarks will be detected by face_alignment library. Set device = 'cuda' if use GPU.
self.fa = face_alignment.FaceAlignment(face_alignment.LandmarksType._2D, device=device, face_detector=detector)
def align(self, image):
landmarks = self.__get_max_face_landmarks(image)
##只加入68个点
# return image, landmarks
#标出68个点,并校正
if landmarks is None:
return None
else:
return self.__rotate(image, landmarks)
##标出68个点,校正,切割
# image, landmarks = self.__rotate(image, landmarks)
# return self.__crop(image, landmarks)
def __get_max_face_landmarks(self, image):
preds = self.fa.get_landmarks(image)
if preds is None:
return None
elif len(preds) == 1:
return preds[0]
else:
# find max face
areas = []
for pred in preds:
landmarks_top = np.min(pred[:, 1])
landmarks_bottom = np.max(pred[:, 1])
landmarks_left = np.min(pred[:, 0])
landmarks_right = np.max(pred[:, 0])
areas.append((landmarks_bottom - landmarks_top) * (landmarks_right - landmarks_left))
max_face_index = np.argmax(areas)
return preds[max_face_index]
@staticmethod
def __rotate(image, landmarks):
# rotation angle
left_eye_corner = landmarks[36]
right_eye_corner = landmarks[45]
radian = np.arctan((left_eye_corner[1] - right_eye_corner[1]) / (left_eye_corner[0] - right_eye_corner[0]))
# image size after rotating
height, width, _ = image.shape
cos = math.cos(radian)
sin = math.sin(radian)
new_w = int(width * abs(cos) + height * abs(sin))
new_h = int(width * abs(sin) + height * abs(cos))
# translation
Tx = new_w // 2 - width // 2
Ty = new_h // 2 - height // 2
# affine matrix
M = np.array([[cos, sin, (1 - cos) * width / 2. - sin * height / 2. + Tx],
[-sin, cos, sin * width / 2. + (1 - cos) * height / 2. + Ty]])
image_rotate = cv2.warpAffine(image, M, (new_w, new_h), borderValue=(255, 255, 255))
landmarks = np.concatenate([landmarks, np.ones((landmarks.shape[0], 1))], axis=1)
landmarks_rotate = np.dot(M, landmarks.T).T
return image_rotate, landmarks_rotate
@staticmethod
def __crop(image, landmarks):
landmarks_top = np.min(landmarks[:, 1])
landmarks_bottom = np.max(landmarks[:, 1])
landmarks_left = np.min(landmarks[:, 0])
landmarks_right = np.max(landmarks[:, 0])
# expand bbox
top = int(landmarks_top - 0.8 * (landmarks_bottom - landmarks_top))
bottom = int(landmarks_bottom + 0.3 * (landmarks_bottom - landmarks_top))
left = int(landmarks_left - 0.3 * (landmarks_right - landmarks_left))
right = int(landmarks_right + 0.3 * (landmarks_right - landmarks_left))
if bottom - top > right - left:
left -= ((bottom - top) - (right - left)) // 2
right = left + (bottom - top)
else:
top -= ((right - left) - (bottom - top)) // 2
bottom = top + (right - left)
image_crop = np.ones((bottom - top + 1, right - left + 1, 3), np.uint8) * 255
h, w = image.shape[:2]
left_white = max(0, -left)
left = max(0, left)
right = min(right, w-1)
right_white = left_white + (right-left)
top_white = max(0, -top)
top = max(0, top)
bottom = min(bottom, h-1)
bottom_white = top_white + (bottom - top)
image_crop[top_white:bottom_white+1, left_white:right_white+1] = image[top:bottom+1, left:right+1].copy()
return image_crop
import argparse
import os
from tqdm import tqdm
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--data_path', type=str, help='photo folder path')
parser.add_argument('--save_path', type=str, help='save folder path')
args = parser.parse_args()
os.makedirs(args.save_path, exist_ok=True)
for idx, image_name in enumerate(tqdm(os.listdir(args.data_path))):
img = cv2.cvtColor(cv2.imread(os.path.join(args.data_path,image_name)), cv2.COLOR_BGR2RGB)
fd = FaceDetect(device='cpu',detector='sfd')
face_info = fd.align(img)
if face_info is not None:
image_align, landmarks_align = face_info
for i in range(landmarks_align.shape[0]):
cv2.circle(image_align, (int(landmarks_align[i][0]), int(landmarks_align[i][1])), 2, (255, 0, 0), -1)
cv2.imwrite(os.path.join(args.save_path,str(idx).zfill(4)+'.jpg'), cv2.cvtColor(image_align, cv2.COLOR_RGB2BGR))
使用68点标记图像时,修改def align(self, image):中的1,2,3点。
当第一点时,在命令行输入python utils/face_detect_dian.py --data_path ./1 --save_path ./2
得到:
当第二点时,在命令行输入python utils/face_detect_dian.py --data_path ./1 --save_path ./3
当第三点时,在命令行输入python utils/face_detect_dian.py --data_path ./1 --save_path ./4
同时将一下代码修改
# if face_info is not None:
# image_align, landmarks_align = face_info
# for i in range(landmarks_align.shape[0]):
# cv2.circle(image_align, (int(landmarks_align[i][0]), int(landmarks_align[i][1])), 2, (255, 0, 0), -1)
if face_info is not None:
image_align = face_info
得到:
最后调用:data_process.py,可以显示最终结果。