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 排除， 由网盘单独上传。
178 lines
7.2 KiB
Python
178 lines
7.2 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
|
|
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()
|