import torch import torch.onnx import torch.nn as nn import sys import os from momocv.resnet import resnet18 import re import numpy as np import cv2 from sklearn import linear_model from tqdm import tqdm from momocv.LeftEye import get_left_eye_symbol from utils import landmark_processor from mtcnn.detector import MTCNNFaceDetector import math class ResNet18(nn.Module): def __init__(self, out_channels): super(ResNet18, self).__init__() self.op_name = 'ResNet18' self.res18 = resnet18(pretrained=False) self.res18.fc = nn.Linear(512, out_channels) def forward(self, x): ret = self.res18(x) return ret def get_full_face_symbol(output_nc=137 * 2): return ResNet18(output_nc) class Model137(nn.Module): def __init__(self, gpu_id=None): super(Model137, self).__init__() self.device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu') self.face_alignment_net = ResNet18(137 * 2) self.model_path, _ = os.path.split(os.path.realpath(__file__)) weights = torch.load(os.path.join(self.model_path, 'FullFace.pth'), map_location=lambda storage, loc: storage) self.face_alignment_net.load_state_dict(weights) self.to(self.device) self.eval() def forward(self, imgs): pred_key_pts = self.face_alignment_net(imgs) return pred_key_pts class MomocvFaceAlignmentFinalV1(object): def __init__(self, predict_eye=False, gpu_id=None): self.gpu_id = gpu_id self.device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu') self.predict_eye = predict_eye self.face_alignment_net = get_full_face_symbol() self.model_path, _ = os.path.split(os.path.realpath(__file__)) weights = torch.load(os.path.join(self.model_path, 'FullFace.pth'), map_location=lambda storage, loc: storage) self.face_alignment_net.load_state_dict(weights) self.face_alignment_net.to(self.device) self.face_alignment_net.eval() if predict_eye: self.eye_alignment_net = get_left_eye_symbol() self.model_path, _ = os.path.split(os.path.realpath(__file__)) weights = torch.load(os.path.join(self.model_path, 'LeftEye.pth'), map_location=lambda storage, loc: storage) self.eye_alignment_net.load_state_dict(weights) self.eye_alignment_net.to(self.device) self.eye_alignment_net.eval() self.trackingFaceRects = [] print('conansherry MomocvFaceAlignmentFinalV1') def forward(self, img_tensor): fullyconnected1 = self.face_alignment_net(img_tensor).detach().cpu().numpy() return fullyconnected1 def detect(self, img, landmarks): dst_size = 256 landmarks_res = [] with torch.no_grad(): input_numpy = np.zeros((len(landmarks), 3, dst_size, dst_size), dtype=np.float32) all_mat = [] for ix, landmark in enumerate(landmarks): M = landmark_processor.get_transform_mat_mmcv_bigger(landmark, dst_size) all_mat.append(M) tmp = cv2.warpAffine(img, M, (dst_size, dst_size)) input_numpy[ix, :, :, :] = tmp.transpose((2, 0, 1)).astype(np.float32) / 255 in_tensor = torch.from_numpy(input_numpy) in_tensor = in_tensor.to(self.device) fullyconnected1 = self.face_alignment_net(in_tensor).detach().cpu().numpy() for ix, pts in enumerate(fullyconnected1): orig_pts = (np.reshape(pts, (2, 137)).transpose((1, 0)) * dst_size) orig_pts = landmark_processor.transform_points(orig_pts, all_mat[ix], invert=True) if self.predict_eye: eye_landmark = self.detect_eye(img, orig_pts) orig_pts[87:104] = eye_landmark[0] orig_pts[104:121] = eye_landmark[1] landmarks_res.append(orig_pts) return landmarks_res def detect_eye(self, img, landmarks): dst_size = 96 src_len = cv2.norm(landmarks[96] - landmarks[88]) dst_len = 96 * 0.7 degree = math.atan2(landmarks[88, 1] - landmarks[96, 1], landmarks[88, 0] - landmarks[96, 0]) src_center = (landmarks[88] + landmarks[96]) / 2 offset = np.array([0.5, 0.5]) * 96 - src_center left_M = cv2.getRotationMatrix2D((src_center[0], src_center[1]), math.degrees(degree), dst_len / src_len) left_M[:, 2] += offset left_eye_img = cv2.warpAffine(img, left_M, (dst_size, dst_size)) dst_size = 96 src_len = cv2.norm(landmarks[105] - landmarks[113]) dst_len = 96 * 0.7 degree = math.atan2(landmarks[113, 1] - landmarks[105, 1], landmarks[113, 0] - landmarks[105, 0]) src_center = (landmarks[105] + landmarks[113]) / 2 offset = np.array([0.5, 0.5]) * 96 - src_center right_M = cv2.getRotationMatrix2D((src_center[0], src_center[1]), math.degrees(degree), dst_len / src_len) right_M[:, 2] += offset right_eye_img = cv2.warpAffine(img, right_M, (dst_size, dst_size)) right_eye_img = cv2.flip(right_eye_img, 1) # cv2.imshow('left_eye_img', left_eye_img) # cv2.imshow('right_eye_img', right_eye_img) with torch.no_grad(): input_numpy = np.zeros((2, 3, dst_size, dst_size), dtype=np.float32) input_numpy[0, :, :, :] = left_eye_img.transpose((2, 0, 1)).astype(np.float32) / 255 input_numpy[1, :, :, :] = right_eye_img.transpose((2, 0, 1)).astype(np.float32) / 255 in_tensor = torch.from_numpy(input_numpy) in_tensor = in_tensor.to(self.device) fullyconnected1 = self.eye_alignment_net(in_tensor).detach().cpu().numpy() landmarks_res = [] all_mat = [left_M, right_M] for ix, pts in enumerate(fullyconnected1): orig_pts = (np.reshape(pts, (2, 17)).transpose((1, 0)) * dst_size) if ix == 1: orig_pts[:, 0] = dst_size - orig_pts[:, 0] orig_pts = landmark_processor.transform_points(orig_pts, all_mat[ix], invert=True) landmarks_res.append(orig_pts) return landmarks_res # cv2.waitKey() def stable_forward(self, image, detected_faces, reset=False): if reset is True: self.trackingFaceRects = [] if len(self.trackingFaceRects) == 0: for face_rect in detected_faces: new_tracking_rect = [face_rect, True, [0, 0], 0, None] self.trackingFaceRects.append(new_tracking_rect) with torch.no_grad(): landmarks = [] eye_landmarks = [] for tracking_face_rect in self.trackingFaceRects: if tracking_face_rect[1] == True: d = tracking_face_rect[0] src_center = np.array([d[2] - (d[2] - d[0]) / 2.0, d[3] - (d[3] - d[1]) / 2.0]) rotate_degree = tracking_face_rect[3] scale = 256 * 0.6 / min(d[2] - d[0], d[3] - d[1]) dst_center = np.array([0.5, 0.5]) * 256 offset = dst_center - src_center M = cv2.getRotationMatrix2D((src_center[0], src_center[1]), rotate_degree, scale) M[:, 2] += offset else: rotate_degree = 0 M = landmark_processor.get_transform_mat_mmcv_bigger(tracking_face_rect[4], 256) inp = cv2.warpAffine(image, M, (256, 256)) cv2.imshow('inp', inp) orig_inp = inp inp = inp.transpose((2, 0, 1)).astype(np.float32) inp = inp[np.newaxis, :, :, :] / 255 t0 = cv2.getTickCount() in_tensor = torch.from_numpy(inp) in_tensor = in_tensor.to(self.device) fullyconnected1 = self.face_alignment_net(in_tensor).detach().cpu().numpy() fullyconnected1 = fullyconnected1[0] orig_pts = (np.reshape(fullyconnected1, (2, 137)).transpose((1, 0)) * 256) t2 = cv2.getTickCount() orig_pts = landmark_processor.transform_points(orig_pts, M, invert=True) if self.predict_eye: eye_landmark = self.detect_eye(image, orig_pts) orig_pts[87:104] = eye_landmark[0] orig_pts[104:121] = eye_landmark[1] eye_landmarks.append(eye_landmark) # orig_pts = orig_pts.transpose((1, 0)) fullyconnected1 = orig_pts # update tracking infos tracking_face_rect[1] = False tracking_face_rect[2] = [fullyconnected1[68], fullyconnected1[74], fullyconnected1[96], fullyconnected1[113]] tracking_face_rect[3] = rotate_degree tracking_face_rect[4] = fullyconnected1 landmarks.append(fullyconnected1) return landmarks, eye_landmarks if __name__=='__main__': gpu_id = 0 net = MomocvFaceAlignmentFinalV1(predict_eye=True, gpu_id=gpu_id) video_name = r'G:\all_online_videos\2019_04_09_21_57_25_10342e28-ed60-4568-8790-5a4431384031_Trim.mp4' cap = cv2.VideoCapture(video_name) face_detector = MTCNNFaceDetector(gpu_id=gpu_id) reset = False bounding_boxes = [] while True: _, in_frame = cap.read() if in_frame is None: cap = cv2.VideoCapture(video_name) _, in_frame = cap.read() if len(bounding_boxes) == 0 or reset: bounding_boxes, landmarks = face_detector.forward(in_frame, min_face_size=100, thresholds=[0.8, 0.9, 0.95]) # for box_score in bounding_boxes: # cv2.rectangle(in_frame, (int(box_score[0]), int(box_score[1])), # (int(box_score[2]), int(box_score[3])), # (0, 255, 0), # 2) # # for pt in landmarks: # for i in range(5): # cv2.circle(in_frame, (int(pt[i]), int(pt[i + 5])), 1, (255, 0, 0), 2) for _ in range(3): landmark137, eye_landmark = net.stable_forward(in_frame, bounding_boxes, reset) reset = False for pts in landmark137: for ix, pt in enumerate(pts): # cv2.putText(in_frame, str(ix), (int(pt[0]), int(pt[1])), cv2.FONT_HERSHEY_SIMPLEX, 0.3, (255, 0, 0)) cv2.circle(in_frame, (pt[0], pt[1]), 1, (0, 255, 0), 1) # for pts in eye_landmark: # for eye in pts: # for pt in eye: # cv2.circle(in_frame, (pt[0], pt[1]), 1, (0, 0, 255), 1) cv2.imshow('in_frame', in_frame) key = cv2.waitKey(10) if key == ord('r'): reset = True bounding_boxes = []