import torch.nn as nn import torch.utils.model_zoo as model_zoo import torch import numpy as np import os from utils import landmark_processor, umeyama import cv2 __all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101', 'resnet152'] model_urls = { 'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth', 'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth', 'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth', 'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth', 'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth', } def conv3x3(in_planes, out_planes, stride=1): """3x3 convolution with padding""" return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False) def conv1x1(in_planes, out_planes, stride=1): """1x1 convolution""" return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False) class BasicBlock(nn.Module): expansion = 1 def __init__(self, inplanes, planes, stride=1, downsample=None): super(BasicBlock, self).__init__() self.conv1 = conv3x3(inplanes, planes, stride) self.bn1 = nn.BatchNorm2d(planes) self.relu = nn.ReLU(inplace=True) self.conv2 = conv3x3(planes, planes) self.bn2 = nn.BatchNorm2d(planes) self.downsample = downsample self.stride = stride def forward(self, x): identity = x out = self.conv1(x) out = self.bn1(out) out = self.relu(out) out = self.conv2(out) out = self.bn2(out) if self.downsample is not None: identity = self.downsample(x) out += identity out = self.relu(out) return out class Bottleneck(nn.Module): expansion = 4 def __init__(self, inplanes, planes, stride=1, downsample=None): super(Bottleneck, self).__init__() self.conv1 = conv1x1(inplanes, planes) self.bn1 = nn.BatchNorm2d(planes) self.conv2 = conv3x3(planes, planes, stride) self.bn2 = nn.BatchNorm2d(planes) self.conv3 = conv1x1(planes, planes * self.expansion) self.bn3 = nn.BatchNorm2d(planes * self.expansion) self.relu = nn.ReLU(inplace=True) self.downsample = downsample self.stride = stride def forward(self, x): identity = x out = self.conv1(x) out = self.bn1(out) out = self.relu(out) out = self.conv2(out) out = self.bn2(out) out = self.relu(out) out = self.conv3(out) out = self.bn3(out) if self.downsample is not None: identity = self.downsample(x) out += identity out = self.relu(out) return out class ResNet(nn.Module): def __init__(self, block, layers, num_classes=1000, is_1k=False, zero_init_residual=False): super(ResNet, self).__init__() self.inplanes = 64 self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False) self.bn1 = nn.BatchNorm2d(64) self.relu = nn.ReLU(inplace=True) self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) self.layer1 = self._make_layer(block, 64, layers[0]) self.layer2 = self._make_layer(block, 128, layers[1], stride=2) self.layer3 = self._make_layer(block, 256, layers[2], stride=2) self.layer4 = self._make_layer(block, 512, layers[3], stride=2) self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) if is_1k: self.fc = nn.Sequential(*[nn.Linear(512 * block.expansion, num_classes), nn.Tanh()]) else: self.fc_key = nn.Sequential(*[nn.Linear(256 * block.expansion, 45 * 2), nn.Tanh()]) self.fc_ctrl = nn.Sequential(*[nn.Linear(256 * block.expansion, 48 * 2), nn.Tanh()]) self.is_1k = is_1k for m in self.modules(): if isinstance(m, nn.Conv2d): nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') elif isinstance(m, nn.BatchNorm2d): nn.init.constant_(m.weight, 1) nn.init.constant_(m.bias, 0) # Zero-initialize the last BN in each residual branch, # so that the residual branch starts with zeros, and each residual block behaves like an identity. # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677 if zero_init_residual: for m in self.modules(): if isinstance(m, Bottleneck): nn.init.constant_(m.bn3.weight, 0) elif isinstance(m, BasicBlock): nn.init.constant_(m.bn2.weight, 0) def _make_layer(self, block, planes, blocks, stride=1): downsample = None if stride != 1 or self.inplanes != planes * block.expansion: downsample = nn.Sequential( conv1x1(self.inplanes, planes * block.expansion, stride), nn.BatchNorm2d(planes * block.expansion), ) layers = [] layers.append(block(self.inplanes, planes, stride, downsample)) self.inplanes = planes * block.expansion for _ in range(1, blocks): layers.append(block(self.inplanes, planes)) return nn.Sequential(*layers) def forward(self, x): x = self.conv1(x) x = self.bn1(x) x = self.relu(x) x = self.maxpool(x) x = self.layer1(x) x = self.layer2(x) x = self.layer3(x) x = self.layer4(x) if self.is_1k: key = self.avgpool(x) key = key.view(key.size(0), -1) key = self.fc(key) return key else: key, ctrl = torch.chunk(x, 2, dim=1) key = self.avgpool(key) key = key.view(key.size(0), -1) key = self.fc_key(key) ctrl = self.avgpool(ctrl) ctrl = ctrl.view(ctrl.size(0), -1) ctrl = self.fc_ctrl(ctrl) return key, ctrl def resnet18(pretrained=False, **kwargs): """Constructs a ResNet-18 model. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet """ model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs) if pretrained: model.load_state_dict(model_zoo.load_url(model_urls['resnet18']), strict=False) return model def resnet34(pretrained=False, **kwargs): """Constructs a ResNet-34 model. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet """ model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs) if pretrained: model.load_state_dict(model_zoo.load_url(model_urls['resnet34'])) return model def resnet50(pretrained=False, **kwargs): """Constructs a ResNet-50 model. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet """ model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs) if pretrained: model.load_state_dict(model_zoo.load_url(model_urls['resnet50'])) return model def resnet101(pretrained=False, **kwargs): """Constructs a ResNet-101 model. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet """ model = ResNet(Bottleneck, [3, 4, 23, 3], **kwargs) if pretrained: model.load_state_dict(model_zoo.load_url(model_urls['resnet101'])) return model def resnet152(pretrained=False, **kwargs): """Constructs a ResNet-152 model. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet """ model = ResNet(Bottleneck, [3, 8, 36, 3], **kwargs) if pretrained: model.load_state_dict(model_zoo.load_url(model_urls['resnet152'])) return model class Model1k(nn.Module): def __init__(self, gpu_id=None): super(Model1k, self).__init__() self.device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu') self.face_alignment_net = resnet18(pretrained=False, num_classes=1000 * 2, is_1k=True) self.model_path, _ = os.path.split(os.path.realpath(__file__)) weights = torch.load(os.path.join(self.model_path, 'face_alignment_1k.pth'), map_location=lambda storage, loc: storage) self.load_state_dict(weights) self.to(self.device) self.eval() def forward(self, imgs): pred_key_pts = self.face_alignment_net(imgs) pred_key_pts = pred_key_pts + 0.5 return pred_key_pts class MomocvFaceAlignment1K(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 = Model1k(gpu_id) # 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 MomocvFaceAlignment1K') 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_full_face(landmark, dst_size) all_mat.append(M) tmp = cv2.warpAffine(img, M, (dst_size, dst_size)) # cv2.imshow('inp', tmp) # cv2.waitKey() 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, 1000)).transpose((1, 0)) * dst_size) orig_pts = landmark_processor.transform_points(orig_pts, all_mat[ix], invert=True) 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 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 = [] for ix, tracking_face_rect in enumerate(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_{}'.format(ix), inp) # cv2.waitKey() orig_inp = inp inp = inp.transpose((2, 0, 1)).astype(np.float32) inp = inp[np.newaxis, :, :, :] / 255 in_tensor = torch.from_numpy(inp) in_tensor = in_tensor.cuda(0) fullyconnected1 = self.forward(in_tensor) fullyconnected1 = fullyconnected1[0] orig_pts = (np.reshape(fullyconnected1, (2, 1000)).transpose((1, 0))) * 256 t2 = cv2.getTickCount() orig_pts = landmark_processor.transform_points(orig_pts, M, invert=True) # orig_pts = orig_pts.transpose((1, 0)) fullyconnected1 = orig_pts # update tracking infos tracking_face_rect[1] = False tracking_face_rect[2] = None tracking_face_rect[3] = rotate_degree tracking_face_rect[4] = fullyconnected1 # fullyconnected1 = landmark_processor.pts_1k_to_137(fullyconnected1) # eye_landmark = self.detect_eye(image, fullyconnected1) # fullyconnected1[87:104] = eye_landmark[0] # fullyconnected1[104:121] = eye_landmark[1] landmarks.append(fullyconnected1) return landmarks def detect_according_5pts(self, img, pts5): dst_size = 256 with torch.no_grad(): input_numpy = np.zeros((1, 3, dst_size, dst_size), dtype=np.float32) eye_dis = 0.34 mouth_dis = 0.34 g_Average_5point_180 = np.array([ eye_dis, 0.3, 1 - eye_dis, 0.3, 0.5, 0.6, mouth_dis, 0.63, 1 - mouth_dis, 0.63 ]) # print(g_Average_5point_180) left_eye = np.array([pts5[0], pts5[5]]) right_eye = np.array([pts5[1], pts5[6]]) nose = np.array([pts5[2], pts5[7]]) left_mouth = np.array([pts5[3], pts5[8]]) right_mouth = np.array([pts5[4], pts5[9]]) pts5_src = np.vstack((left_eye, right_eye, nose, left_mouth, right_mouth)) pts5_src = np.array(pts5_src).astype(np.int32) pts5_dst = g_Average_5point_180.reshape((5, -1)) * dst_size mat = umeyama(pts5_src, pts5_dst, True)[0:2] tmp = cv2.warpAffine(img, mat, (dst_size, dst_size)) # cv2.imshow("tmp", tmp) # cv2.waitKey() input_numpy[0, :, :, :] = 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() orig_pts = (np.reshape(fullyconnected1[0], (2, 1000)).transpose((1, 0)) * dst_size) orig_pts = landmark_processor.transform_points(orig_pts, mat, invert=True) return orig_pts