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 from mathlib.umeyama import umeyama 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.replace('conv', 'relu'), nn.LeakyReLU(inplace=False, negative_slope=5e-11)), ) def conv(name, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1): return op_name(name, nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, True)) def bn_conv_relu(name, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, momentum = 0.9, track_running_stats=True): return nn.Sequential( op_name(name + '_bn1', nn.BatchNorm2d(in_channels, momentum=momentum, eps=2e-5, track_running_stats=track_running_stats)), op_name(name + '_conv1', nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, True)), op_name(name + '_relu', nn.LeakyReLU(inplace=False, negative_slope=5e-11)), ) def bn(name, in_channels, momentum = 0.9, track_running_stats=True): return op_name(name, nn.BatchNorm2d(in_channels, momentum=momentum, eps=2e-5, track_running_stats=track_running_stats)) class BasicBlock(nn.Module): def __init__(self, stage, unit, inplanes, outplanes, stride): super(BasicBlock, self).__init__() self.bn = bn('stage{}_unit{}_bn1'.format(stage, unit), inplanes) self.conv1 = conv_relu('stage{}_unit{}_conv1'.format(stage, unit), inplanes, outplanes, kernel_size=3, stride=1, padding=1) self.conv2 = conv('stage{}_unit{}_conv2'.format(stage, unit), outplanes, outplanes, kernel_size=3, stride=stride, padding=1) if inplanes != outplanes or stride != 1: self.conv_sc = conv('stage{}_unit{}_conv1sc'.format(stage, unit), inplanes, outplanes, kernel_size=1, stride=stride, padding=0) self.inplanes = inplanes self.outplanes = outplanes self.stride = stride def forward(self, x): residule = x out = self.bn(x) out = self.conv1(out) out = self.conv2(out) if self.inplanes != self.outplanes or self.stride != 1: residule = self.conv_sc(residule) ret = out + residule return ret class ResnetBlock(nn.Module): def __init__(self, stage, inplanes, outplanes, stride=2, n_blocks=1): super(ResnetBlock, self).__init__() self.conv = [] for m in range(n_blocks): if m == 0: self.conv.append(BasicBlock(stage, m + 1, inplanes, outplanes, stride)) else: self.conv.append(BasicBlock(stage, m + 1, outplanes, outplanes, 1)) self.conv = nn.Sequential(*self.conv) def forward(self, x): ret = self.conv(x) return ret class FaceRecognitionServer(nn.Module): def __init__(self): super(FaceRecognitionServer, self).__init__() ch_num = [64, 64, 128, 256, 512] # ch_num = [64, 64] strides = [2, 2, 2, 2] n_blocks = [3, 4, 14, 3] op_list = [conv_relu('conv0', 3, 64, 3, 1, 1)] op_list += [ResnetBlock(i + 1, ch_num[i], ch_num[i + 1], strides[i], n_blocks[i]) for i in range(len(ch_num) - 1)] op_list += [bn('bn1', 512)] op_list += [flatten('flatten', 1)] op_list += [op_name('pre_fc1', nn.Linear(25088, 512))] self.features = nn.Sequential(*op_list) model_path, _ = os.path.split(os.path.realpath(__file__)) weights = torch.load(os.path.join(model_path, 'MMCVFaceRecognitionServer.pth'), map_location=lambda storage, loc: storage) self.load_state_dict(weights) self.eval() def forward(self, x): features = self.features(x) return features class MomocvFaceRecognitionServer(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_recognition_net = FaceRecognitionServer() # self.model_path, _ = os.path.split(os.path.realpath(__file__)) # weights = torch.load(os.path.join(self.model_path, 'MMCVFaceRecognitionServer.pth'), # map_location=lambda storage, loc: storage) # self.face_recognition_net.load_state_dict(weights) self.face_recognition_net.to(self.device) # self.face_recognition_net.eval() def forward(self, images, landmarks): dst_size = 112 with torch.no_grad(): input_numpy = np.zeros((len(images), 3, dst_size, dst_size), dtype=np.float32) assert len(images) == len(landmarks) for ix, img in enumerate(images): landmark = landmarks[ix] mat = landmark_processor.get_transform_mat_for_face_recognition(landmark, dst_size)[0:2] tmp = cv2.warpAffine(img, mat, (dst_size, dst_size)) input_numpy[ix, :, :, :] = (tmp.transpose((2, 0, 1)).astype(np.float32) - 127.5) / 128 # cv2.imshow('tmp', tmp) # cv2.waitKey() in_tensor = torch.from_numpy(input_numpy) in_tensor = in_tensor.to(self.device) features = self.face_recognition_net(in_tensor) features = features.cpu().numpy() norm_factor = np.sqrt(np.sum(features ** 2, axis=1)) norm_factor = norm_factor.reshape(-1, 1) features /= norm_factor return features if __name__ == '__main__': all_jpegs = glob.glob(r'D:\data\deepface_example\data\02b59e75ce91bda300ff827a85a74687d633cdfb5d830e7d3855e31b75201fa8\*.jpg') for s_filename_path in all_jpegs: img = cv2.imread(s_filename_path) # cv2.imshow('img', img) # cv2.waitKey() dflpng = DATAIMG(str(s_filename_path), print_on_no_embedded_data=True) if dflpng is None: print('ERROR') landmarks = dflpng.get_landmarks_mmcv_137() mmcv = MomocvFaceRecognitionServer() features = mmcv.forward([img, img], [landmarks, landmarks]) print(features) print('conansherry') # 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()