class FaceAnti:
def __init__(self):
# 准备模型
self.net = Net(num_class=2, is_first_bn=True) # 网络结构加载
model_path = 'model_A_color_48/checkpoint/global_min_acer_model.pth' # 权重加载
if torch.cuda.is_available():
state_dict = torch.load(model_path, map_location='cuda')
else:
state_dict = torch.load(model_path, map_location='cpu')
new_state_dict = OrderedDict()
for k, v in state_dict.items():
name = k[7:] # remove 'module'
new_state_dict[name] = v
self.net.load_state_dict(new_state_dict)
if torch.cuda.is_available():
self.net = self.net.cuda()
self.net.eval() # bn 层输出锁定
def classify(self, color):
return self.detect(color)
def detect(self, color):
color = cv2.resize(color, (RESIZE_SIZE, RESIZE_SIZE))
def color_augmentor(image, target_shape=(64, 64, 3), is_infer=False):
if is_infer:
augment_img = iaa.Sequential([iaa.Fliplr(0)])
image = augment_img.augment_image(image)
image = TTA_36_cropps(image, target_shape)
return image
color = color_augmentor(color, target_shape=(64, 64, 3), is_infer=True)
n = len(color)
color = np.concatenate(color, axis=0)
image = color
image = np.transpose(image, (0, 3, 1, 2))
image = image.astype(np.float32)
image = image / 255.0
input_image = torch.FloatTensor(image)
if (len(input_image.size()) == 4) and torch.cuda.is_available():
input_image = input_image.unsqueeze(0).cuda()
elif (len(input_image.size()) == 4) and not torch.cuda.is_available():
input_image = input_image.unsqueeze(0)
b, n, c, w, h = input_image.size()
input_image = input_image.view(b * n, c, w, h)
if torch.cuda.is_available():
input_image = input_image.cuda()
with torch.no_grad():
logit, _, _ = self.net(input_image)
logit = logit.view(b, n, 2)
logit = torch.mean(logit, dim=1, keepdim=False)
prob = F.softmax(logit, 1)
print('probabilistic: ', prob)
print('predict: ', np.argmax(prob.detach().cpu().numpy()))
return np.argmax(prob.detach().cpu().numpy())
class FaceBagNet:
def __init__(self, model_path, patch_size=48, torch_device="cpu"):
# TODO: bn, id_class?
self.model_path = model_path
self.patch_size = patch_size
self.neural_net = Net(num_class=2, id_class=300, is_first_bn=True)
self.neural_net.load_pretrain(self.model_path)
self.neural_net = torch.nn.DataParallel(self.neural_net)
self.neural_net.to(torch_device)
self.torch_device = torch_device
# TODO: this line
self.neural_net.eval()
self.augmentor = color_augumentor
# returns probability that the image is genuine (not presented)
def predict(self, full_size_image):
image = deepcopy(full_size_image) # TODO: remove copying
image = self.augmentor(image,
target_shape=(self.patch_size, self.patch_size,
3),
is_infer=True)
n = len(image)
image = np.concatenate(image, axis=0)
image = np.transpose(image, (0, 3, 1, 2))
image = image.astype(np.float32)
image = image.reshape([n, 3, self.patch_size, self.patch_size])
image = np.array([image])
image = image / 255.0
input_tensor = torch.FloatTensor(image)
shape = input_tensor.shape
b, n, c, w, h = shape
# print(b, n, c, w, h)
input_tensor = input_tensor.view(b * n, c, w, h)
# inpt = inpt.cuda() if torch.cuda.is_available() else inpt.cpu()
input_tensor = input_tensor.to(self.torch_device)
# print(input_tensor)
with torch.no_grad():
logit, _, _ = self.neural_net(input_tensor)
logit = logit.view(b, n, 2)
logit = torch.mean(logit, dim=1, keepdim=False)
prob = F.softmax(logit, 1)
is_real = list(prob.data.cpu().numpy()[:, 1])[0]
return is_real