colomi
0eb61f3e60
包含： - 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 排除， 由网盘单独上传。
383 lines
17 KiB
Python
383 lines
17 KiB
Python
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()
|