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' 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)