xsl
443cfa298f
包含： - hair_service_sd: 主服务（换发型/换发色/生发，端口8801） - photo_service: LoRA调度+训练（端口32678） - hair_grow_service: 调试测试页（端口8888，含4个测试页） - 批量训练脚本（batch_train_hairstyles.py） - 发际线mask自动识别（hairline_mask.py，4种方案） - 手绘mask换发型（hair_swap_manual.py） - 文档：README.md + LARGE_FILES.md + docs/ 大文件（模型权重200G、训练数据123G）已排除，见 LARGE_FILES.md OSS/COS密钥已脱敏为环境变量，原文件备份在本地
251 lines
11 KiB
Python
251 lines
11 KiB
Python
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 = []
|