import torch import torch.nn as nn import math import os import cv2 import numpy as np # from utils.DATAIMG import DATAIMG from utils import landmark_processor import glob def op_name(op_name, m): m.op_name = op_name return m class Flatten(nn.Module): def __init__(self, axis=1): super(Flatten, self).__init__() self.axis = axis def forward(self, x): assert self.axis == 1 x = x.reshape(x.shape[0], -1) return x def flatten(name, axis=1): return op_name(name, Flatten(axis)) def conv_relu(name, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1): return nn.Sequential( op_name(name, nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, True)), op_name(name + '/relu', nn.ReLU()), ) class BasicBlock(nn.Module): def __init__(self, name, inplanes, planes, stride=2): super(BasicBlock, self).__init__() self.op_name = name self.conv1 = conv_relu(name + '/conv1', inplanes, planes, kernel_size=3, stride=stride, padding=1) self.conv2 = conv_relu(name + '/conv2', planes, planes, kernel_size=3, stride=1, padding=1) self.downsample = conv_relu(name + '/sc_conv', inplanes, planes, kernel_size=1, stride=stride) def forward(self, x): residual = x out = self.conv1(x) out = self.conv2(out) if self.downsample is not None: residual = self.downsample(x) ret = out + residual return ret class BigResNetStable(nn.Module): def __init__(self, name, in_channels, out_channels): super(BigResNetStable, self).__init__() self.op_name = name op_list = [] op_list += [conv_relu(name + '/first_conv', in_channels, 32, kernel_size=5, stride=2, padding=2)] ch_num = [32, 48, 64, 96, 128] op_list += [BasicBlock(name + '/stage%d' % (i + 1), ch_num[i], ch_num[i + 1]) for i in range(len(ch_num) - 1)] op_list += [flatten(name + '/flatten', 1)] op_list1 = [op_name(name + '/FC1/FC', nn.Linear(2048, 512)), op_name(name + '/FC1/relu', nn.ReLU())] fullyconnected1 = [op_name(name + '/FC2', nn.Linear(512, out_channels))] poselayer = [op_name('poselayer', nn.Linear(512, 3))] tracking_probe = [op_name('tracking_probe', nn.Linear(2048, 1))] occlusion_probe = [op_name('occlusion_probe', nn.Linear(2048, 87))] self.features = nn.Sequential(*op_list) self.fc1 = nn.Sequential(*op_list1) self.fc2 = nn.Sequential(*fullyconnected1) self.poselayer = nn.Sequential(*poselayer) self.tracking_probe = nn.Sequential(*tracking_probe) self.occlusion_probe = nn.Sequential(*occlusion_probe) def forward(self, x): features = self.features(x) fc1 = self.fc1(features) fullyconnected1 = self.fc2(fc1) poselayer = self.poselayer(fc1) tracking_probe = self.tracking_probe(features) occlusion_probe = self.occlusion_probe(features) return fullyconnected1.cpu().numpy(), poselayer.cpu().numpy(), tracking_probe.cpu().numpy(), occlusion_probe.cpu().numpy() class MomocvFaceAlignment(object): def __init__(self, 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.face_alignment_net = BigResNetStable('BigResNetStable', 3, 174) self.model_path, _ = os.path.split(os.path.realpath(__file__)) weights = torch.load(os.path.join(self.model_path, 'BigResNetStable.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() self.trackingFaceRects = [] def crop_img(self, image, labels=None,img_size = 128): ''' :param imageimgs: Original graph :param labels: The boxes of the original picture ; labels is a ndarray: [list,list] :return: Coordinates and categories relative to the original ''' # Preprocessing labels = np.array([labels]) boxes_ret = np.zeros(labels.shape) crop_imgs = [] ret_M = [] for i in range(labels.shape[0]): box_orig = np.array(labels[i, :]) box = box_orig.astype(np.int32).copy() center = np.array([(box[0] + box[2]) / 2, (box[1] + box[3]) / 2]).astype(np.int32) max_lenth = int(max(box[3] - box[1], box[2] - box[0]) / 2 * 1.0) up = int(center[1] - max_lenth) down = int(center[1] + max_lenth) left = int(center[0] - max_lenth) right = int(center[0] + max_lenth) if up < 0: up = 0 down = max_lenth * 2 if down > image.shape[0]: down = image.shape[0] if left < 0: left = 0 right = max_lenth * 2 if right > image.shape[1]: right = image.shape[1] crop_img = image[up:down, left:right, :].copy() crop_img = cv2.resize(crop_img, (img_size, img_size)) box_orig[0] -= left box_orig[2] -= left box_orig[1] -= up box_orig[3] -= up box_orig[0] *= img_size / (right - left) box_orig[2] *= img_size / (right - left) box_orig[1] *= img_size / (down - up) box_orig[3] *= img_size / (down - up) crop_imgs.append(crop_img) ret_M.append(np.array([img_size / (right - left), img_size / (down - up), left, up])) boxes_ret[i, :] = box_orig return crop_imgs, boxes_ret, ret_M def detect_from_bbox(self, img, bboxs): dst_size = 128 landmarks_res = [] with torch.no_grad(): input_numpy = np.zeros((len(bboxs), 3, dst_size, dst_size), dtype=np.float32) for ix, bbox in enumerate(bboxs): crop_imgs, boxes_ret, ret_M = self.crop_img(img, bbox, dst_size) input_numpy[ix, :, :, :] = crop_imgs[0].transpose((2, 0, 1)).astype(np.float32) # cv2.imshow('tmp', tmp) # cv2.waitKey() in_tensor = torch.from_numpy(input_numpy) in_tensor = in_tensor.to(self.device) fullyconnected1, poselayer, tracking_probe, occlusion_probe = self.face_alignment_net(in_tensor) for ix, pts in enumerate(fullyconnected1): orig_pts = ((np.reshape(pts, (2, 87)).transpose((1, 0))) * dst_size) # print('pts',orig_pts) # landmark2 = landmark2.cpu().detach().numpy()[0].reshape(2, -1).transpose((1, 0)).reshape(-1) # orig_pts = ((orig_pts + 0.5) * 128) orig_pts[:, 0] = orig_pts[:, 0] / ret_M[0][0] + ret_M[0][2] orig_pts[:, 1] = orig_pts[:, 1] / ret_M[0][1] + ret_M[0][3] # orig_pts[0] = orig_pts[0] / ret_M[0][0] + ret_M[0][2] # orig_pts[2] = orig_pts[2] / ret_M[0][0] + ret_M[0][2] # orig_pts[1] = orig_pts[1] / ret_M[0][1] + ret_M[0][3] # orig_pts[3] = orig_pts[3] / ret_M[0][1] + ret_M[0][3] # cur_landmark2.append(orig_pts) # orig_pts = landmark_processor.transform_points(orig_pts, all_mat[ix], invert=True) landmarks_res.append(orig_pts) tracking_probe = 1 / (1 + np.exp(-tracking_probe)) occlusion_probe = 1 / (1 + np.exp(-occlusion_probe)) return landmarks_res, poselayer, tracking_probe, occlusion_probe def detect(self, img, landmarks): dst_size = 128 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(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) # # cv2.imshow('tmp', tmp) # cv2.waitKey() in_tensor = torch.from_numpy(input_numpy) in_tensor = in_tensor.to(self.device) fullyconnected1, poselayer, tracking_probe, occlusion_probe = self.face_alignment_net(in_tensor) for ix, pts in enumerate(fullyconnected1): orig_pts = (np.reshape(pts, (2, 87)).transpose((1, 0)) * dst_size) orig_pts = landmark_processor.transform_points(orig_pts, all_mat[ix], invert=True) landmarks_res.append(orig_pts) tracking_probe = 1 / (1 + np.exp(-tracking_probe)) occlusion_probe = 1 / (1 + np.exp(-occlusion_probe)) return landmarks_res, poselayer, tracking_probe, occlusion_probe def forward(self, images, landmarks): dst_size = 128 landmarks_res = [] with torch.no_grad(): input_numpy = np.zeros((len(images), 3, dst_size, dst_size), dtype=np.float32) assert len(images) == len(landmarks) all_mat = [] for ix, img in enumerate(images): landmark = landmarks[ix] M = landmark_processor.get_transform_mat_mmcv(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) # # cv2.imshow('tmp', tmp) # cv2.waitKey() in_tensor = torch.from_numpy(input_numpy) in_tensor = in_tensor.to(self.device) fullyconnected1, poselayer, tracking_probe, occlusion_probe = self.face_alignment_net(in_tensor) for ix, pts in enumerate(fullyconnected1): orig_pts = (np.reshape(pts, (2, 87)).transpose((1, 0)) * dst_size) orig_pts = landmark_processor.transform_points(orig_pts, all_mat[ix], invert=True) landmarks_res.append(orig_pts) tracking_probe = 1 / (1 + np.exp(-tracking_probe)) occlusion_probe = 1 / (1 + np.exp(-occlusion_probe)) return landmarks_res, poselayer, tracking_probe, occlusion_probe def stable_forward(self, image, detected_faces): for face_rect in detected_faces: if len(self.trackingFaceRects) == 0: new_tracking_rect = [face_rect, True, [0, 0], 0] self.trackingFaceRects.append(new_tracking_rect) with torch.no_grad(): 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 = 128 * 0.8 / min(d[2] - d[0], d[3] - d[1]) dst_center = np.array([0.5, 0.5]) * 128 offset = dst_center - src_center print('hello') else: dst_left_anchor = np.array([0.395, 0.52]) * 128 dst_right_anchor = np.array([1 - 0.395, 0.52]) * 128 # use last anchors src_center = (tracking_face_rect[2][0] + tracking_face_rect[2][1]) / 2 rotate_radian = math.atan2(tracking_face_rect[2][1][1] - tracking_face_rect[2][0][1], tracking_face_rect[2][1][0] - tracking_face_rect[2][0][0]) rotate_degree = rotate_radian / math.pi * 180 print('degree', rotate_degree) dst_anchor_len = cv2.norm(dst_left_anchor - dst_right_anchor) src_anchor_len = cv2.norm(tracking_face_rect[2][0], tracking_face_rect[2][1]) scale = dst_anchor_len / src_anchor_len dst_center = (dst_left_anchor + dst_right_anchor) / 2 offset = dst_center - src_center M = cv2.getRotationMatrix2D((src_center[0], src_center[1]), rotate_degree, scale) M[:, 2] += offset inp = cv2.warpAffine(image, M, (128, 128)) cv2.imshow('inp_stable', inp) # cv2.waitKey() inp = inp.transpose((2, 0, 1)).astype(np.float32) inp = inp[np.newaxis, :, :, :] in_tensor = torch.from_numpy(inp) in_tensor = in_tensor.to(self.device) fullyconnected1, poselayer, tracking_probe, occlusion_probe = self.face_alignment_net(in_tensor) fullyconnected1 = fullyconnected1[0] poselayer = poselayer[0] tracking_probe = tracking_probe[0] occlusion_probe = occlusion_probe[0] orig_pts = (np.reshape(fullyconnected1, (2, 87)).transpose((1, 0)) * 128) orig_pts = landmark_processor.transform_points(orig_pts, M, invert=True) # orig_pts = orig_pts.transpose((1, 0)) fullyconnected1 = orig_pts tracking_probe = 1 / (1 + np.exp(-tracking_probe)) occlusion_probe = 1 / (1 + np.exp(-occlusion_probe)) # update tracking infos tracking_face_rect[1] = False tracking_face_rect[2] = [fullyconnected1[51], fullyconnected1[57]] tracking_face_rect[3] = rotate_degree landmarks.append(fullyconnected1) return landmarks if __name__ == '__main__': all_jpegs = glob.glob(r'E:\deepfacelab_data\expression_dst\7201806132018061208311920180612083119\*.jpg') for s_filename_path in all_jpegs: img = cv2.imread(s_filename_path) dflpng = DATAIMG(str(s_filename_path), print_on_no_embedded_data=True) if dflpng is None: print('ERROR') landmarks = dflpng.get_landmarks() mmcv = MomocvFaceAlignment() fullyconnected1, poselayer, tracking_probe, occlusion_probe = mmcv.forward([img], [landmarks]) fullyconnected1 = fullyconnected1[0] fullyconnected1 = (np.reshape(fullyconnected1, (2, 87)).transpose((1, 0))).astype(np.int32) occlusion_probe = occlusion_probe[0] for ix, pt in enumerate(fullyconnected1): cv2.circle(img, (int(pt[0]), int(pt[1])), 1, (0, 255, 0) if occlusion_probe[ix] > 0.1 else (0, 0, 255), 2) # cv2.putText(img, str(ix), (int(pt[0]), int(pt[1])), cv2.FONT_HERSHEY_SIMPLEX, 0.3, (0, 255, 0), 1) cv2.imshow('img', img) cv2.waitKey() # anchor_dis = 0.445 # dst_size = 128 # anchors_template = np.array([[anchor_dis, 0.52], [1 - anchor_dis, 0.52]]) * dst_size # # src_center = (landmarks[31] + landmarks[35]) / 2 # src_left_anchor = landmarks[31] # src_right_anchor = landmarks[35] # rotate_radian = math.atan2(src_right_anchor[1] - src_left_anchor[1], src_right_anchor[0] - src_left_anchor[0]) # rotate_degree = rotate_radian / math.pi * 180 # dst_eye_len = np.sqrt(np.sum((anchors_template[0] - anchors_template[1]) ** 2)) # src_eye_len = np.sqrt(np.sum((src_left_anchor - src_right_anchor) ** 2)) # scale = dst_eye_len / src_eye_len # dst_center = (anchors_template[0] + anchors_template[1]) / 2 # offset = dst_center - src_center # # M = cv2.getRotationMatrix2D((src_center[0], src_center[1]), rotate_degree, scale) # M[:, 2] += offset # # tmp = cv2.warpAffine(img, M, (dst_size, dst_size)) # cv2.imshow('tmp', tmp) # cv2.waitKey() # # mmcv = MomocvFaceAlignment() # # tmp_input = tmp.transpose((2, 0, 1))[np.newaxis, :, : :].astype(np.float32) # # # tmp_input = cv2.imread(r'E:\deepfacelab_data\workspace\input.png') # # tmp_ori = tmp_input # # tmp_input = tmp_input.transpose((2, 0, 1))[np.newaxis, :, :, :].astype(np.float32) # fullyconnected1, poselayer, tracking_probe, occlusion_probe = mmcv.forward(torch.from_numpy(tmp_input)) # fullyconnected1 = fullyconnected1[0] # # fullyconnected1 = (np.reshape(fullyconnected1, (2, 87)).transpose((1, 0)) * dst_size).astype(np.int32) # for i in range(87): # cv2.circle(tmp, (int(fullyconnected1[i] * 128), int(fullyconnected1[i + 87] * 128)), 1, (255, 0, 0), 1) # cv2.imshow('tmp_ori', tmp) # cv2.waitKey() # # for ix, pt in enumerate(landmarks): # cv2.circle(img, (int(pt[0]), int(pt[1])), 1, (255, 0, 0), 2) # cv2.putText(img, str(ix), (int(pt[0]), int(pt[1])), cv2.FONT_HERSHEY_SIMPLEX, 0.3, (0, 255, 0), 1) # cv2.imshow('img', img) # cv2.waitKey()