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 排除， 由网盘单独上传。
98 lines
3.4 KiB
Python
98 lines
3.4 KiB
Python
import os
|
|
import cv2
|
|
import glob
|
|
import numpy as np
|
|
from utils import landmark_processor
|
|
|
|
from face_enhance.face_gan_pt import FaceGAN
|
|
from time import time
|
|
|
|
class FaceEnhancement(object):
|
|
def __init__(self, size=512, gpu_id=0):
|
|
self.facegan = FaceGAN(size, gpu_id)
|
|
self.size = size
|
|
self.threshold = 0.9
|
|
|
|
# the mask for pasting restored faces back
|
|
self.mask = np.zeros((512, 512), np.float32)
|
|
cv2.rectangle(self.mask, (26, 26), (486, 486), (1, 1, 1), -1, cv2.LINE_AA)
|
|
self.mask = cv2.GaussianBlur(self.mask, (101, 101), 11)
|
|
self.mask = cv2.GaussianBlur(self.mask, (101, 101), 11)
|
|
|
|
self.kernel = np.array((
|
|
[0.0625, 0.125, 0.0625],
|
|
[0.125, 0.25, 0.125],
|
|
[0.0625, 0.125, 0.0625]), dtype="float32")
|
|
|
|
def process(self, img, landmarks1k):
|
|
|
|
assert len(landmarks1k) == 1000
|
|
|
|
image_to_face_mat = landmark_processor.get_transform_mat_face_restore(landmarks1k, self.size)
|
|
tfm_inv = cv2.invertAffineTransform(image_to_face_mat)
|
|
|
|
height, width = img.shape[:2]
|
|
full_mask = np.zeros((height, width), dtype=np.float32)
|
|
full_img = np.zeros(img.shape, dtype=np.uint8)
|
|
|
|
fh, fw = (landmarks1k[0][1]-landmarks1k[154][1]), (landmarks1k[95][0]-landmarks1k[215][0])
|
|
|
|
of = cv2.warpAffine(img, image_to_face_mat, (self.size, self.size), flags=cv2.INTER_LANCZOS4, borderMode=cv2.BORDER_CONSTANT, borderValue=[0, 0, 0])
|
|
|
|
# enhance the face
|
|
ef = self.facegan.process(of)
|
|
tmp_mask = self.mask
|
|
tmp_mask = cv2.resize(tmp_mask, ef.shape[:2])
|
|
tmp_mask = cv2.warpAffine(tmp_mask, tfm_inv, (width, height), flags=3)
|
|
|
|
if min(fh, fw)<100: # gaussian filter for small faces
|
|
ef = cv2.filter2D(ef, -1, self.kernel)
|
|
|
|
# tmp_img = cv2.warpAffine(ef, tfm_inv, (width, height), flags=3)
|
|
tmp_img = cv2.warpAffine(ef, tfm_inv, (width, height), dst=img.copy(), borderMode=cv2.BORDER_TRANSPARENT)
|
|
|
|
# cv2.imshow("tmp_img: ", tmp_img)
|
|
|
|
mask = tmp_mask - full_mask
|
|
full_mask[np.where(mask>0)] = tmp_mask[np.where(mask>0)]
|
|
full_img[np.where(mask>0)] = tmp_img[np.where(mask>0)]
|
|
|
|
full_mask = full_mask[:, :, np.newaxis]
|
|
img = cv2.convertScaleAbs(img*(1-full_mask) + full_img*full_mask)
|
|
|
|
# cv2.imshow("img: ", img)
|
|
|
|
# cv2.waitKey()
|
|
|
|
return img
|
|
|
|
if __name__=='__main__':
|
|
|
|
indir = '/media/DATA_4T/zao_data/tencent_con_0701/ref_c/res_contrast_blur13_notps_yunfu_res'
|
|
outdir = '/media/DATA_4T/zao_data/tencent_con_0701/ref_c/res_contrast_blur13_notps_yunfu_res_outs2'
|
|
os.makedirs(outdir, exist_ok=True)
|
|
|
|
faceenhancer = FaceEnhancement(base_dir="./", size=512, model="GPEN-512", channel_multiplier=2)
|
|
|
|
files = sorted(glob.glob(os.path.join(indir, '*.*g')))
|
|
for n, file in enumerate(files[:]):
|
|
filename = os.path.basename(file)
|
|
txtname = file.replace(".jpg", "_landmark1k.txt")
|
|
|
|
im = cv2.imread(file, cv2.IMREAD_COLOR) # BGR
|
|
print(txtname)
|
|
landmark = np.loadtxt(txtname)
|
|
if not isinstance(im, np.ndarray): print(filename, 'error'); continue
|
|
|
|
start = time()
|
|
|
|
img = faceenhancer.process(im, landmark)
|
|
|
|
end = time()
|
|
|
|
print("Time cost: {:.4f}".format(end - start))
|
|
|
|
cv2.imwrite(os.path.join(outdir, '.'.join(filename.split('.')[:-1])+'_2.jpg'), img)
|
|
|
|
|