初始化：换发型/换发色/训练发型服务
包含： - 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密钥已脱敏为环境变量，原文件备份在本地
This commit is contained in:
xsl
@@ -0,0 +1,97 @@
|
||||
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)
|
||||
|
||||
|
||||
@@ -0,0 +1,66 @@
|
||||
'''
|
||||
@paper: GAN Prior Embedded Network for Blind Face Restoration in the Wild (CVPR2021)
|
||||
@author: yangxy (yangtao9009@gmail.com)
|
||||
'''
|
||||
import torch
|
||||
import os
|
||||
import cv2
|
||||
import numpy as np
|
||||
# modelRoot = "/home/yangchaojie/Desktop/hairstyle/hairstyle_infer/weights"
|
||||
modelRoot = "./weights"
|
||||
class FaceGAN(object):
|
||||
def __init__(self, size=512, gpu_id=0):
|
||||
# self.mfile = os.path.join(base_dir, model+'.pth')
|
||||
self.n_mlp = 8
|
||||
self.resolution = size
|
||||
self.device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
|
||||
|
||||
self.load_model()
|
||||
|
||||
def load_model(self):
|
||||
self.model_dir = modelRoot
|
||||
self.face_gan_model = os.path.join(self.model_dir, "face_enhance_0630.pt")
|
||||
|
||||
self.model = torch.jit.load(self.face_gan_model).to(self.device)
|
||||
|
||||
self.model.eval()
|
||||
|
||||
def process_o(self, img):
|
||||
img = cv2.resize(img, (self.resolution, self.resolution))
|
||||
img_t = self.img2tensor(img)
|
||||
|
||||
with torch.no_grad():
|
||||
out, __ = self.model(img_t)
|
||||
|
||||
out = self.tensor2img(out)
|
||||
|
||||
return out
|
||||
|
||||
def process(self, img):
|
||||
img = cv2.resize(img, (self.resolution, self.resolution))
|
||||
img_t = self.img2tensor(img)
|
||||
|
||||
with torch.no_grad():
|
||||
out = self.forward(img_t)
|
||||
|
||||
out = self.tensor2img(out)
|
||||
|
||||
return out
|
||||
|
||||
def forward(self, img_t):
|
||||
with torch.no_grad():
|
||||
out = self.model(img_t)
|
||||
|
||||
return out
|
||||
|
||||
def img2tensor(self, img):
|
||||
img_t = (torch.from_numpy(img).to(self.device)/255. - 0.5) / 0.5
|
||||
img_t = img_t.permute(2, 0, 1).unsqueeze(0).flip(1) # BGR->RGB
|
||||
return img_t
|
||||
|
||||
def tensor2img(self, image_tensor, pmax=255.0, imtype=np.uint8):
|
||||
image_tensor = image_tensor * 0.5 + 0.5
|
||||
image_tensor = image_tensor.squeeze(0).permute(1, 2, 0).flip(2) # RGB->BGR
|
||||
image_numpy = np.clip(image_tensor.float().cpu().numpy(), 0, 1) * pmax
|
||||
|
||||
return image_numpy.astype(imtype)
|
||||
@@ -0,0 +1,102 @@
|
||||
#!/usr/bin/env python
|
||||
# -*- coding: utf-8 -*-
|
||||
|
||||
# @File : setup.py
|
||||
# @Time : 2020/1/15
|
||||
# @Author : yangchaojie (yangchaojie@immomo.com)
|
||||
|
||||
import os
|
||||
import sys
|
||||
import shutil
|
||||
import numpy
|
||||
import tempfile
|
||||
|
||||
from setuptools import setup
|
||||
from setuptools.extension import Extension
|
||||
|
||||
from Cython.Build import cythonize
|
||||
from Cython.Distutils import build_ext
|
||||
|
||||
import platform
|
||||
|
||||
|
||||
def get_root_path(root):
|
||||
if os.path.dirname(root) in ['', '.']:
|
||||
return os.path.basename(root)
|
||||
else:
|
||||
return get_root_path(os.path.dirname(root))
|
||||
|
||||
|
||||
def copy_file(src, dest):
|
||||
if os.path.exists(dest):
|
||||
return
|
||||
|
||||
if not os.path.exists(os.path.dirname(dest)):
|
||||
os.makedirs(os.path.dirname(dest))
|
||||
if os.path.isdir(src):
|
||||
shutil.copytree(src, dest)
|
||||
else:
|
||||
shutil.copyfile(src, dest)
|
||||
|
||||
|
||||
def touch_init_file():
|
||||
init_file_name = os.path.join(tempfile.mkdtemp(), '__init__.py')
|
||||
with open(init_file_name, 'w'):
|
||||
pass
|
||||
return init_file_name
|
||||
|
||||
|
||||
|
||||
|
||||
def compose_extensions(root='.'):
|
||||
for file_ in os.listdir(root):
|
||||
abs_file = os.path.join(root, file_)
|
||||
|
||||
if os.path.isfile(abs_file):
|
||||
if abs_file.endswith('.py'):
|
||||
extensions.append(Extension(get_root_path(abs_file) + '.*', [abs_file]))
|
||||
elif abs_file.endswith('.c') or abs_file.endswith('.pyc'):
|
||||
continue
|
||||
else:
|
||||
copy_file(abs_file, os.path.join(build_root_dir, abs_file))
|
||||
if abs_file.endswith('__init__.py'):
|
||||
copy_file(init_file, os.path.join(build_root_dir, abs_file))
|
||||
|
||||
else:
|
||||
if os.path.basename(abs_file) in ignore_folders :
|
||||
continue
|
||||
if os.path.basename(abs_file) in conf_folders:
|
||||
copy_file(abs_file, os.path.join(build_root_dir, abs_file))
|
||||
compose_extensions(abs_file)
|
||||
|
||||
|
||||
|
||||
# if __name__ == '__main__':
|
||||
build_root_dir = 'build/lib.' + platform.system().lower() + '-' + platform.machine() + '-' + str(
|
||||
sys.version_info.major) + '.' + str(sys.version_info.minor)
|
||||
|
||||
print(build_root_dir)
|
||||
|
||||
extensions = []
|
||||
ignore_folders = ['build', 'new_ref_zao_color_0818', 'ref_hair_online_0703_local', 'ref_分好类别', '.git']
|
||||
conf_folders = ['conf']
|
||||
|
||||
|
||||
init_file = touch_init_file()
|
||||
print(init_file)
|
||||
|
||||
|
||||
compose_extensions()
|
||||
os.remove(init_file)
|
||||
|
||||
setup(
|
||||
name='moxie_hairstyle',
|
||||
version='1.0',
|
||||
ext_modules=cythonize(
|
||||
extensions,
|
||||
nthreads=16,
|
||||
compiler_directives=dict(always_allow_keywords=True),
|
||||
include_path=[numpy.get_include()]),
|
||||
cmdclass=dict(build_ext=build_ext))
|
||||
|
||||
# python setup.py build_ext
|
||||
Reference in New Issue
Block a user