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 排除， 由网盘单独上传。
248 lines
10 KiB
Python
Executable File
248 lines
10 KiB
Python
Executable File
import cv2
|
|
import os
|
|
import numpy as np
|
|
import tqdm
|
|
|
|
from utils import landmark_processor
|
|
import base64
|
|
import requests
|
|
from PIL import Image
|
|
import io
|
|
|
|
def encode_numpy_to_base64(img):
|
|
retval, bytes = cv2.imencode('.png', img)
|
|
encoded_image = base64.b64encode(bytes).decode('utf-8')
|
|
return encoded_image
|
|
|
|
|
|
def webui_img2img(img, mask, prompt=''):
|
|
url = "http://127.0.0.1:57860/sdapi/v1/img2img"
|
|
request_dict = {
|
|
"prompt": prompt,
|
|
"negative_prompt": '(nsfw:1.5), ng_deepnegative_v1_75t, (badhandv4:1.2), (worst quality:2), (low quality:2), (normal quality:2), lowres, bad anatomy, '
|
|
'bad hands, ((monochrome)), ((grayscale)) watermark, large breast, big breast, bad_pictures,easynegative, faceless, no human, white background, simple background, ',
|
|
"sampler_name": "DPM++ 2M Karras",
|
|
"batch_size": 1,
|
|
"steps": 20,
|
|
"width": img.shape[1],
|
|
"height": img.shape[0],
|
|
"cfg_scale": 7.0,
|
|
"seed": 123456789,
|
|
"mask_blur": 11,
|
|
"init_images": [
|
|
encode_numpy_to_base64(img)
|
|
],
|
|
"inpaint_full_res": False,
|
|
"inpainting_fill": 1,
|
|
"inpainting_mask_invert": 0,
|
|
"mask": encode_numpy_to_base64(mask),
|
|
# "refiner_checkpoint":"majicmixRealistic_v7.safetensors",
|
|
# "refiner_switch_at": 0.4,
|
|
"denoising_strength": 0.7,
|
|
"alwayson_scripts": {
|
|
# "controlnet": {
|
|
# "args": [
|
|
# {
|
|
# "enabled": True,
|
|
# "module": "openpose_full",
|
|
# "model": "openpose",
|
|
# "weight": 1.0,
|
|
# # "image": self.read_image(),
|
|
# "resize_mode": "Crop and Resize",
|
|
# "low_vram": False,
|
|
# "processor_res": 512,
|
|
# "guidance_start": 0.0,
|
|
# "guidance_end": 1.0,
|
|
# "control_mode": "Balanced",
|
|
# "pixel_perfect": False
|
|
# }
|
|
# ]
|
|
# }
|
|
}
|
|
}
|
|
response = requests.post(url=url, json=request_dict)
|
|
ret_json = response.json()
|
|
result = ret_json['images'][0]
|
|
img = cv2.imdecode(np.frombuffer(base64.b64decode(result.split(",", 1)[0]), np.uint8), cv2.IMREAD_COLOR)
|
|
return img
|
|
|
|
def webui_super_res_img(img, ratio):
|
|
url = "http://127.0.0.1:57860/sdapi/v1/extra-single-image"
|
|
request_dict = {
|
|
"resize_mode": 0,
|
|
"show_extras_results": False,
|
|
"gfpgan_visibility": 0,
|
|
"codeformer_visibility": 1,
|
|
"codeformer_weight": 1,
|
|
"upscaling_resize": ratio,
|
|
"upscaler_1": "8x_NMKD-Superscale_150000_G",
|
|
"upscale_first": False,
|
|
"image": encode_numpy_to_base64(img)
|
|
}
|
|
response = requests.post(url=url, json=request_dict)
|
|
ret_json = response.json()
|
|
result = ret_json['image']
|
|
img = cv2.imdecode(np.frombuffer(base64.b64decode(result), np.uint8), cv2.IMREAD_COLOR)
|
|
return img
|
|
|
|
|
|
def webui_tag_by_clip(img):
|
|
url = "http://127.0.0.1:57860/sdapi/v1/interrogate"
|
|
request_dict = {
|
|
"image": encode_numpy_to_base64(img),
|
|
"model": "clip"
|
|
}
|
|
response = requests.post(url=url, json=request_dict)
|
|
ret_json = response.json()
|
|
return ret_json['caption']
|
|
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
hair_dir = '/mnt/database2/jiangqian/0808/exp2-data-zrn-0808/1816523294655647746'
|
|
|
|
data2process_list = []
|
|
|
|
# 遍历查找hair_dir下的所有npy文件
|
|
for root, dirs, files in os.walk(hair_dir):
|
|
for file in files:
|
|
if file.endswith('.npy'):
|
|
# 关键点路径
|
|
pt1k_path = os.path.join(root, file)
|
|
origin_img_path = pt1k_path[:-4] + '.png'
|
|
origin_matting_path = pt1k_path[:-4] + '_origin_matting.png'
|
|
new_matting_path = pt1k_path[:-4] + '_new_matting.png'
|
|
result_img_path = pt1k_path[:-4] + '_res.png'
|
|
# 判断上面的文件是否存在
|
|
if (not os.path.exists(origin_img_path) or not os.path.exists(origin_matting_path)
|
|
or not os.path.exists(new_matting_path) or not os.path.exists(result_img_path)):
|
|
continue
|
|
|
|
ref_hair_path = pt1k_path[:-4] + '_orig_hair.png'
|
|
# if not os.path.exists(ref_hair_path):
|
|
# continue
|
|
lora_model_path = pt1k_path[:-4] + '_hairstyle_lora.safetensors'
|
|
# if not os.path.exists(lora_model_path):
|
|
# continue
|
|
data2process_list.append([pt1k_path, origin_img_path, origin_matting_path,
|
|
new_matting_path, result_img_path, ref_hair_path, lora_model_path])
|
|
|
|
crop_size = 768
|
|
|
|
webui_lora_dir = '/home/student/Documents/workspace_cxt_tianjing_hair/miaoya/webui_home/stable-diffusion-webui/models/Lora'
|
|
for pt1k_path, origin_img_path, origin_matting_path, new_matting_path, result_img_path, ref_hair_path, lora_model_path in tqdm.tqdm(data2process_list):
|
|
# if '508417f3-2c71-45cf-a75b-969b27ec7d8f' not in pt1k_path: continue
|
|
# 读取关键点
|
|
pt1k = np.load(pt1k_path)
|
|
# 读取原图
|
|
origin_img = cv2.imread(origin_img_path)
|
|
tmp_scale = 1920 / max(origin_img.shape[0], origin_img.shape[1])
|
|
if tmp_scale < 1.0:
|
|
origin_img = cv2.resize(origin_img, (0, 0), fx=tmp_scale, fy=tmp_scale, interpolation=cv2.INTER_LANCZOS4)
|
|
|
|
print("origin_img shape:", origin_img.shape)
|
|
# cv2.imshow("origin_img", origin_img)
|
|
|
|
# 读取原图抠图
|
|
origin_matting = cv2.imread(origin_matting_path, cv2.IMREAD_GRAYSCALE)
|
|
# 读取新图抠图
|
|
new_matting = cv2.imread(new_matting_path, cv2.IMREAD_GRAYSCALE)
|
|
# 读取结果图
|
|
result_img = cv2.imread(result_img_path)
|
|
print("result_img shape:", result_img.shape)
|
|
# cv2.imshow("result_img", result_img)
|
|
# # 读取参考头发
|
|
# ref_hair = cv2.imread(ref_hair_path)
|
|
# if max(ref_hair.shape[:2]) < 300: continue
|
|
|
|
|
|
# 如何图像不清晰，进行超分辨率处理
|
|
# if max(origin_img.shape[:2]) < 1500:
|
|
# scale_ratio = 2000 / max(origin_img.shape[:2])
|
|
# result_img = webui_super_res_img(result_img, scale_ratio)
|
|
# origin_img = cv2.resize(origin_img, (result_img.shape[1], result_img.shape[0]),
|
|
# interpolation=cv2.INTER_LANCZOS4)
|
|
# origin_matting = cv2.resize(origin_matting, (result_img.shape[1], result_img.shape[0]))
|
|
# new_matting = cv2.resize(new_matting, (result_img.shape[1], result_img.shape[0]))
|
|
# pt1k = pt1k * scale_ratio
|
|
|
|
# 获取头发处理的局部区域图像
|
|
# M = landmark_processor.get_transform_mat_hair_ratio_v1(pt1k, crop_size, ratio=0.30, h_offset=0.32)
|
|
|
|
scale = 768 / max(origin_img.shape[0], origin_img.shape[1])
|
|
M = cv2.getRotationMatrix2D((0, 0), 0, scale)
|
|
|
|
dst_size = (int(origin_img.shape[1] * scale), int(origin_img.shape[0] * scale))
|
|
|
|
# 高质量的从原图中截取头发区域
|
|
crop_origin = landmark_processor.high_quality_warpAffine(origin_img, M, dst_size)
|
|
cv2.imwrite("./crop_origin.png", crop_origin)
|
|
|
|
crop_result = landmark_processor.high_quality_warpAffine(result_img, M, dst_size)
|
|
cv2.imwrite("./crop_result.png", crop_result)
|
|
# tmp = cv2.warpAffine(origin_img, M, dst_size, flags=cv2.INTER_AREA)
|
|
|
|
# 构造重绘的mask
|
|
matting_merge = np.concatenate([origin_matting[:,:, np.newaxis], new_matting[:,:, np.newaxis]], axis=2)
|
|
matting_merge = np.max(matting_merge, axis=2)
|
|
# matting_merge = new_matting
|
|
crop_matting = cv2.warpAffine(matting_merge, M, dst_size)
|
|
mask = (crop_matting > 10).astype(np.float32)
|
|
mask_dilate = cv2.dilate(mask, np.ones((3, 11), np.uint8))
|
|
final_img = crop_result
|
|
|
|
mask_dilate = np.clip(mask_dilate * 255, 0, 255).astype(np.uint8)
|
|
|
|
# file_name = os.path.basename(pt1k_path)[:-4]
|
|
# save_dir = '/home/chinatszrn/Downloads/abc/ref_hair/dst_res/style3_tmp'
|
|
# cv2.imwrite(os.path.join(save_dir, file_name + '.png'), final_img)
|
|
# cv2.imwrite(os.path.join(save_dir, file_name + '_mask.png'), mask_dilate)
|
|
# # cv2.imwrite(os.path.join(save_dir, file_name + '_ref_hair.png'), ref_hair)
|
|
# continue
|
|
|
|
# # 拷贝lora模型
|
|
# os.system('cp {} {}'.format(lora_model_path, webui_lora_dir))
|
|
# # 构建prompt提示词
|
|
# lora_model_name = os.path.basename(pt1k_path)[:-4]
|
|
|
|
# 对final_img进行打标
|
|
# tag_result = webui_tag_by_clip(final_img)
|
|
tag_result = ''
|
|
|
|
# 开始重绘
|
|
prompt = f'<lora:1816523294655647746_hairstyle_lora:1> titor hairstyle, easyphoto, ' + tag_result
|
|
# 对发型进行重绘
|
|
|
|
# cv2.imshow("final_img_0", final_img)
|
|
# cv2.imshow("mask_dilate", mask_dilate)
|
|
# cv2.waitKey(100)
|
|
|
|
sd_result = webui_img2img(final_img, mask_dilate, prompt)
|
|
|
|
final_img = origin_img.copy()
|
|
# 将重绘结果恢复到原图
|
|
M_inv = cv2.invertAffineTransform(M)
|
|
cv2.warpAffine(sd_result, M_inv, (final_img.shape[1], final_img.shape[0]), dst=final_img,
|
|
borderMode=cv2.BORDER_TRANSPARENT, flags=cv2.INTER_LANCZOS4)
|
|
|
|
# cv2.imshow("final_img", final_img)
|
|
# cv2.waitKey(0)
|
|
cv2.imwrite(result_img_path[:-4] + '_sd.png', final_img)
|
|
|
|
# ref_hair_scale = final_img.shape[0] / ref_hair.shape[0]
|
|
# ref_hair = cv2.resize(ref_hair, (0, 0), fx=ref_hair_scale, fy=ref_hair_scale, interpolation=cv2.INTER_LANCZOS4)
|
|
# img2show = np.concatenate([origin_img, ref_hair, final_img], axis=1)
|
|
#
|
|
# # 显示结果
|
|
# save_dir = '/home/chinatszrn/Downloads/exp'
|
|
# cv2.imwrite(os.path.join(save_dir, lora_model_name + '.png'), img2show)
|
|
# # cv2.imshow("origin_img", cv2.resize(img2show, (0, 0), fx=0.3, fy=0.3, interpolation=cv2.INTER_AREA))
|
|
# cv2.imshow("sd_result", cv2.resize(sd_result, (0, 0), fx=0.3, fy=0.3, interpolation=cv2.INTER_AREA))
|
|
# cv2.waitKey(1000)
|
|
|
|
|
|
|
|
|
|
|