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密钥已脱敏为环境变量，原文件备份在本地
75 lines
3.0 KiB
Python
75 lines
3.0 KiB
Python
import os
|
|
import pickle
|
|
|
|
import torch
|
|
from faceseg.u2net import U2NET
|
|
from utils import landmark_processor
|
|
import cv2
|
|
import numpy as np
|
|
|
|
|
|
class FaceSeg:
|
|
def __init__(self, gpu_id = 0):
|
|
model = U2NET(in_ch=4, out_ch=1)
|
|
weights = torch.load('weights/20210927_01.pth', map_location='cpu')
|
|
model_dict = model.state_dict()
|
|
pretrained_dict = {}
|
|
for ix, (k, v) in enumerate(model_dict.items()):
|
|
if k in weights and weights[k].data.shape == v.data.shape:
|
|
pretrained_dict[k] = weights[k]
|
|
else:
|
|
print('ignore {}'.format(k))
|
|
model_dict.update(pretrained_dict)
|
|
model.load_state_dict(model_dict)
|
|
print('update success')
|
|
model.cuda(gpu_id)
|
|
model.eval()
|
|
self.model = model
|
|
self.last_mask = None
|
|
self.output_img_size = 320
|
|
self.gpu_id = gpu_id
|
|
|
|
def inference(self, frame, pt1k, video_mode=False):
|
|
image_to_face_mat = landmark_processor.get_transform_mat_full_face(pt1k, self.output_img_size)
|
|
face_image = cv2.warpAffine(frame, image_to_face_mat, (self.output_img_size, self.output_img_size), flags=cv2.INTER_LANCZOS4)
|
|
if face_image.dtype == np.uint8: face_image = face_image.astype(np.float32) / 255
|
|
if video_mode and self.last_mask is not None:
|
|
last_small_mask = cv2.warpAffine(self.last_mask, image_to_face_mat,
|
|
(self.output_img_size, self.output_img_size), flags=cv2.INTER_LANCZOS4)[:,:,np.newaxis]
|
|
input_img = np.concatenate([face_image, last_small_mask], axis=2)
|
|
else:
|
|
zero_mask = np.zeros((face_image.shape[1], face_image.shape[0], 1), dtype=np.float32)
|
|
input_img = np.concatenate([face_image, zero_mask], axis=2)
|
|
|
|
face_image_tensor = input_img.transpose((2, 0, 1))[np.newaxis]
|
|
face_image_tensor = torch.from_numpy(face_image_tensor).cuda(self.gpu_id)
|
|
|
|
mask = self.model.test(face_image_tensor)
|
|
mask = mask[0].detach().cpu().numpy().transpose((1, 2, 0))
|
|
|
|
origin_mask = cv2.warpAffine(mask, image_to_face_mat, (frame.shape[1], frame.shape[0]),
|
|
flags=cv2.WARP_INVERSE_MAP|cv2.INTER_LANCZOS4)[:, :, np.newaxis]
|
|
if video_mode: self.last_mask = origin_mask.copy()
|
|
|
|
return origin_mask
|
|
|
|
|
|
if __name__ == '__main__':
|
|
face_segmentor = FaceSeg(gpu_id=0)
|
|
testdata_dir = "/mnt/DataDisk/my_projects/faceswap_hq/train_data/example"
|
|
for picname in os.listdir(testdata_dir):
|
|
img_path = os.path.join(testdata_dir, picname)
|
|
pkl_path = img_path[:-4]+".pkl"
|
|
if not picname.endswith(".jpg"):
|
|
continue
|
|
if not os.path.exists(pkl_path):
|
|
continue
|
|
img = cv2.imread(img_path)
|
|
with open(pkl_path, "rb") as fp:
|
|
info = pickle.load(fp)
|
|
pt1k = info["human_pt1k"]
|
|
|
|
face_seg_mask = face_segmentor.inference(img, pt1k, video_mode=False)
|
|
cv2.imshow("face_seg_mask", face_seg_mask)
|
|
cv2.imshow("img", img)
|
|
cv2.waitKey() |