初始化换发型项目:3个微服务代码 + 部署脚本

包含:
- hair_service_sd: 换发型/换发色算法服务 (端口 8801)
- photo_service: LoRA 训练调度服务 (端口 32678)
- stable-diffusion-webui: SD WebUI 推理服务 (端口 57860)
- kohya_ss_home: 训练环境代码
- meidaojia: 监控测试脚本
- setup.sh: 一键部署脚本 (conda环境恢复 + 配置生成 + 完整性检查)
- start_all_services.sh: 启动3个服务
- configure.ini.template: 路径模板化 (BASE_DIR自动推导)
- conda_envs/py310.yml: py310 环境定义

大文件 (weights/, models/, data/, conda_envs/*.tar.gz 等) 通过 .gitignore 排除,
由网盘单独上传。
This commit is contained in:
colomi
2026-07-11 18:11:49 +08:00
commit 0eb61f3e60
628 changed files with 120882 additions and 0 deletions
@@ -0,0 +1,3 @@
version https://git-lfs.github.com/spec/v1
oid sha256:ece946bd0ae8e4142eb8667f8b273b75a9564a93b08656ae50ff064c7cff08bf
size 7375394
+382
View File
@@ -0,0 +1,382 @@
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()
+3
View File
@@ -0,0 +1,3 @@
version https://git-lfs.github.com/spec/v1
oid sha256:a34ffdb26cd6780d41cb3ed313785c09bb7d60634a3f24237303fe7b2746d7e5
size 12271767
+3
View File
@@ -0,0 +1,3 @@
version https://git-lfs.github.com/spec/v1
oid sha256:4236333ca60a7bffcb13e22e869f738593362e4d407e87e55b579401eeb05be8
size 45329800
+250
View File
@@ -0,0 +1,250 @@
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 = []
+3
View File
@@ -0,0 +1,3 @@
version https://git-lfs.github.com/spec/v1
oid sha256:30238ab182eba7291094d6be816b7ca4062f9abf616a101eb0c91424cab03b59
size 3216228
+94
View File
@@ -0,0 +1,94 @@
import sys
import os
import cv2
import torch
import torch.nn as nn
def op_name(op_name, m):
m.op_name = op_name
return m
def conv_bn_relu(name, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, momentum = 0.1, track_running_stats=True):
return nn.Sequential(
op_name(name, nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, False)),
op_name(name + '/bn', nn.BatchNorm2d(out_channels, momentum = momentum, track_running_stats = track_running_stats)),
op_name(name + '/relu', nn.ReLU(inplace=True)),
)
# Define a Basic resnet block
class BasicResnetBlock(nn.Module):
def __init__(self, name, in_channels, out_channels, kernel_size=3, stride=2, padding=1, direct_plus=False):
super(BasicResnetBlock, self).__init__()
self.op_name = name
self.conv_block = nn.Sequential(
conv_bn_relu(name=name + '/conv1', in_channels=in_channels, out_channels=out_channels,
kernel_size=kernel_size, stride=stride, padding=padding),
conv_bn_relu(name=name + '/conv2', in_channels=out_channels, out_channels=out_channels,
kernel_size=3, stride=1, padding=1)
)
if direct_plus and (in_channels == out_channels) and (stride == 1):
self.sc_block = None
else:
self.sc_block = conv_bn_relu(name=name + '/sc_conv', in_channels=in_channels, out_channels=out_channels,
kernel_size=1, stride=stride, padding=0)
def forward(self, x):
if self.sc_block is None:
out = x + self.conv_block(x)
else:
out = self.conv_block(x) + self.sc_block(x)
return out
class Flatten(nn.Module):
def __init__(self, axis):
super(Flatten, self).__init__()
self.axis = axis
def forward(self, x):
assert self.axis == 1
x = x.reshape(x.shape[0], -1)
# x = x.view(-1, 1152)
return x
def flatten(name, axis):
return op_name(name, Flatten(axis))
def linear_bn_relu(name, in_features, out_features, momentum = 0.1, track_running_stats=True):
return nn.Sequential(
op_name(name + '/FC', nn.Linear(in_features, out_features)),
op_name(name + '/bn', nn.BatchNorm1d(out_features, momentum = momentum, track_running_stats = track_running_stats)),
op_name(name + '/relu', nn.ReLU(inplace=True)),
)
class LeftEye(nn.Module):
def __init__(self, name, in_channels, out_channels):
super(LeftEye, self).__init__()
self.op_name = name
op_list = []
op_list += [conv_bn_relu(name + '/first_conv', in_channels, 24, kernel_size = 5, stride = 2, padding = 2) ]
ch_num = [24, 32, 64, 96, 128]
op_list += [BasicResnetBlock(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),
linear_bn_relu(name + '/FC1', 1152, 256),
op_name(name + '/FC2', nn.Linear(256, out_channels))]
self.conv_block = nn.Sequential(*op_list)
def forward(self, x):
return self.conv_block(x)
def get_left_eye_symbol(symbol_name='BigResNet', input_nc = 3, output_nc = 17 * 2):
return LeftEye(symbol_name, input_nc, output_nc)
if __name__=='__main__':
pass
@@ -0,0 +1,147 @@
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.PReLU(out_channels)),
)
class BasicBlock(nn.Module):
def __init__(self, i, j, inplanes):
super(BasicBlock, self).__init__()
self.conv1 = conv_relu('conv{}_{}'.format(i, 2 * j), inplanes, inplanes, kernel_size=3, stride=1, padding=1)
self.conv2 = conv_relu('conv{}_{}'.format(i, 2 * j + 1), inplanes, inplanes, kernel_size=3, stride=1, padding=1)
def forward(self, x):
residule = x
out = self.conv1(x)
out = self.conv2(out)
ret = out + residule
return ret
class ResnetBlock(nn.Module):
def __init__(self, i, inplanes, outplanes, stride=2, n_blocks=1):
super(ResnetBlock, self).__init__()
self.conv = [conv_relu('conv{}_{}'.format(i, 1), inplanes, outplanes, kernel_size=3, stride=stride, padding=1)]
for m in range(n_blocks):
self.conv.append(BasicBlock(i, m + 1, outplanes))
self.conv = nn.Sequential(*self.conv)
def forward(self, x):
ret = self.conv(x)
return ret
class FaceRecognition(nn.Module):
def __init__(self):
super(FaceRecognition, self).__init__()
op_list = []
ch_num = [3, 32, 64, 128, 128]
strides = [2, 2, 2, 2]
n_blocks = [1, 2, 4, 1]
op_list = []
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 += [flatten('flatten', 1)]
op_list += [op_name('fc5', nn.Linear(4608, 256))]
self.features = nn.Sequential(*op_list)
model_path, _ = os.path.split(os.path.realpath(__file__))
weights = torch.load(os.path.join(model_path, 'Face_ResNet.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 MomocvFaceRecognition(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 = FaceRecognition()
# self.model_path, _ = os.path.split(os.path.realpath(__file__))
# weights = torch.load(os.path.join(self.model_path, 'Face_ResNet.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 = 90
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) - 128) / 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'E:\deepfacelab_data\expression_dst\7201806132018061208311920180612083119\*.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 = MomocvFaceRecognition()
features = mmcv.forward([img, img], [landmarks, 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()
@@ -0,0 +1,3 @@
version https://git-lfs.github.com/spec/v1
oid sha256:752d0396a3a8ce943d7a915bb15063f27112ccb2dd8f4afa2bb9a42824eafbd9
size 174375530
@@ -0,0 +1,177 @@
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()
+311
View File
@@ -0,0 +1,311 @@
import torch.nn as nn
import torch.utils.model_zoo as model_zoo
import torch
import os
import numpy as np
import cv2
from utils import landmark_processor, utils_3ddfa, params_3ddfa
__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, no_branch=False, no_activate=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 no_activate:
if no_branch:
self.fc = nn.Sequential(*[nn.Linear(512 * block.expansion, num_classes)])
else:
self.fc_key = nn.Sequential(*[nn.Linear(256 * block.expansion, 45 * 2)])
self.fc_ctrl = nn.Sequential(*[nn.Linear(256 * block.expansion, 48 * 2)])
else:
if no_branch:
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.no_branch = no_branch
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.no_branch:
key = self.avgpool(x)
key = key.view(-1, 512)
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 Model_3DDFA(nn.Module):
def __init__(self, gpu_id=None):
super(Model_3DDFA, self).__init__()
self.face_alignment_net = resnet18(pretrained=False, num_classes=76, no_branch=True, no_activate=True)
model_path, _ = os.path.split(os.path.realpath(__file__))
weights = torch.load(os.path.join(model_path, 'face_3ddfa.pth'), map_location=lambda storage, loc: storage)
self.load_state_dict(weights)
self.eval()
self.device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
self.to(self.device)
def forward(self, imgs):
pred_pose_shape_exp = self.face_alignment_net(imgs)
return pred_pose_shape_exp
def forward_np(self, imgs):
pred_pose_shape_exp = self.face_alignment_net(imgs)
return pred_pose_shape_exp.detach().cpu().numpy()
def detect(self, images, landmarks):
dst_size = 256
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 = []
all_res = []
all_height = []
for ix, img in enumerate(images):
landmark = landmarks[ix]
mat = landmark_processor.get_transform_mat_full_face(landmark, dst_size)
all_mat.append(mat)
all_height.append(img.shape[0])
tmp = cv2.warpAffine(img, mat, (dst_size, dst_size))
input_numpy[ix, :, :, :] = tmp.transpose((2, 0, 1)).astype(np.float32) / 255
# cv2.imshow('3ddfa_', tmp)
# cv2.waitKey()
in_tensor = torch.from_numpy(input_numpy)
in_tensor = in_tensor.to(self.device)
params = self.face_alignment_net(in_tensor)
params = params.cpu().numpy()
for ix, param in enumerate(params):
param[0] = param[0] / params_3ddfa.SCALE_F
param[1:4] = param[1:4] / params_3ddfa.SCALE_ROTATE
param[4:6] = param[4:6] / params_3ddfa.SCALE_OFFSET
param[6:56] = (param[6:56] / params_3ddfa.SCALE_SHAPE)
param[56:] = (param[56:] / params_3ddfa.SCALE_EXP)
new_param = utils_3ddfa.transform_params(param, cv2.invertAffineTransform(all_mat[ix]), all_height[ix], dst_size)
all_res.append(new_param)
return all_res
@@ -0,0 +1,452 @@
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
+229
View File
@@ -0,0 +1,229 @@
import torch
import torch.nn as nn
import math
import sys
import os
from utils import landmark_processor
import cv2
import numpy as np
output_points = 137 * 2
input_size = 192
def op_name(op_name, m):
m.op_name = op_name
return m
class Flatten(nn.Module):
def __init__(self, axis):
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):
return op_name(name, Flatten(axis))
def conv_bn_relu(name, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, momentum = 0.1, track_running_stats=True):
return nn.Sequential(
op_name(name, nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, False)),
op_name(name + '/bn', nn.BatchNorm2d(out_channels, momentum = momentum, track_running_stats = track_running_stats)),
op_name(name + '/relu', nn.ReLU(inplace=True)),
)
def linear_bn_relu(name, in_features, out_features, momentum = 0.1, track_running_stats=True):
return nn.Sequential(
op_name(name + '/FC', nn.Linear(in_features, out_features)),
op_name(name + '/bn', nn.BatchNorm1d(out_features, momentum = momentum, track_running_stats = track_running_stats)),
op_name(name + '/relu', nn.ReLU(inplace=True)),
)
# Define a Basic resnet block
class BasicResnetBlock(nn.Module):
def __init__(self, name, in_channels, out_channels, kernel_size=3, stride=2, padding=1, direct_plus=False):
super(BasicResnetBlock, self).__init__()
self.op_name = name
self.conv_block = nn.Sequential(
conv_bn_relu(name=name + '/conv1', in_channels=in_channels, out_channels=out_channels,
kernel_size=kernel_size, stride=stride, padding=padding),
conv_bn_relu(name=name + '/conv2', in_channels=out_channels, out_channels=out_channels,
kernel_size=3, stride=1, padding=1)
)
if direct_plus and (in_channels == out_channels) and (stride == 1):
self.sc_block = None
else:
self.sc_block = conv_bn_relu(name=name + '/sc_conv', in_channels=in_channels, out_channels=out_channels,
kernel_size=1, stride=stride, padding=0)
def forward(self, x):
if self.sc_block is None:
out = x + self.conv_block(x)
else:
out = self.conv_block(x) + self.sc_block(x)
return out
#define SuperBigResNet Frame
class SuperBigResNet(nn.Module):
def __init__(self, name='SuperBigResNet', in_channels=3, out_channels=137 * 2):
super(SuperBigResNet, self).__init__()
self.op_name = name
op_list = []
stride_cnt = 2
op_list += [conv_bn_relu(name + '/first_conv', in_channels, 48, kernel_size=5, stride=2, padding=2)]
ch_num = [48, 64, 96, 128, 192]
for i in range(len(ch_num) - 1):
if ch_num[i] == ch_num[i+1]:
op_list += [BasicResnetBlock(name + '/stage%d'%(i + 1), ch_num[i], ch_num[i+1], kernel_size=3, stride=1, direct_plus=True)]
else:
stride_cnt = stride_cnt * 2
op_list += [BasicResnetBlock(name + '/stage%d'%(i + 1), ch_num[i], ch_num[i+1], kernel_size=3, stride=2)]
item_cnt = int((input_size / stride_cnt) * (input_size / stride_cnt) * ch_num[-1])
op_list += [flatten(name + '/flatten', 1),
linear_bn_relu(name + '/FC1', item_cnt, 512),
op_name(name + '/FC2', nn.Linear(512, out_channels))]
self.conv_block = nn.Sequential(*op_list)
def forward(self, x):
res = self.conv_block(x)
return res.cpu().numpy()
class MomocvFaceAlignmentBigger(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 = SuperBigResNet('SuperBigResNet', 3, 137 * 2)
self.model_path, _ = os.path.split(os.path.realpath(__file__))
weights = torch.load(os.path.join(self.model_path, 'SuperBigResNet_latest.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 detect(self, img, landmarks):
dst_size = 192
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)
# cv2.imshow('tmp', tmp)
# cv2.waitKey()
in_tensor = torch.from_numpy(input_numpy)
in_tensor = in_tensor.to(self.device)
fullyconnected1 = self.face_alignment_net(in_tensor)
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)
landmarks_res.append(orig_pts)
return landmarks_res
def stable_forward(self, image, detected_faces, reset=False):
if reset is True:
self.trackingFaceRects = []
for face_rect in detected_faces:
if len(self.trackingFaceRects) == 0:
new_tracking_rect = [face_rect, True, [0, 0], 0, None]
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 = 192 * 0.6 / min(d[2] - d[0], d[3] - d[1])
dst_center = np.array([0.5, 0.5]) * 192
offset = dst_center - src_center
print('hello')
M = cv2.getRotationMatrix2D((src_center[0], src_center[1]), rotate_degree, scale)
M[:, 2] += offset
else:
# dst_left_anchor = np.array([0.403, 0.52]) * 192
# dst_right_anchor = np.array([1 - 0.403, 0.52]) * 192
# # 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][2], tracking_face_rect[2][3])
# scale = 71.0 / src_anchor_len
#
# # bounding_box = cv2.boundingRect(np.array(tracking_face_rect[2]))
# # a = bounding_box[2] * bounding_box[3]
# # print('a', a)
# # scale = float(65*65) / (bounding_box[2] * bounding_box[3])
# # print('scale', scale)
#
# dst_center = (dst_left_anchor + dst_right_anchor) / 2
# offset = dst_center - src_center
rotate_degree = 0
M = landmark_processor.get_transform_mat_mmcv_bigger(tracking_face_rect[4], 192)
inp = cv2.warpAffine(image, M, (192, 192))
orig_inp = inp
inp = inp.transpose((2, 0, 1)).astype(np.float32)
inp = inp[np.newaxis, :, :, :]
t0 = cv2.getTickCount()
in_tensor = torch.from_numpy(inp)
in_tensor = in_tensor.to(self.device)
fullyconnected1 = self.face_alignment_net(in_tensor)
fullyconnected1 = fullyconnected1[0]
orig_pts = (np.reshape(fullyconnected1, (2, 137)).transpose((1, 0)) * 192)
t2 = cv2.getTickCount()
# print('cost', (t2 - t0) / cv2.getTickFrequency() * 1000)
# for ix, pt in enumerate(orig_pts):
# cv2.putText(orig_inp, str(ix), (int(pt[0]), int(pt[1])), cv2.FONT_HERSHEY_SIMPLEX, 0.3, (255, 0, 0))
# cv2.circle(orig_inp, (int(pt[0]), int(pt[1])), 1, (0, 255, 0), 1)
# cv2.imshow('inp_stable', orig_inp)
# cv2.waitKey()
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] = [fullyconnected1[68], fullyconnected1[74], fullyconnected1[96], fullyconnected1[113]]
tracking_face_rect[3] = rotate_degree
tracking_face_rect[4] = fullyconnected1
landmarks.append(fullyconnected1)
return landmarks
if __name__=='__main__':
pass
+102
View File
@@ -0,0 +1,102 @@
import torch
import torch.nn as nn
import math
import torch.optim as optim
import sys
import os
def op_name(op_name, m):
m.op_name = op_name
return m
class Flatten(nn.Module):
def __init__(self, axis):
super(Flatten, self).__init__()
def forward(self, x):
x = x.view(-1, 2560)
return x
def flatten(name, axis):
return op_name(name, Flatten(axis))
def conv_bn_relu(name, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, momentum = 0.1, track_running_stats=True):
return nn.Sequential(
op_name(name, nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, False)),
op_name(name + '/bn', nn.BatchNorm2d(out_channels, momentum = momentum, track_running_stats = track_running_stats)),
op_name(name + '/relu', nn.ReLU(inplace=True)),
)
def linear_bn_relu(name, in_features, out_features, momentum = 0.1, track_running_stats=True):
return nn.Sequential(
op_name(name + '/FC', nn.Linear(in_features, out_features)),
op_name(name + '/bn', nn.BatchNorm1d(out_features, momentum = momentum, track_running_stats = track_running_stats)),
op_name(name + '/relu', nn.ReLU(inplace=True)),
)
# Define a Basic resnet block
class BasicResnetBlock(nn.Module):
def __init__(self, name, in_channels, out_channels, kernel_size=3, stride=2, padding=1, direct_plus=False):
super(BasicResnetBlock, self).__init__()
self.op_name = name
self.conv_block = nn.Sequential(
conv_bn_relu(name=name + '/conv1', in_channels=in_channels, out_channels=out_channels,
kernel_size=kernel_size, stride=stride, padding=padding),
conv_bn_relu(name=name + '/conv2', in_channels=out_channels, out_channels=out_channels,
kernel_size=3, stride=1, padding=1)
)
if direct_plus and (in_channels == out_channels) and (stride == 1):
self.sc_block = None
else:
self.sc_block = conv_bn_relu(name=name + '/sc_conv', in_channels=in_channels, out_channels=out_channels,
kernel_size=1, stride=stride, padding=0)
def forward(self, x):
if self.sc_block is None:
out = x + self.conv_block(x)
else:
out = self.conv_block(x) + self.sc_block(x)
return out
class SuperBigResNetV2(nn.Module):
def __init__(self, name, in_channels, out_channels):
super(SuperBigResNetV2, self).__init__()
self.op_name = name
input_size = 256
op_list = []
stride_cnt = 4
op_list += [conv_bn_relu(name + '/first_conv', in_channels, 48, kernel_size=5, stride=4, padding=2)]
ch_num = [48, 64, 96, 128, 160]
for i in range(len(ch_num) - 1):
if ch_num[i] == ch_num[i+1]:
op_list += [BasicResnetBlock(name + '/stage%d'%(i + 1), ch_num[i], ch_num[i+1], kernel_size=3, stride=1, direct_plus=True)]
else:
stride_cnt = stride_cnt * 2
op_list += [BasicResnetBlock(name + '/stage%d'%(i + 1), ch_num[i], ch_num[i+1], kernel_size=3, stride=2)]
item_cnt = int((input_size / stride_cnt) * (input_size / stride_cnt) * ch_num[-1])
op_list += [flatten(name + '/flatten', 1),
linear_bn_relu(name + '/FC1', item_cnt, 512),
op_name(name + '/FC2', nn.Linear(512, out_channels))]
self.conv_block = nn.Sequential(*op_list)
self.model_path, _ = os.path.split(os.path.realpath(__file__))
self.weights = torch.load(os.path.join(self.model_path, 'SuperBigResNetV2_latest.pth'),
map_location=lambda storage, loc: storage)
def forward(self, x):
return self.conv_block(x)
if __name__=='__main__':
pass
@@ -0,0 +1,3 @@
version https://git-lfs.github.com/spec/v1
oid sha256:82a5c042c29002a009f83345875a73d39209a435459a1fa46a95f30392d07e92
size 9536531
@@ -0,0 +1,3 @@
version https://git-lfs.github.com/spec/v1
oid sha256:7206aaa727f403c6f8e660e9bf9d9967e6fdb3dced1f77232320e0ba0b5b2515
size 19020307
+3
View File
@@ -0,0 +1,3 @@
version https://git-lfs.github.com/spec/v1
oid sha256:c0d5cc765269c4237fb29d2a825da9d0a39e59813fe2f681003cb1ac4b081f17
size 44926130
@@ -0,0 +1,3 @@
version https://git-lfs.github.com/spec/v1
oid sha256:89042256066e0e4ff692f4cfa8a4727ee55c6ccf8dccb8d1a1907484da283582
size 48873911
+223
View File
@@ -0,0 +1,223 @@
import torch.nn as nn
import torch.utils.model_zoo as model_zoo
__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, 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))
self.fc = nn.Linear(512 * block.expansion, num_classes)
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)
x = self.avgpool(x)
x = x.view(-1, 512)
return self.fc(x)
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']))
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