import cv2 import math import copy import numpy as np import argparse import torch # private package from external.landmark_detection.lib import utility from external.landmark_detection.FaceBoxesV2.faceboxes_detector import * class GetCropMatrix(): """ from_shape -> transform_matrix """ def __init__(self, image_size, target_face_scale, align_corners=False): self.image_size = image_size self.target_face_scale = target_face_scale self.align_corners = align_corners def _compose_rotate_and_scale(self, angle, scale, shift_xy, from_center, to_center): cosv = math.cos(angle) sinv = math.sin(angle) fx, fy = from_center tx, ty = to_center acos = scale * cosv asin = scale * sinv a0 = acos a1 = -asin a2 = tx - acos * fx + asin * fy + shift_xy[0] b0 = asin b1 = acos b2 = ty - asin * fx - acos * fy + shift_xy[1] rot_scale_m = np.array([ [a0, a1, a2], [b0, b1, b2], [0.0, 0.0, 1.0] ], np.float32) return rot_scale_m def process(self, scale, center_w, center_h): if self.align_corners: to_w, to_h = self.image_size - 1, self.image_size - 1 else: to_w, to_h = self.image_size, self.image_size rot_mu = 0 scale_mu = self.image_size / (scale * self.target_face_scale * 200.0) shift_xy_mu = (0, 0) matrix = self._compose_rotate_and_scale( rot_mu, scale_mu, shift_xy_mu, from_center=[center_w, center_h], to_center=[to_w / 2.0, to_h / 2.0]) return matrix class TransformPerspective(): """ image, matrix3x3 -> transformed_image """ def __init__(self, image_size): self.image_size = image_size def process(self, image, matrix): return cv2.warpPerspective( image, matrix, dsize=(self.image_size, self.image_size), flags=cv2.INTER_LINEAR, borderValue=0) class TransformPoints2D(): """ points (nx2), matrix (3x3) -> points (nx2) """ def process(self, srcPoints, matrix): # nx3 desPoints = np.concatenate([srcPoints, np.ones_like(srcPoints[:, [0]])], axis=1) desPoints = desPoints @ np.transpose(matrix) # nx3 desPoints = desPoints[:, :2] / desPoints[:, [2, 2]] return desPoints.astype(srcPoints.dtype) class Alignment: def __init__(self, args, model_path, dl_framework, device_ids): self.input_size = 256 self.target_face_scale = 1.0 self.dl_framework = dl_framework # model if self.dl_framework == "pytorch": # conf self.config = utility.get_config(args) self.config.device_id = device_ids[0] # set environment # utility.set_environment(self.config) # self.config.init_instance() # if self.config.logger is not None: # self.config.logger.info("Loaded configure file %s: %s" % (args.config_name, self.config.id)) # self.config.logger.info("\n" + "\n".join(["%s: %s" % item for item in self.config.__dict__.items()])) net = utility.get_net(self.config) if device_ids == [-1]: checkpoint = torch.load(model_path, map_location="cpu") else: checkpoint = torch.load(model_path) net.load_state_dict(checkpoint["net"]) if self.config.device_id == -1: net = net.cpu() else: net = net.to(self.config.device_id) net.eval() self.alignment = net else: assert False self.getCropMatrix = GetCropMatrix(image_size=self.input_size, target_face_scale=self.target_face_scale, align_corners=True) self.transformPerspective = TransformPerspective(image_size=self.input_size) self.transformPoints2D = TransformPoints2D() def norm_points(self, points, align_corners=False): if align_corners: # [0, SIZE-1] -> [-1, +1] return points / torch.tensor([self.input_size - 1, self.input_size - 1]).to(points).view(1, 1, 2) * 2 - 1 else: # [-0.5, SIZE-0.5] -> [-1, +1] return (points * 2 + 1) / torch.tensor([self.input_size, self.input_size]).to(points).view(1, 1, 2) - 1 def denorm_points(self, points, align_corners=False): if align_corners: # [-1, +1] -> [0, SIZE-1] return (points + 1) / 2 * torch.tensor([self.input_size - 1, self.input_size - 1]).to(points).view(1, 1, 2) else: # [-1, +1] -> [-0.5, SIZE-0.5] return ((points + 1) * torch.tensor([self.input_size, self.input_size]).to(points).view(1, 1, 2) - 1) / 2 def preprocess(self, image, scale, center_w, center_h): matrix = self.getCropMatrix.process(scale, center_w, center_h) input_tensor = self.transformPerspective.process(image, matrix) input_tensor = input_tensor[np.newaxis, :] input_tensor = torch.from_numpy(input_tensor) input_tensor = input_tensor.float().permute(0, 3, 1, 2) input_tensor = input_tensor / 255.0 * 2.0 - 1.0 if self.config.device_id == -1: input_tensor = input_tensor.cpu() else: input_tensor = input_tensor.to(self.config.device_id) return input_tensor, matrix def postprocess(self, srcPoints, coeff): # dstPoints = self.transformPoints2D.process(srcPoints, coeff) # matrix^(-1) * src = dst # src = matrix * dst dstPoints = np.zeros(srcPoints.shape, dtype=np.float32) for i in range(srcPoints.shape[0]): dstPoints[i][0] = coeff[0][0] * srcPoints[i][0] + coeff[0][1] * srcPoints[i][1] + coeff[0][2] dstPoints[i][1] = coeff[1][0] * srcPoints[i][0] + coeff[1][1] * srcPoints[i][1] + coeff[1][2] return dstPoints def analyze(self, image, scale, center_w, center_h): input_tensor, matrix = self.preprocess(image, scale, center_w, center_h) if self.dl_framework == "pytorch": with torch.no_grad(): output = self.alignment(input_tensor) landmarks = output[-1][0] else: assert False landmarks = self.denorm_points(landmarks) landmarks = landmarks.data.cpu().numpy()[0] landmarks = self.postprocess(landmarks, np.linalg.inv(matrix)) return landmarks # parser = argparse.ArgumentParser(description="Evaluation script") # args = parser.parse_args() # image_path = './rgb.png' # image = cv2.imread(image_path) # # use_gpu = False # ########### face detection ############ # if use_gpu: # device = torch.device("cuda:0") # else: # device = torch.device("cpu") # # detector = FaceBoxesDetector('FaceBoxes', 'FaceBoxesV2/weights/FaceBoxesV2.pth', use_gpu, device) # # ########### facial alignment ############ # model_path = './weights/68_keypoints_model.pkl' # # if use_gpu: # device_ids = [0] # else: # device_ids = [-1] # # args.config_name = 'alignment' # alignment = Alignment(args, model_path, dl_framework="pytorch", device_ids=device_ids) # image_draw = copy.deepcopy(image) # # ########### inference ############ # ldk_list = [] # # detections, _ = detector.detect(image, 0.9, 1) # for idx in range(len(detections)): # x1_ori = detections[idx][2] # y1_ori = detections[idx][3] # x2_ori = x1_ori + detections[idx][4] # y2_ori = y1_ori + detections[idx][5] # # scale = max(x2_ori - x1_ori, y2_ori - y1_ori) / 180 # center_w = (x1_ori + x2_ori) / 2 # center_h = (y1_ori + y2_ori) / 2 # scale, center_w, center_h = float(scale), float(center_w), float(center_h) # # landmarks_pv = alignment.analyze(image, scale, center_w, center_h) # # for num in range(landmarks_pv.shape[0]): # cv2.circle(image_draw, (round(landmarks_pv[num][0]), round(landmarks_pv[num][1])), # 2, (0, 255, 0), -1) # # ldk_list.append([round(landmarks_pv[num][0]), round(landmarks_pv[num][1])]) # # cv2.imshow("win", image_draw) # # # ldk_img = cv2.imread('/home/gyalex/Desktop/image_landmark_149/all.jpg') # # cv2.imshow("win1", ldk_img) # # cv2.waitKey(0) # # with open('./cord.txt', 'w') as f: # for num in range(len(ldk_list)): # s = str(ldk_list[num][0]) + ' ' + str(ldk_list[num][1]) + '\n' # f.write(s) # # f.close()