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 排除， 由网盘单独上传。
69 lines
2.6 KiB
Python
69 lines
2.6 KiB
Python
import os
|
|
import time
|
|
import torch
|
|
import torch.nn.parallel
|
|
|
|
import numpy as np
|
|
from utils import landmark_processor
|
|
import cv2
|
|
|
|
# modelRoot = "/home/yangchaojie/Desktop/hairstyle/hairstyle_infer/weights"
|
|
modelRoot = "./weights"
|
|
|
|
label_map = [
|
|
[0, 0, 0],
|
|
[255, 0, 0],
|
|
[0, 0, 255],
|
|
[0, 255, 0],
|
|
[0, 255, 255],
|
|
# [0, 255, 0]
|
|
]
|
|
class Generator_BaldSeg_5c(object):
|
|
def __init__(self, gpu_flag, gpu_id):
|
|
|
|
if not gpu_flag:
|
|
self.device = torch.device("cpu")
|
|
else:
|
|
self.device = torch.device('cuda:{0}'.format(gpu_id))
|
|
|
|
# load seg model
|
|
self.model_dir = modelRoot
|
|
self.pre_trained_model = os.path.join(self.model_dir, "ori_hair_checkpoint_7660_0611.pt")
|
|
self.net = torch.jit.load(self.pre_trained_model, map_location=self.device).to(self.device)
|
|
self.net.eval()
|
|
|
|
self.output_img_size = 512
|
|
self.img_ratio = 0.4
|
|
|
|
def label_to_mask(self, label_np):
|
|
label_np = label_np.astype(np.int32)[:, :, np.newaxis]
|
|
mask = np.zeros((label_np.shape[0], label_np.shape[1], 3), dtype=np.uint8)
|
|
for id, color in enumerate(label_map):
|
|
index = (label_np == id).all(axis=2)
|
|
mask[index] = color
|
|
return mask
|
|
def forward(self, image, alpha, landmarks1k):
|
|
image_to_face_mat = landmark_processor.get_transform_mat_full_face_ratio_deeplab(landmarks1k, self.output_img_size, self.img_ratio)
|
|
img_ = (image * (1 - alpha[:, :, np.newaxis].astype(np.float32) / 255)).astype(np.uint8)
|
|
img_cuted = cv2.warpAffine(img_, image_to_face_mat, (self.output_img_size, self.output_img_size), flags=cv2.INTER_LANCZOS4, borderMode=cv2.BORDER_CONSTANT, borderValue=[255, 255, 255])
|
|
inv_image_to_face_mat = cv2.invertAffineTransform(image_to_face_mat)
|
|
|
|
img = (img_cuted.astype(np.float32) / 255).transpose((2, 0, 1))
|
|
img = np.expand_dims(img, axis=0)
|
|
img = torch.from_numpy(img).to(self.device) #cuda(self.gpu_id)
|
|
|
|
with torch.no_grad():
|
|
output = self.net(img)
|
|
|
|
pred = output.detach().cpu().numpy().squeeze().astype(np.float32) #torch.max(output[:1], 1)[1].detach().cpu().numpy().squeeze().astype(np.float32)
|
|
mask = self.label_to_mask(pred)
|
|
|
|
# cv2.imshow("img_cuted: ", img_cuted)
|
|
# cv2.imshow("mask: ", mask)
|
|
# cv2.waitKey()
|
|
|
|
black_img = np.zeros(image.shape).astype(np.uint8)
|
|
cv2.warpAffine(mask, inv_image_to_face_mat, (image.shape[1], image.shape[0]),
|
|
dst=black_img, flags=cv2.INTER_NEAREST, borderMode=cv2.BORDER_TRANSPARENT)
|
|
|
|
return black_img |