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