import pickle import time import cv2 import numpy as np import torch import json import os import configparser from process_modules import PersonProcessor_yolov5,KeypointsProcessor,Human_Keypoints,pt_conv_25_to_17 from process_modules import Get_Landmark, Process_Data, Generator_Hair, Generator_Fusion_Res, Change_Hair_Color, BodySeg from prepare_ref_hairstyle_data import GenderClassifyProcessor from face_enhance.face_enhancement import FaceEnhancement from utils import landmark_processor from common.logger import LogFactory from faceseg.face_seg import FaceSeg from hair_init import HairInit config = configparser.ConfigParser() # 创建对象 config.read("config/configure.ini", encoding="utf-8") # 读取配置文件，如果配置文件不存在则创建 class HairStyle_Model_Infer(object): def __init__(self, gpu=True, use_enhance=False): hair_init = HairInit() self.use_enhance = use_enhance self.gpu_index = hair_init.gpu_index self.get_landmark = hair_init.get_landmark self.hair_size = hair_init.hair_size self.process_data_infer = hair_init.process_data_infer self.generator_hair = hair_init.generator_hair self.hair_fusion = hair_init.hair_fusion self.logger_call = hair_init.logger_call self.face_seg = hair_init.face_seg self.face_enhance = hair_init.face_enhance self.change_haircolor = hair_init.change_haircolor self.gender_classify = hair_init.gender_classify self.person_processor = hair_init.person_processor self.keypoints_processor = hair_init.keypoints_processor self.human_keypoint = hair_init.human_keypoint # if gpu and torch.cuda.is_available(): # self.device = torch.device("cuda:%d" % device_id if device_id >= 0 else "cpu") # self.cuda = True # else: # self.device = torch.device("cpu") # self.cuda = False # # self.get_landmark = Get_Landmark(gpu_id=device_id) # self.hair_size = 768 # self.process_data = Process_Data(gpu, device_id) # self.generator_hair = Generator_Hair(gpu, device_id) # self.hair_fusion = Generator_Fusion_Res(gpu, device_id) # self.logger_call = LogFactory.getLogger("call") # self.face_seg = FaceSeg(device_id) # self.use_enhance = use_enhance # if self.use_enhance: # self.face_enhance = FaceEnhancement(512, device_id) # self.table_enlight = cv2.imread('./data/convert_enlight.png') # self.change_haircolor = Change_Hair_Color(gpu, device_id) # self.gender_classify = GenderClassifyProcessor(gpu_id=device_id) # # # body process # self.person_processor = PersonProcessor_yolov5(gpu_id=device_id) # self.keypoints_processor = KeypointsProcessor(gpu_id=device_id) # self.human_keypoint = Human_Keypoints(gpu=True, device_id=device_id) # print("Load model finish ... ") def infer_haircolor_v4(self, user_rgb_8uc3_orisize, target_hair_color, haircolor_dir): # haircolor_dir = '/home/data/hair/data/ref_color/3628746832766' face_base, hair_matting, status = self.infer_haircolor_new(user_rgb_8uc3_orisize, haircolor_dir,target_hair_color, return_matting=True) user_rgb_8uc3_orisize = face_base # 构建一个色板 r, g, b = target_hair_color tar_color = np.zeros_like(user_rgb_8uc3_orisize) tar_color[:, :, 0] = b tar_color[:, :, 1] = g tar_color[:, :, 2] = r # 将色板和原图转换到LAB空间 origin_img_LAB = cv2.cvtColor(user_rgb_8uc3_orisize, cv2.COLOR_BGR2LAB) target_color_LAB = cv2.cvtColor(tar_color, cv2.COLOR_BGR2LAB) # img_res_LAB取色板的色度通道和原图的亮度通道 img_res_LAB = np.concatenate([origin_img_LAB[:, :, 0:1], target_color_LAB[:, :, 1:3]], axis=2) ret_color_img = cv2.cvtColor(img_res_LAB, cv2.COLOR_LAB2BGR) # 将头发的matting转换为三通道 user_matting_mask_fc32_orisize = np.repeat(hair_matting[:, :, np.newaxis], 3, axis=2).astype(np.float32) / 255 ### change ycj img_res_change_hsv_change_bgr_fc32 = ret_color_img.astype(np.float32) / 255 ref_rgb_avg_std = ([b / 255.0, g / 255.0, r / 255.0], [0, 0, 0]) face_base_new_fc32_orisize_new, src_avg_std = self.reinhard_rgb(img_res_change_hsv_change_bgr_fc32, user_matting_mask_fc32_orisize, ref_rgb_avg_std, ratio=1.0) face_base_new_fc32_orisize_new = face_base_new_fc32_orisize_new * user_matting_mask_fc32_orisize + \ (user_rgb_8uc3_orisize.astype(np.float32) / 255) * ( 1 - user_matting_mask_fc32_orisize) face_base_new_new = (face_base_new_fc32_orisize_new * 255).astype(np.uint8) return face_base_new_new, hair_matting, 0 def infer_hairstyle_random_ref(self, origin_img, ref_img): ratio = 1 if ratio == 0: gender = "boy" else: gender = "girl" user_rgb_8uc3_orisize = origin_img landmarks_origin_img_1k = self.get_landmark.forward(origin_img) if landmarks_origin_img_1k is None: return None, 10001 user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize, \ user_baldseg_8uc3_768, user_bald_8uc3_768, user_landmark_f1k2_768, user_hairstyle_M = self.process_data.get_prepare_user_768_data(user_rgb_8uc3_orisize, landmarks_origin_img_1k, ratio=ratio) # cv2.imshow("user_bald_res_8uc3_orisize", user_bald_res_8uc3_orisize) # cv2.waitKey() # show_concat = np.concatenate((user_rgb_8uc3_orisize, user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize), axis=1) # ratio = 1536. / max(show_concat.shape[:2]) # show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) # cv2.imshow("show_concat", show_concat) # cv2.waitKey() ref_landmark_1k2_f_orisize = self.get_landmark.forward(ref_img) ref_rgb_8uc3_768, ref_matting_fg_8uc3_768, ref_matting_8uc3_768, ref_baldseg_8uc3_768, ref_landmark_f1k2_768 = \ self.process_data.get_prepare_ref_768_data(ref_img, ref_landmark_1k2_f_orisize, ratio=2) another_pose_hair_image = self.process_data.Generator_reftensor(ref_rgb_8uc3_768, ref_matting_8uc3_768, ref_baldseg_8uc3_768, ref_landmark_f1k2_768) if another_pose_hair_image is None: return None, 10002 # 换发型 hair_gene_8uc3_768 = self.generator_hair.Generator_Hair_inference_use_pref(another_pose_hair_image, user_baldseg_8uc3_768, user_bald_8uc3_768, user_landmark_f1k2_768, gender) hair_gene_fusion_8uc3_orisize, hair_gene_matte_8uc3_orisize = self.process_data.get_fusion_res_hairpaste( user_bald_res_8uc3_orisize, hair_gene_8uc3_768, user_landmark_f1k2_768, user_hairstyle_M) user_res_8uc3_orisize = self.hair_fusion.inference(hair_gene_fusion_8uc3_orisize, hair_gene_matte_8uc3_orisize, landmarks_origin_img_1k, user_baldseg_8uc3_orisize, ratio) # if self.use_enhance: user_res_8uc3_orisize = self.face_enhance.process(user_res_8uc3_orisize, landmarks_origin_img_1k) # for pt in landmarks_origin_img_1k.astype(np.int32): # cv2.circle(origin_img, (pt[0], pt[1]), 1, (0, 255, 0), -1) # show_concat_orisize = np.concatenate((origin_img, user_res_8uc3_orisize), axis=1) # ratio = 1024. / max(show_concat_orisize.shape[:2]) # show_concat_orisize = cv2.resize(show_concat_orisize, (0, 0), fx=ratio, fy=ratio) # ratio = 1024. / max(show_concat_768.shape[:2]) # show_concat_768 = cv2.resize(show_concat_768, (0, 0), fx=ratio, fy=ratio) # cv2.imshow("show_concat_orisize", show_concat_orisize) # cv2.imshow("show_concat_768", show_concat_768) # cv2.waitKey() return user_res_8uc3_orisize, 0 def draw_hull_mask(self, fc_landmark, w=256, h=256, is_gray=False): hull_mask = np.zeros((h, w), dtype=np.float32) if len(fc_landmark) == 1000: cv2.fillConvexPoly(hull_mask, cv2.convexHull(fc_landmark), (1)) else: raise Exception('landmark should be 1000') return hull_mask def get_1k(self, origin_img): landmarks_origin_img_1k = self.get_landmark.forward(origin_img) return landmarks_origin_img_1k def get_body_info(self, img_res): with torch.no_grad(): img_w, img_h = 768, 1024 ori_h, ori_w, _ = img_res.shape res = self.person_processor.forward(img_res) if len(res['boxes']) == 0: print('No person detected') return [0,0,ori_w, ori_h] filter_boxes = [] for ix, box_score in enumerate(res['scores']): box_ = res['boxes'][ix] box_w = box_[2] - box_[0] box_h = box_[3] - box_[1] max_box_len = max(box_w, box_h) if box_score > 0.5 and max_box_len > 150: filter_boxes.append([box_, box_w * box_h, box_h / img_res.shape[0]]) filter_boxes.sort(key=lambda x: x[1], reverse=True) filter_boxes = list(filter(lambda x: x[2] > 0.2, filter_boxes)) if len(filter_boxes) == 0: return [0, 0, ori_w, ori_h] filter_boxes = [item[0] for item in filter_boxes] # return filter_boxes, res # step2 person keypoints (version: hrnet 17pt) keypoints = self.keypoints_processor.forward(img_res, filter_boxes) # step3 human keypoints: hands, face, body keypoints(25pt) person_keypoints = keypoints[0] human_box = filter_boxes[0].reshape((-1, 2)) human_box[0][1] = max(0, human_box[0][1] - ori_h * 0.1) # human_box[0][0] = max(0, human_box[0][0]-ori_w*0.1) # human_box[1][0] = min(ori_w, human_box[1][0]+ori_w*0.1) knee_keypoints_left = person_keypoints[13] knee_keypoints_right = person_keypoints[14] if knee_keypoints_left[2] > 0.4 or knee_keypoints_right[2] > 0.4: src_body_kpnts_select = person_keypoints[np.where(person_keypoints[:, 2] > 0.2)][:12][:, :2] src_body_box = cv2.boundingRect(src_body_kpnts_select[np.newaxis, :, :]) kpnts_bbox_tlx, kpnts_bbox_tly = src_body_box[:2] kpnts_bbox_brx, kpnts_bbox_bry = kpnts_bbox_tlx + src_body_box[2] - 1, kpnts_bbox_tly + src_body_box[ 3] - 1 kpnts_bbox_tlx, kpnts_bbox_tly = min(kpnts_bbox_tlx, human_box[0][0]), min(kpnts_bbox_tly, human_box[0][1]) kpnts_bbox_brx, kpnts_bbox_bry = min(kpnts_bbox_brx, human_box[1][0]), min(kpnts_bbox_bry, human_box[1][1]) box_center_x = int(kpnts_bbox_tlx / 2 + kpnts_bbox_brx / 2) box_hight_now = kpnts_bbox_bry - kpnts_bbox_tly kpnts_bbox_tlx = max(int(box_center_x - box_hight_now * 3 / 8), 0) kpnts_bbox_brx = min(int(box_center_x + box_hight_now * 3 / 8), ori_w) # cv2.rectangle(img_res, (kpnts_bbox_tlx, kpnts_bbox_tly), (kpnts_bbox_brx, kpnts_bbox_bry), (0, 255, 0), # 2) waist_keypoints_mean = (person_keypoints[11] + person_keypoints[12]) / 2 crop_img = img_res[kpnts_bbox_tly:kpnts_bbox_bry, kpnts_bbox_tlx:kpnts_bbox_brx] return [kpnts_bbox_tlx, kpnts_bbox_tly,kpnts_bbox_brx, kpnts_bbox_bry] return [0, 0, ori_w, ori_h] def infer_hairstyle(self, origin_img, hairstyle_dir, return_pt1k=False,use_enhance=False): config_path = os.path.join(hairstyle_dir, "config.json") if not os.path.exists(config_path): return None, 10001 config_info = json.load(open(config_path, "rb")) gender = config_info["gender"] ratio = int(config_info["ratio"]) if ratio == 2: ratio = 3 print("gender: ", gender, " ratio: ", ratio) user_rgb_8uc3_orisize = origin_img landmarks_origin_img_1k = self.get_landmark.forward_infer(origin_img) # 根据人脸关键点画出人脸区域的mask hull_mask = self.draw_hull_mask(landmarks_origin_img_1k.astype(np.int32), w=origin_img.shape[1], h=origin_img.shape[0]) user_mask_save_path = os.path.join(hairstyle_dir, "hull_mask.png") cv2.imwrite(user_mask_save_path, hull_mask * 255) # cv2.imshow('hull_mask', hull_mask) # cv2.waitKey(0) if landmarks_origin_img_1k is None: return None, 10001 user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize, \ user_baldseg_8uc3_768, user_bald_8uc3_768, user_landmark_f1k2_768, user_hairstyle_M, user_matting_8uc3_bald_orisize = self.process_data_infer.get_prepare_user_768_data(user_rgb_8uc3_orisize, landmarks_origin_img_1k, ratio=ratio) # show_concat = np.concatenate((user_rgb_8uc3_orisize, user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize), axis=1) # ratio = 1536. / max(show_concat.shape[:2]) # show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) # cv2.imshow("show_concat", show_concat) # cv2.waitKey() # cv2.imshow("user_matting_8uc3_bald_orisize", user_matting_8uc3_bald_orisize) # cv2.waitKey(0) user_orig_mask_path = os.path.join(hairstyle_dir, "user_orig_mask.png") cv2.imwrite(user_orig_mask_path, user_matting_8uc3_bald_orisize) another_pose_hair_img_path = os.path.join(hairstyle_dir, "input_another_pose_hair_image.npy") if not os.path.exists(another_pose_hair_img_path): return None, 10002 another_pose_hair_image = np.load(another_pose_hair_img_path) # 换发型 hair_gene_8uc3_768 = self.generator_hair.Generator_Hair_inference_use_pref(another_pose_hair_image, user_baldseg_8uc3_768, user_bald_8uc3_768, user_landmark_f1k2_768, gender) hair_gene_fusion_8uc3_orisize, hair_gene_matte_8uc3_orisize = self.process_data_infer.get_fusion_res_hairpaste( user_bald_res_8uc3_orisize, hair_gene_8uc3_768, user_landmark_f1k2_768, user_hairstyle_M) # cv2.imshow("hair_gene_fusion_8uc3_orisize:", hair_gene_fusion_8uc3_orisize) # cv2.imshow("hair_gene_matte_8uc3_orisize:", hair_gene_matte_8uc3_orisize) # cv2.waitKey(0) gen_hair_mask_path = os.path.join(hairstyle_dir, "hair_mask.png") cv2.imwrite(gen_hair_mask_path, hair_gene_matte_8uc3_orisize * 255) gen_hair_mask_path_2 = os.path.join(hairstyle_dir, "hair_mask_2.png") cv2.imwrite(gen_hair_mask_path_2, hair_gene_matte_8uc3_orisize) if use_enhance: user_res_8uc3_orisize = self.hair_fusion.inference(hair_gene_fusion_8uc3_orisize, hair_gene_matte_8uc3_orisize, landmarks_origin_img_1k, user_baldseg_8uc3_orisize, ratio) hair_gene_fusion_8uc3_orisize_LAB = cv2.cvtColor(hair_gene_fusion_8uc3_orisize, cv2.COLOR_BGR2LAB) user_res_8uc3_orisize_LAB = cv2.cvtColor(user_res_8uc3_orisize, cv2.COLOR_BGR2LAB) mix_ratio = 0.2 user_res_8uc3_orisize_LAB[:, :, 0] = hair_gene_fusion_8uc3_orisize_LAB[:, :, 0] * mix_ratio + user_res_8uc3_orisize_LAB[:, :, 0] * (1 - mix_ratio) user_res_8uc3_orisize = cv2.cvtColor(user_res_8uc3_orisize_LAB, cv2.COLOR_LAB2BGR) # show_concat = np.concatenate((hair_gene_fusion_8uc3_orisize, hair_gene_matte_8uc3_orisize, # user_baldseg_8uc3_orisize, user_res_8uc3_orisize), axis=1) # ratio = 1536. / max(show_concat.shape[:2]) # show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) # cv2.imshow("mid_concat", show_concat) # cv2.waitKey() else: user_res_8uc3_orisize = hair_gene_fusion_8uc3_orisize # face_bbox = cv2.boundingRect(landmarks_origin_img_1k[np.newaxis, :, :]) # face_max_len = max(face_bbox[2], face_bbox[3]) # erode_kernel_size = int(face_max_len * 0.11) # if erode_kernel_size % 2 == 0: # erode_kernel_size += 1 # blur_kernel_size = int(face_max_len * 0.05) # if blur_kernel_size % 2 == 0: # blur_kernel_size += 1 # # print("blur_kernel_size: ", blur_kernels # hair_gene_matte_8uc3_hard = hair_gene_matte_8uc3_orisize.copy() # hair_gene_matte_8uc3_hard[hair_gene_matte_8uc3_orisize[:, :, 0] > 0] = 255 # hair_gene_matte_8uc3_orisize_erode = cv2.erode(hair_gene_matte_8uc3_hard, np.ones((erode_kernel_size, erode_kernel_size), np.uint8), iterations=1) # hair_gene_matte_fc32_orisize_erode = hair_gene_matte_8uc3_orisize_erode.astype(np.float32) / 255 # hair_gene_matte_fc32_orisize = hair_gene_matte_8uc3_orisize.astype(np.float32) / 255 # hair_gene_matte_fc32_orisize_circle = hair_gene_matte_fc32_orisize * (1 - hair_gene_matte_fc32_orisize_erode) # hair_gene_matte_fc32_orisize_circle = cv2.GaussianBlur(hair_gene_matte_fc32_orisize_circle, (blur_kernel_size, blur_kernel_size), 0, 0) # # hair_gene_fusion_8uc3_orisize_LAB = cv2.cvtColor(hair_gene_fusion_8uc3_orisize, cv2.COLOR_BGR2LAB) # user_res_8uc3_orisize_LAB = cv2.cvtColor(user_res_8uc3_orisize, cv2.COLOR_BGR2LAB) # mix_ratio = 0.5 # user_res_8uc3_orisize_LAB_new = user_res_8uc3_orisize_LAB.copy() # user_res_8uc3_orisize_LAB_new[:, :, 0] = hair_gene_fusion_8uc3_orisize_LAB[:, :, 0] * mix_ratio + user_res_8uc3_orisize_LAB[:, :, 0] * (1 - mix_ratio) # user_res_8uc3_orisize_LAB[:, :, 0] = user_res_8uc3_orisize_LAB_new[:, :, 0] * hair_gene_matte_fc32_orisize_circle[:, :, 0] + user_res_8uc3_orisize_LAB[:, :, 0] * (1 - hair_gene_matte_fc32_orisize_circle[:, :, 0]) # user_res_8uc3_orisize = cv2.cvtColor(user_res_8uc3_orisize_LAB, cv2.COLOR_LAB2BGR) # mid_show = np.concatenate((hair_gene_matte_fc32_orisize, hair_gene_matte_fc32_orisize_erode, hair_gene_matte_fc32_orisize_circle), axis=1) # ratio = 1536. / max(mid_show.shape[:2]) # mid_show = cv2.resize(mid_show, (0, 0), fx=ratio, fy=ratio) # cv2.imshow("mid_show", mid_show) # # cv2.imshow("user_res_8uc3_orisize", user_res_8uc3_orisize) # cv2.waitKey() if use_enhance: user_res_8uc3_orisize_enhance = self.face_enhance.process(user_res_8uc3_orisize, landmarks_origin_img_1k) user_res_8uc3_orisize = user_bald_res_8uc3_orisize.astype(np.float32) / 255. * ( 1. - hair_gene_matte_8uc3_orisize.astype( np.float32) / 255.) + user_res_8uc3_orisize_enhance.astype( np.float32) / 255. * hair_gene_matte_8uc3_orisize.astype(np.float32) / 255. user_res_8uc3_orisize = (user_res_8uc3_orisize * 255).astype(np.uint8) # for pt in landmarks_origin_img_1k.astype(np.int32): # cv2.circle(origin_img, (pt[0], pt[1]), 1, (0, 255, 0), -1) # show_concat_orisize = np.concatenate((origin_img, user_res_8uc3_orisize), axis=1) # ratio = 1024. / max(show_concat_orisize.shape[:2]) # show_concat_orisize = cv2.resize(show_concat_orisize, (0, 0), fx=ratio, fy=ratio) # ratio = 1024. / max(show_concat_768.shape[:2]) # show_concat_768 = cv2.resize(show_concat_768, (0, 0), fx=ratio, fy=ratio) # cv2.imshow("show_concat_orisize", show_concat_orisize) # cv2.imshow("show_concat_768", show_concat_768) # cv2.waitKey() if return_pt1k: ret_dict = dict(user_res_8uc3_orisize=user_res_8uc3_orisize, user_bald_res_8uc3_orisize=user_bald_res_8uc3_orisize, landmarks_origin_img_1k=landmarks_origin_img_1k) return ret_dict, 0 else: return user_res_8uc3_orisize, 0 def infer_hairstyle_fix_face(self, origin_img, hairstyle_dir): config_path = os.path.join(hairstyle_dir, "config.json") if not os.path.exists(config_path): return None, 10001 config_info = json.load(open(config_path, "rb")) gender = config_info["gender"] ratio = int(config_info["ratio"]) # print("gender: ", gender, " ratio: ", ratio) user_rgb_8uc3_orisize = origin_img landmarks_origin_img_1k = self.get_landmark.forward(origin_img) if landmarks_origin_img_1k is None: return None, 10001 user_hair_matting_fg_orisize, user_hair_matting_8uc1_orisize = self.process_data.generator_matte.matte_inference(user_rgb_8uc3_orisize, landmarks_origin_img_1k) user_hair_matting_fc32_orisize = user_hair_matting_8uc1_orisize[:, :, np.newaxis].astype(np.float32) / 255 user_face_mask_fc32_orisize = self.face_seg.inference(origin_img, landmarks_origin_img_1k) user_face_mask_fc32_orisize = np.repeat(user_face_mask_fc32_orisize, 3, axis=2) user_face_mask_fc32_orisize = np.clip(user_face_mask_fc32_orisize, 0, 1) user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize, \ user_baldseg_8uc3_768, user_bald_8uc3_768, user_landmark_f1k2_768, user_hairstyle_M,_ = self.process_data.get_prepare_user_768_data(user_rgb_8uc3_orisize, landmarks_origin_img_1k, ratio=ratio) # show_concat = np.concatenate((user_rgb_8uc3_orisize, user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize), axis=1) # ratio = 1536. / max(show_concat.shape[:2]) # show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) # cv2.imshow("show_concat", show_concat) # cv2.waitKey() another_pose_hair_img_path = os.path.join(hairstyle_dir, "input_another_pose_hair_image.npy") if not os.path.exists(another_pose_hair_img_path): return None, 10002 another_pose_hair_image = np.load(another_pose_hair_img_path) # 换发型 hair_gene_8uc3_768 = self.generator_hair.Generator_Hair_inference_use_pref(another_pose_hair_image, user_baldseg_8uc3_768, user_bald_8uc3_768, user_landmark_f1k2_768, gender) hair_gene_fusion_8uc3_orisize, hair_gene_matte_8uc3_orisize = self.process_data.get_fusion_res_hairpaste( user_bald_res_8uc3_orisize, hair_gene_8uc3_768, user_landmark_f1k2_768, user_hairstyle_M) user_res_8uc3_orisize = self.hair_fusion.inference(hair_gene_fusion_8uc3_orisize, hair_gene_matte_8uc3_orisize, landmarks_origin_img_1k, user_baldseg_8uc3_orisize, ratio) # for pt in landmarks_origin_img_1k.astype(np.int32): # cv2.circle(origin_img, (pt[0], pt[1]), 1, (0, 255, 0), -1) # show_concat_orisize = np.concatenate((origin_img, (user_face_mask_fc32_orisize*255).astype(np.uint8), hair_gene_fusion_8uc3_orisize, hair_gene_matte_8uc3_orisize, user_baldseg_8uc3_orisize, user_res_8uc3_orisize), axis=1) # ratio = 1536. / max(show_concat_orisize.shape[:2]) # show_concat_orisize = cv2.resize(show_concat_orisize, (0, 0), fx=ratio, fy=ratio) # cv2.imshow("show_concat_orisize", show_concat_orisize) # cv2.waitKey() if self.use_enhance: user_res_8uc3_orisize_enhance = self.face_enhance.process(user_res_8uc3_orisize, landmarks_origin_img_1k) # show_concat_orisize = np.concatenate((origin_img, user_res_8uc3_orisize, user_res_8uc3_orisize_enhance), axis=1) # ratio = 1536. / max(show_concat_orisize.shape[:2]) # show_concat_orisize = cv2.resize(show_concat_orisize, (0, 0), fx=ratio, fy=ratio) # cv2.imshow("show_concat_orisize", show_concat_orisize) # cv2.waitKey() user_res_8uc3_orisize = user_res_8uc3_orisize_enhance user_res_fc32_orisize = user_res_8uc3_orisize.astype(np.float32) / 255 hair_gene_matte_fc32_orisize = hair_gene_matte_8uc3_orisize.astype(np.float32) / 255 user_rgb_fc32_orisize = user_rgb_8uc3_orisize.astype(np.float32) / 255 user_hair_matting_fc32_orisize = cv2.dilate(user_hair_matting_fc32_orisize[:, :, 0], np.ones((35, 35), np.uint8))[:, :, np.newaxis] user_face_mask_fc32_orisize = user_face_mask_fc32_orisize * (1 - hair_gene_matte_fc32_orisize) * (1 - user_hair_matting_fc32_orisize) user_res_fc32_orisize = user_res_fc32_orisize * (1 - user_face_mask_fc32_orisize) + user_rgb_fc32_orisize * user_face_mask_fc32_orisize show_concat_orisize = np.concatenate((user_res_fc32_orisize, user_face_mask_fc32_orisize, user_rgb_fc32_orisize), axis=1) ratio = 1536. / max(show_concat_orisize.shape[:2]) show_concat_orisize = cv2.resize(show_concat_orisize, (0, 0), fx=ratio, fy=ratio) cv2.imshow("show_concat_orisize", show_concat_orisize) cv2.waitKey(10) return (user_res_fc32_orisize*255).astype(np.uint8), 0 # return user_res_8uc3_orisize, 0 def infer_hairstyle_finetune(self, origin_img, hairstyle_dir): config_path = os.path.join(hairstyle_dir, "config.json") if not os.path.exists(config_path): return None, 10001 config_info = json.load(open(config_path, "rb")) gender = config_info["gender"] ratio = int(config_info["ratio"]) # print("gender: ", gender, " ratio: ", ratio) user_rgb_8uc3_orisize = origin_img landmarks_origin_img_1k = self.get_landmark.forward(origin_img) if landmarks_origin_img_1k is None: return None, 10001 user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize, \ user_baldseg_8uc3_768, user_bald_8uc3_768, user_landmark_f1k2_768, user_hairstyle_M,_ = self.process_data.get_prepare_user_768_data(user_rgb_8uc3_orisize, landmarks_origin_img_1k, ratio=ratio) # show_concat = np.concatenate((user_rgb_8uc3_orisize, user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize), axis=1) # ratio = 1536. / max(show_concat.shape[:2]) # show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) # cv2.imshow("show_concat", show_concat) # cv2.waitKey() another_pose_hair_img_path = os.path.join(hairstyle_dir, "input_another_pose_hair_image.npy") if not os.path.exists(another_pose_hair_img_path): return None, 10002 another_pose_hair_image = np.load(another_pose_hair_img_path) # 换发型 hair_gene_8uc3_768 = self.generator_hair.Generator_Hair_inference_use_pref(another_pose_hair_image, user_baldseg_8uc3_768, user_bald_8uc3_768, user_landmark_f1k2_768, gender) hair_gene_fusion_8uc3_orisize, hair_gene_matte_8uc3_orisize = self.process_data.get_fusion_res_hairpaste( user_bald_res_8uc3_orisize, hair_gene_8uc3_768, user_landmark_f1k2_768, user_hairstyle_M) # cv2.imshow("hair_gene_fusion_8uc3_orisize", hair_gene_fusion_8uc3_orisize) # cv2.imshow("hair_gene_matte_8uc3_orisize", hair_gene_matte_8uc3_orisize) # cv2.imshow("user_bald_res_8uc3_orisize", user_bald_res_8uc3_orisize) # cv2.waitKey() user_res_8uc3_orisize = self.hair_fusion.inference(hair_gene_fusion_8uc3_orisize, hair_gene_matte_8uc3_orisize, landmarks_origin_img_1k, user_baldseg_8uc3_orisize, ratio) # cv2.imshow("hair_gene_fusion_8uc3_orisize", hair_gene_fusion_8uc3_orisize) # cv2.imshow("hair_gene_matte_8uc3_orisize", hair_gene_matte_8uc3_orisize) # cv2.imshow("user_res_8uc3_orisize", user_res_8uc3_orisize) # cv2.waitKey() if self.use_enhance: user_res_8uc3_orisize = self.face_enhance.process(user_res_8uc3_orisize, landmarks_origin_img_1k) # for pt in landmarks_origin_img_1k.astype(np.int32): # cv2.circle(origin_img, (pt[0], pt[1]), 1, (0, 255, 0), -1) # show_concat_orisize = np.concatenate((origin_img, user_res_8uc3_orisize), axis=1) # ratio = 1024. / max(show_concat_orisize.shape[:2]) # show_concat_orisize = cv2.resize(show_concat_orisize, (0, 0), fx=ratio, fy=ratio) # ratio = 1024. / max(show_concat_768.shape[:2]) # show_concat_768 = cv2.resize(show_concat_768, (0, 0), fx=ratio, fy=ratio) # cv2.imshow("show_concat_orisize", show_concat_orisize) # cv2.imshow("show_concat_768", show_concat_768) # cv2.waitKey() return user_res_8uc3_orisize, 0 def change_image_hair_color(self, image, color=[20, 20, 200]): image = self.lookup(self.table_enlight, image) b, g, r = color # [10, 50, 250] # [10, 250, 10] tar_color = np.zeros_like(image) tar_color[:, :, 0] = b tar_color[:, :, 1] = g tar_color[:, :, 2] = r image_hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV) tar_hsv = cv2.cvtColor(tar_color, cv2.COLOR_BGR2HSV) image_hsv[:, :, 0:2] = tar_hsv[:, :, 0:2] changed = cv2.cvtColor(image_hsv.astype(np.uint8), cv2.COLOR_HSV2BGR) # changed_sharpen = sharpen(changed) return changed def lookup(self, luv, img): table = luv b, g, r = img[:, :, 0] / 255.0, img[:, :, 1] / 255.0, img[:, :, 2] / 255.0 b = b * 63 x1 = np.floor(np.floor(b) / 8.) y1 = np.floor(b) - x1 * 8.0 idx_x1 = x1 * 0.125 + 0.5 / 512.0 + (0.125 - 1.0 / 512.0) * g idx_y1 = y1 * 0.125 + 0.5 / 512.0 + (0.125 - 1.0 / 512.0) * r x2 = np.floor(np.ceil(b) / 8.) y2 = np.ceil(b) - x2 * 8.0 idx_x2 = x2 * 0.125 + 0.5 / 512.0 + (0.125 - 1.0 / 512.0) * g idx_y2 = y2 * 0.125 + 0.5 / 512.0 + (0.125 - 1.0 / 512.0) * r fract = np.mod(b * 63, 1.0)[:, :, np.newaxis] idx_x1 = (idx_x1 * 512).astype(np.int32).reshape(-1) idx_y1 = (idx_y1 * 512).astype(np.int32).reshape(-1) idx_x2 = (idx_x2 * 512).astype(np.int32).reshape(-1) idx_y2 = (idx_y2 * 512).astype(np.int32).reshape(-1) color1 = table[idx_x1, idx_y1].reshape(img.shape[0], img.shape[1], 3) color2 = table[idx_x2, idx_y2].reshape(img.shape[0], img.shape[1], 3) final_img = np.clip(color1 * (1 - fract) + color2 * fract, 0, 255) return final_img.astype(np.uint8) def infer_hairtiaoran(self, origin_img, color, input_mask=None): landmarks_origin_img_1k = self.get_landmark.forward(origin_img) if landmarks_origin_img_1k is None: return None, 10001 hair_gene_matte_fg_8uc3_orisize, hair_gene_matte_8uc1_orisize = self.process_data.get_matte_img(origin_img, landmarks_origin_img_1k) hair_matting_mask_8uc3_orisize = np.repeat(hair_gene_matte_8uc1_orisize[:, :, np.newaxis], 3, axis=2) if input_mask is None: input_mask = np.zeros_like(origin_img) cv2.circle(input_mask, (int(landmarks_origin_img_1k[154, 0]), int(landmarks_origin_img_1k[154, 1])), int(origin_img.shape[0]*0.1), (255, 255, 255), -1) origin_img_haircolor = self.change_image_hair_color(origin_img, color) input_mask_blur = cv2.GaussianBlur(input_mask, (55, 55), 0, 0) input_mask_blur_fc32 = input_mask_blur.astype(np.float32) / 255 hair_matting_mask_fc32_orisize = hair_matting_mask_8uc3_orisize.astype(np.float32) / 255 fusion_mask_fc32 = input_mask_blur_fc32 *hair_matting_mask_fc32_orisize origin_img_haircolor_fc32 = origin_img_haircolor.astype(np.float32) / 255 origin_img_fc32 = origin_img.astype(np.float32) / 255 img_res_fc32 = origin_img_haircolor_fc32 * fusion_mask_fc32 + origin_img_fc32 * (1 - fusion_mask_fc32) img_res = (img_res_fc32 * 255).astype(np.uint8) # for pt in landmarks_origin_img_1k.astype(np.int32): # cv2.circle(origin_img, (pt[0], pt[1]), 2, (0, 255, 0), -1) # show_concat = np.concatenate((origin_img, input_mask, hair_matting_mask_8uc3_orisize, img_res), axis=1) # ratio = 1024. / max(show_concat.shape[:2]) # show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) # cv2.imshow("show_concat", show_concat) # cv2.waitKey() return img_res, 0 def calculate_hair_ratio_after_align(self, hair_mask, origin_landmark1k, img_size=768): image_to_face_mat = landmark_processor.get_transform_mat_hair_ratio_v1(origin_landmark1k, 768, ratio=0.35, h_offset=0.32) hair_mask_align = cv2.warpAffine(hair_mask, image_to_face_mat, (img_size, img_size)) hair_rect = cv2.boundingRect(hair_mask_align[:, :, :1]) hair_mask_ratio = hair_rect[2] * hair_rect[3] / (img_size * img_size) return hair_mask_ratio def Generator_reftensor(self, ref_rgb_8uc3_768, ref_matting_8uc3_768, ref_baldseg_8uc3_768, ref_landmark_f1k2_768): """ input： 图像尺寸基于： 人脸 512， 图像均为3通道 ref_rgb_8uc3_512: 参考图， size 512， uint8 （0-255） ref_matting_8uc3_512：参考图 matting alpha 值， uint8 （0-255） ref_baldseg_8uc3_512: 参考图 光头分割， uint8 （0-255） ref_landmark_f1k2_512： 参考图 关键点 1k*2 float32 output： input_another_pose_hair_image: 参考图 条件图， float32 （0-255） """ # 8 ********************** another pose hair_image ********************** another_pose_image = ref_rgb_8uc3_768.copy() another_nohair_pose_mask = ref_baldseg_8uc3_768.copy() # cv2.imshow("another_nohair_pose_mask", another_nohair_pose_mask) # cv2.waitKey() another_nohair_pose_mask[(np.abs(another_nohair_pose_mask - [0, 0, 255]) < 50).all(axis=2)] = [0, 255, 255] another_nohair_pose_mask[(another_nohair_pose_mask == [0, 0, 0]).all(axis=2)] = [255, 255, 255] another_pose_pts137 = landmark_processor.pts_1k_to_137(ref_landmark_f1k2_768).astype(np.int32) cv2.fillPoly(another_nohair_pose_mask, np.concatenate((another_pose_pts137[47:35:-1], another_pose_pts137[56:64], [another_pose_pts137[48], another_pose_pts137[22]]))[ np.newaxis, :, :], (255, 255, 0)) cv2.fillPoly(another_nohair_pose_mask, np.concatenate((another_pose_pts137[22:37], another_pose_pts137[56:47:-1]))[np.newaxis, :, :], (255, 0, 128)) cv2.fillPoly(another_nohair_pose_mask, another_pose_pts137[88:104][np.newaxis, :, :], (255, 0, 255)) # eye cv2.fillPoly(another_nohair_pose_mask, another_pose_pts137[105:121][np.newaxis, :, :], (255, 0, 255)) # eye # Label nose cv2.fillPoly(another_nohair_pose_mask, another_pose_pts137[64:79][np.newaxis, :, :], (255, 255, 255)) # Label eyebrow cv2.fillPoly(another_nohair_pose_mask, another_pose_pts137[121:129][np.newaxis, :, :], (255, 128, 0)) cv2.fillPoly(another_nohair_pose_mask, another_pose_pts137[129:137][np.newaxis, :, :], (255, 128, 0)) another_pose_hair_alpha = ref_matting_8uc3_768.copy() / 255. another_pose_hair_image = (another_pose_image * another_pose_hair_alpha + another_nohair_pose_mask * ( 1 - another_pose_hair_alpha)).astype(np.uint8) input_another_pose_hair_image = another_pose_hair_image.astype(np.float32) / 255 return input_another_pose_hair_image def infer_hairstyle_diy(self, user_rgb_8uc3_orisize, ref_rgb_8uc3_orisize): ref_landmark_1k2_f_orisize = self.get_landmark.forward(ref_rgb_8uc3_orisize) ref_landmark_137kpts_f_orisize = landmark_processor.pts_1k_to_137(ref_landmark_1k2_f_orisize) gender_res = self.gender_classify.forward(ref_rgb_8uc3_orisize, ref_landmark_137kpts_f_orisize) hair_ratio = self.calculate_hair_ratio_after_align(ref_rgb_8uc3_orisize, ref_landmark_1k2_f_orisize) if hair_ratio > 0.3: ratio = 2 else: if gender_res: ratio = 1 else: ratio = 0 if gender_res: gender = "girl" else: gender = "boy" ref_rgb_8uc3_768, ref_matting_fg_8uc3_768, ref_matting_8uc3_768, ref_baldseg_8uc3_768, ref_landmark_f1k2_768 = \ self.process_data.get_prepare_ref_768_data(ref_rgb_8uc3_orisize, ref_landmark_1k2_f_orisize, ratio) another_pose_hair_image = self.Generator_reftensor(ref_rgb_8uc3_768, ref_matting_8uc3_768, ref_baldseg_8uc3_768, ref_landmark_f1k2_768) landmarks_origin_img_1k = self.get_landmark.forward(user_rgb_8uc3_orisize) if landmarks_origin_img_1k is None: return None, 10001 user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize, \ user_baldseg_8uc3_768, user_bald_8uc3_768, user_landmark_f1k2_768, user_hairstyle_M,_ = self.process_data.get_prepare_user_768_data(user_rgb_8uc3_orisize, landmarks_origin_img_1k, ratio=ratio) # show_concat = np.concatenate((user_rgb_8uc3_orisize, user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize), axis=1) # ratio = 1536. / max(show_concat.shape[:2]) # show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) # cv2.imshow("show_concat", show_concat) # cv2.waitKey() # 换发型 hair_gene_8uc3_768 = self.generator_hair.Generator_Hair_inference_use_pref(another_pose_hair_image, user_baldseg_8uc3_768, user_bald_8uc3_768, user_landmark_f1k2_768, gender) hair_gene_fusion_8uc3_orisize, hair_gene_matte_8uc3_orisize = self.process_data.get_fusion_res_hairpaste( user_bald_res_8uc3_orisize, hair_gene_8uc3_768, user_landmark_f1k2_768, user_hairstyle_M) user_res_8uc3_orisize = self.hair_fusion.inference(hair_gene_fusion_8uc3_orisize, hair_gene_matte_8uc3_orisize, landmarks_origin_img_1k, user_baldseg_8uc3_orisize, ratio) if self.use_enhance: user_res_8uc3_orisize = self.face_enhance.process(user_res_8uc3_orisize, landmarks_origin_img_1k) # for pt in landmarks_origin_img_1k.astype(np.int32): # cv2.circle(origin_img, (pt[0], pt[1]), 1, (0, 255, 0), -1) # show_concat_orisize = np.concatenate((origin_img, user_res_8uc3_orisize), axis=1) # ratio = 1024. / max(show_concat_orisize.shape[:2]) # show_concat_orisize = cv2.resize(show_concat_orisize, (0, 0), fx=ratio, fy=ratio) # ratio = 1024. / max(show_concat_768.shape[:2]) # show_concat_768 = cv2.resize(show_concat_768, (0, 0), fx=ratio, fy=ratio) # cv2.imshow("show_concat_orisize", show_concat_orisize) # cv2.imshow("show_concat_768", show_concat_768) # cv2.waitKey() return user_res_8uc3_orisize, 0 def preprocess_haircolor(self, origin_img, origin_mask, ref_img, ref_mask): ref_img_hsv = cv2.cvtColor(ref_img, cv2.COLOR_BGR2HSV) loc_index = ref_mask[:, :, 0].nonzero() color_val = ref_img_hsv[loc_index] ref_img_mean_hsv = np.mean(color_val, axis=0) origin_img_hsv = cv2.cvtColor(origin_img, cv2.COLOR_BGR2HSV) loc_index = origin_mask[:, :, 0].nonzero() color_val = origin_img_hsv[loc_index] origin_img_mean_hsv = np.mean(color_val, axis=0) origin_img_hsv[:, :, :1] = ref_img_mean_hsv[:1] # origin_img_hsv[:, :, 1:3] = ref_img_mean_hsv[1:3] res_img = cv2.cvtColor(origin_img_hsv.astype(np.uint8), cv2.COLOR_HSV2BGR) origin_mask_fc32 = origin_mask.astype(np.float32) / 255 res_img = (res_img.astype(np.float32) * origin_mask_fc32 + origin_img.astype(np.float32) * (1 - origin_mask_fc32)).astype(np.uint8) mid_show = np.concatenate((origin_img, origin_mask, res_img), axis=1) ratio = 1536. / max(mid_show.shape[:2]) mid_show = cv2.resize(mid_show, (0, 0), fx=ratio, fy=ratio) cv2.imshow("mid_show", mid_show) cv2.waitKey(0) return res_img def rein_hard_raw(self, user_res_8uc3_orisize, user_matting_8uc3_bald_orisize, ref_rgb_8uc3_change_color_768, ref_matting_8uc3_change_color_768): ref_rgb_8uc3_change_color_768_hsv = cv2.cvtColor(ref_rgb_8uc3_change_color_768, cv2.COLOR_BGR2HSV) loc_index = ref_matting_8uc3_change_color_768[:, :, 0].nonzero() color_val = ref_rgb_8uc3_change_color_768_hsv[loc_index] mean_ref_hair_hsv = np.mean(color_val, axis=0) print("mean_ref_hair_hsv: ", mean_ref_hair_hsv) user_res_8uc3_orisize_hsv = cv2.cvtColor(user_res_8uc3_orisize, cv2.COLOR_BGR2HSV) loc_index = user_matting_8uc3_bald_orisize[:, :, 0].nonzero() color_val = user_res_8uc3_orisize_hsv[loc_index] mean_color_hair_hsv = np.mean(color_val, axis=0) print("mean_color_hair_hsv: ", mean_color_hair_hsv) user_res_8uc3_orisize_hsv[:, :, 1:3] = mean_ref_hair_hsv[1:3] img_res = cv2.cvtColor(user_res_8uc3_orisize_hsv, cv2.COLOR_HSV2BGR) user_res_8uc3_orisize_LAB = cv2.cvtColor(user_res_8uc3_orisize, cv2.COLOR_BGR2LAB) img_res_LAB = cv2.cvtColor(img_res, cv2.COLOR_BGR2LAB) img_res_LAB[:, :, 0] = user_res_8uc3_orisize_LAB[:, :, 0] img_res = cv2.cvtColor(img_res_LAB, cv2.COLOR_LAB2BGR) return img_res def rein_hard(self, user_res_8uc3_orisize, user_matting_8uc3_bald_orisize, ref_rgb_8uc3_change_color_768, ref_matting_8uc3_change_color_768): ref_rgb_8uc3_change_color_768_hsv = cv2.cvtColor(ref_rgb_8uc3_change_color_768, cv2.COLOR_BGR2HSV) loc_index = ref_matting_8uc3_change_color_768[:, :, 0].nonzero() color_val = ref_rgb_8uc3_change_color_768_hsv[loc_index] mean_ref_hair_hsv = np.mean(color_val, axis=0) print("mean_ref_hair_hsv: ", mean_ref_hair_hsv) user_res_8uc3_orisize_hsv = cv2.cvtColor(user_res_8uc3_orisize, cv2.COLOR_BGR2HSV) loc_index = user_matting_8uc3_bald_orisize[:, :, 0].nonzero() color_val = user_res_8uc3_orisize_hsv[loc_index] mean_color_hair_hsv = np.mean(color_val, axis=0) print("mean_color_hair_hsv: ", mean_color_hair_hsv) # user_res_8uc3_orisize_hsv[:, :, 0:1] = user_res_8uc3_orisize_hsv[:, :, 0:1] + mean_ref_hair_hsv[0] - mean_color_hair_hsv[0] user_res_8uc3_orisize_hsv[:, :, 0:1] = mean_ref_hair_hsv[0] # user_res_8uc3_orisize_hsv[:, :, 1:2] = user_res_8uc3_orisize_hsv[:, :, 1:2] + (mean_ref_hair_hsv[1] - mean_color_hair_hsv[1]) * 0.1 # user_res_8uc3_orisize_hsv[:, :, 0:1] = np.clip(user_res_8uc3_orisize_hsv[:, :, 0:1], 0, 180) # user_res_8uc3_orisize_hsv[:, :, 1:2] = user_res_8uc3_orisize_hsv[:, :, 1:2] * mean_ref_hair_hsv[1] / mean_color_hair_hsv[1] ratio_s = max(1.05, mean_ref_hair_hsv[1] * 0.85 / mean_color_hair_hsv[1]) print("ratio_s: ", ratio_s) user_res_8uc3_orisize_hsv[:, :, 1:2] = np.clip(user_res_8uc3_orisize_hsv[:, :, 1:2] * ratio_s, 0, 255) # user_res_8uc3_orisize_hsv[:, :, 2:3] = user_res_8uc3_orisize_hsv[:, :, 2:3] + (mean_ref_hair_hsv[2] - mean_color_hair_hsv[2]) * 0.4 # user_res_8uc3_orisize_hsv[:, :, 2:3] = np.clip(user_res_8uc3_orisize_hsv[:, :, 2:3], 0, 255) # user_res_8uc3_orisize_hsv[:, :, 2:3] = mean_ref_hair_hsv[2] ratio_v = 0.5 user_res_8uc3_orisize_hsv[:, :, 2:3] = user_res_8uc3_orisize_hsv[:, :, 2:3] * ratio_v + mean_ref_hair_hsv[2] * (1 - ratio_v) user_res_8uc3_orisize_hsv[:, :, 2:3] = np.clip(user_res_8uc3_orisize_hsv[:, :, 2:3], 0, 255) # ratio_v = mean_ref_hair_hsv[2] / mean_color_hair_hsv[2] # ratio_v = ratio_v * 0.8 + 0.2 # print("ratio_v: ", ratio_v) # user_res_8uc3_orisize_hsv[:, :, 2:3] = np.clip(user_res_8uc3_orisize_hsv[:, :, 2:3] * ratio_v, 0, 255) img_res = cv2.cvtColor(user_res_8uc3_orisize_hsv, cv2.COLOR_HSV2BGR) user_res_8uc3_orisize_LAB = cv2.cvtColor(user_res_8uc3_orisize, cv2.COLOR_BGR2LAB) img_res_LAB = cv2.cvtColor(img_res, cv2.COLOR_BGR2LAB) # img_res_LAB[:, :, 0] = user_res_8uc3_orisize_LAB[:, :, 0] mix_ratio = 0.5 img_res_LAB[:, :, 0] = user_res_8uc3_orisize_LAB[:, :, 0] * mix_ratio + img_res_LAB[:, :, 0] * (1 - mix_ratio) img_res = cv2.cvtColor(img_res_LAB, cv2.COLOR_LAB2BGR) # middle_show = np.concatenate((user_res_8uc3_orisize, img_res), axis=1) # ratio = 1536. / max(middle_show.shape[:2]) # middle_show = cv2.resize(middle_show, (0, 0), fx=ratio, fy=ratio) # cv2.imshow('middle_show', middle_show) # cv2.waitKey(0) return img_res def infer_haircolor(self, user_rgb_8uc3_orisize, haircolor_dir, return_matting=False): landmarks_origin_img_1k = self.get_landmark.forward(user_rgb_8uc3_orisize) if landmarks_origin_img_1k is None: return None, 10001 _, user_matting_8uc1_bald_orisize = self.process_data.generator_matte.matte_inference(user_rgb_8uc3_orisize, landmarks_origin_img_1k) user_matting_8uc3_bald_orisize = np.repeat(user_matting_8uc1_bald_orisize[:, :, np.newaxis], 3, axis=2) ref_rgb_8uc3_change_color_768 = cv2.imread(os.path.join(haircolor_dir, "ref_rgb_8uc3_color_768.png")) ref_matting_8uc3_change_color_768 = cv2.imread(os.path.join(haircolor_dir, "ref_matting_8uc3_color_768.png")) # user_rgb_8uc3_orisize = self.preprocess_haircolor(user_rgb_8uc3_orisize, user_matting_8uc3_bald_orisize, # ref_rgb_8uc3_change_color_768, ref_matting_8uc3_change_color_768) user_rgb_8uc3_change_color_768, user_matting_8uc3_change_color_768, user_hair_color_M = \ self.process_data.get_prepare_hair_color_user_data(user_rgb_8uc3_orisize, landmarks_origin_img_1k) # cv2.imshow("user_matting_8uc3_change_color_768", user_matting_8uc3_change_color_768) # cv2.waitKey() start_time = time.time() hair_gene_color_8uc3_768, user_matting_mask_8uc3_768 = self.change_haircolor.Change_Hair_inference( user_rgb_8uc3_change_color_768, user_matting_8uc3_change_color_768, ref_rgb_8uc3_change_color_768, ref_matting_8uc3_change_color_768) # print("Change_Hair_inference cost: ", time.time()-start_time) # cv2.imshow("hair_gene_color_8uc3_768: ", hair_gene_color_8uc3_768) # cv2.imshow("user_matting_mask_8uc3_768: ", user_matting_mask_8uc3_768) # cv2.waitKey() user_matting_8uc3_orisize = cv2.warpAffine(user_matting_mask_8uc3_768, cv2.invertAffineTransform(user_hair_color_M), (user_rgb_8uc3_orisize.shape[1], user_rgb_8uc3_orisize.shape[0]), flags=cv2.INTER_CUBIC) hair_gene_color_8uc3_orisize = cv2.warpAffine(hair_gene_color_8uc3_768, cv2.invertAffineTransform(user_hair_color_M), (user_rgb_8uc3_orisize.shape[1], user_rgb_8uc3_orisize.shape[0]), dst=user_rgb_8uc3_orisize.copy(), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_TRANSPARENT) # ref_rgb_8uc3_change_color_orisize = cv2.resize(ref_rgb_8uc3_change_color_768, # (user_rgb_8uc3_orisize.shape[1], user_rgb_8uc3_orisize.shape[0]), # interpolation=cv2.INTER_CUBIC) # user_res_8uc3_orisize = user_rgb_8uc3_orisize * ( # 1 - user_matting_8uc3_change_color_orisize / 255) + hair_gene_color_8uc3_orisize * ( # user_matting_8uc3_change_color_orisize / 255) # user_res_8uc3_orisize = (np.clip(user_res_8uc3_orisize, 0, 255)).astype(np.uint8) # cv2.imshow("user_rgb_8uc3_orisize", user_rgb_8uc3_orisize) # cv2.imshow("hair_gene_color_8uc3_orisize", hair_gene_color_8uc3_orisize) # cv2.imshow("user_matting_8uc3_change_color_orisize", user_matting_8uc3_change_color_orisize) # cv2.waitKey() # if self.use_enhance: user_res_8uc3_orisize = self.face_enhance.process(hair_gene_color_8uc3_orisize, landmarks_origin_img_1k) user_res_8uc3_orisize_hsv = cv2.cvtColor(user_res_8uc3_orisize, cv2.COLOR_BGR2HSV) user_res_8uc3_orisize_hsv[:, :, 1:2] = np.clip(user_res_8uc3_orisize_hsv[:, :, 1:2] * 1.15, 0, 255) user_res_8uc3_orisize = cv2.cvtColor(user_res_8uc3_orisize_hsv, cv2.COLOR_HSV2BGR) user_res_8uc3_orisize2 = user_rgb_8uc3_orisize * ( 1 - user_matting_8uc3_orisize / 255) + user_res_8uc3_orisize * ( user_matting_8uc3_orisize / 255) user_res_8uc3_orisize2 = (np.clip(user_res_8uc3_orisize2, 0, 255)).astype(np.uint8) # user_res_8uc3_orisize_hsv = cv2.cvtColor(user_res_8uc3_orisize, cv2.COLOR_BGR2HSV) # user_res_8uc3_orisize_hsv[:, :, 1:2] = np.clip(user_res_8uc3_orisize_hsv[:, :, 1:2] * 1.1, 0, 255) # user_res_8uc3_orisize = cv2.cvtColor(user_res_8uc3_orisize_hsv, cv2.COLOR_HSV2BGR) # user_res_8uc3_orisize = self.rein_hard_raw(user_res_8uc3_orisize, user_matting_8uc3_bald_orisize, # ref_rgb_8uc3_change_color_768, ref_matting_8uc3_change_color_768) # user_res_8uc3_orisize = self.rein_hard(user_res_8uc3_orisize, user_matting_8uc3_bald_orisize, # ref_rgb_8uc3_change_color_768, ref_matting_8uc3_change_color_768) # user_res_8uc3_orisize = self.rein_hard(user_res_8uc3_orisize, user_matting_8uc3_bald_orisize, # ref_rgb_8uc3_change_color_768, ref_matting_8uc3_change_color_768) # user_res_8uc3_orisize = user_rgb_8uc3_orisize * ( # 1 - user_matting_8uc3_bald_orisize / 255) + user_res_8uc3_orisize * ( # user_matting_8uc3_bald_orisize / 255) # user_res_8uc3_orisize = (np.clip(user_res_8uc3_orisize, 0, 255)).astype(np.uint8) # middle_show = np.concatenate( # (user_rgb_8uc3_orisize, user_matting_8uc3_bald_orisize, user_res_8uc3_orisize), axis=1) # ratio = 1536. / max(middle_show.shape[:2]) # middle_show = cv2.resize(middle_show, (0, 0), fx=ratio, fy=ratio) # cv2.imshow('middle_show', middle_show) # cv2.waitKey(0) # user_res_8uc3_orisize = np.concatenate((user_res_8uc3_orisize, user_matting_8uc3_bald_orisize), axis=1) if return_matting: return user_res_8uc3_orisize2, user_matting_8uc1_bald_orisize, 0 else: return user_res_8uc3_orisize2, 0 def get_hair_color(self, input_img, input_mask): # 确保图像和掩膜具有相同的尺寸 if input_img.shape[:2] != input_mask.shape[:2]: return None # 如果没有找到头发区域 # 创建掩膜，提取头发区域 hair_region = cv2.bitwise_and(input_img, input_img, mask=input_mask) # 计算头发区域的颜色（去除黑色区域） hair_pixels = hair_region[hair_region != 0] if len(hair_pixels) == 0: return None # 如果没有找到头发区域 # 计算平均颜色 average_color = np.mean(hair_pixels, axis=0) is_black = (average_color[2] * 0.299 + average_color[1] * 0.587 + average_color[1] * 0.114) < 60 return is_black def infer_haircolor_tj(self, user_rgb_8uc3_orisize, haircolor_dir): landmarks_origin_img_1k= self.get_landmark.forward(user_rgb_8uc3_orisize) if landmarks_origin_img_1k is None: return None, 10001 t0 = time.time() _, user_matting_8uc1_bald_orisize = self.process_data_infer.generator_matte.matte_inference(user_rgb_8uc3_orisize, landmarks_origin_img_1k) user_matting_mask_8uc3_orisize = np.repeat(user_matting_8uc1_bald_orisize[:, :, np.newaxis], 3, axis=2) ref_rgb_8uc3_change_color_768 = cv2.imread(os.path.join(haircolor_dir, "ref_rgb_8uc3_color_768.png")) ref_matting_8uc3_change_color_768 = cv2.imread(os.path.join(haircolor_dir, "ref_matting_8uc3_color_768.png")) user_rgb_8uc3_change_color_768, user_matting_8uc3_change_color_768, user_hair_color_M = \ self.process_data_infer.get_prepare_hair_color_user_data(user_rgb_8uc3_orisize,landmarks_origin_img_1k) # start_time = time.time() hair_gene_color_8uc3_768, user_matting_mask_8uc3_768 = self.change_haircolor.Change_Hair_inference( user_rgb_8uc3_change_color_768, user_matting_8uc3_change_color_768, ref_rgb_8uc3_change_color_768, ref_matting_8uc3_change_color_768) hair_gene_color_8uc3_orisize = cv2.warpAffine(hair_gene_color_8uc3_768, cv2.invertAffineTransform(user_hair_color_M), (user_rgb_8uc3_orisize.shape[1], user_rgb_8uc3_orisize.shape[0]), dst=user_rgb_8uc3_orisize.copy(), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_TRANSPARENT) user_res_8uc3_orisize = self.face_enhance.process(hair_gene_color_8uc3_orisize, landmarks_origin_img_1k) user_res_fc32_orisize = user_res_8uc3_orisize.astype(np.float32) / 255 # user_res_fc32_orisize_hsv = cv2.cvtColor(user_res_fc32_orisize, cv2.COLOR_BGR2HSV) ref_rgb_fc32_change_color_768 = ref_rgb_8uc3_change_color_768.astype(np.float32) / 255 ref_matting_fc32_change_color_768 = ref_matting_8uc3_change_color_768.astype(np.float32) / 255 ref_rgb_avg_std = self.getavgstd(ref_rgb_fc32_change_color_768, ref_matting_fc32_change_color_768) user_matting_mask_fc32_orisize = user_matting_mask_8uc3_orisize.astype(np.float32) / 255 user_res_fc32_orisize, src_avg_std = self.reinhard_rgb(user_res_fc32_orisize, user_matting_mask_fc32_orisize, ref_rgb_avg_std, ratio=1.0) user_res_8uc3_orisize_reinhard = (user_res_fc32_orisize * 255).astype(np.uint8) user_res_8uc3_orisize_reinhard_LAB = cv2.cvtColor(user_res_8uc3_orisize_reinhard, cv2.COLOR_BGR2LAB) user_res_8uc3_orisize_LAB = cv2.cvtColor(user_res_8uc3_orisize, cv2.COLOR_BGR2LAB) mix_ratio = 0.1 user_res_8uc3_orisize_reinhard_LAB[:, :, 0] = user_res_8uc3_orisize_LAB[:, :, 0] * mix_ratio + user_res_8uc3_orisize_reinhard_LAB[:, :, 0] * ( 1 - mix_ratio) user_res_8uc3_orisize = cv2.cvtColor(user_res_8uc3_orisize_reinhard_LAB, cv2.COLOR_LAB2BGR) user_res_8uc3_orisize2 = user_rgb_8uc3_orisize * ( 1 - user_matting_mask_8uc3_orisize / 255) + user_res_8uc3_orisize * ( user_matting_mask_8uc3_orisize / 255) user_res_8uc3_orisize2 = (np.clip(user_res_8uc3_orisize2, 0, 255)).astype(np.uint8) # print('change color, last proces,', time.time() - t3) return user_res_8uc3_orisize2, 0 def infer_haircolor_new(self, user_rgb_8uc3_orisize, haircolor_dir, target_hair_color, return_matting=False): landmarks_origin_img_1k= self.get_landmark.forward(user_rgb_8uc3_orisize) if landmarks_origin_img_1k is None: return None, 10001 need_process = (target_hair_color[0] * 0.299 + target_hair_color[1] * 0.587 + target_hair_color[2] * 0.114) > 150 _, user_matting_8uc1_bald_orisize = self.process_data_infer.generator_matte.matte_inference(user_rgb_8uc3_orisize, landmarks_origin_img_1k) user_matting_mask_8uc3_orisize = np.repeat(user_matting_8uc1_bald_orisize[:, :, np.newaxis], 3, axis=2) ref_rgb_8uc3_change_color_768 = cv2.imread(os.path.join(haircolor_dir, "ref_rgb_8uc3_color_768.png")) ref_matting_8uc3_change_color_768 = cv2.imread(os.path.join(haircolor_dir, "ref_matting_8uc3_color_768.png")) user_rgb_8uc3_change_color_768, user_matting_8uc3_change_color_768, user_hair_color_M = \ self.process_data_infer.get_prepare_hair_color_user_data(user_rgb_8uc3_orisize, landmarks_origin_img_1k) # cv2.imshow("user_matting_8uc3_change_color_768", user_matting_8uc3_change_color_768) # cv2.waitKey() start_time = time.time() need_process=True if need_process: haircolor_dir_tmp = os.path.join(config.get('default', "haircolorDir"), config.get('default', "baseColor_ID")) face_base, status1 = self.infer_haircolor_tj(user_rgb_8uc3_orisize, haircolor_dir_tmp) # cv2.imwrite('/home/student/Desktop/tmp_color/need/face_base.png', face_base) if status1 == 0: return face_base, user_matting_8uc1_bald_orisize, 0 else: need_process = False if not need_process: hair_gene_color_8uc3_768, user_matting_mask_8uc3_768 = self.change_haircolor.Change_Hair_inference( user_rgb_8uc3_change_color_768, user_matting_8uc3_change_color_768, ref_rgb_8uc3_change_color_768, ref_matting_8uc3_change_color_768) # cv2.imwrite("hair_gene_color_8uc3_768.png", hair_gene_color_8uc3_768) # print("Change_Hair_inference cost: ", time.time()-start_time) # cv2.imshow("hair_gene_color_8uc3_768: ", hair_gene_color_8uc3_768) # cv2.imshow("user_matting_mask_8uc3_768: ", user_matting_mask_8uc3_768) # cv2.waitKey() # user_matting_mask_8uc3_orisize = cv2.warpAffine(user_matting_mask_8uc3_768, # cv2.invertAffineTransform(user_hair_color_M), # (user_rgb_8uc3_orisize.shape[1], # user_rgb_8uc3_orisize.shape[0]), flags=cv2.INTER_CUBIC) hair_gene_color_8uc3_orisize = cv2.warpAffine(hair_gene_color_8uc3_768, cv2.invertAffineTransform(user_hair_color_M), (user_rgb_8uc3_orisize.shape[1], user_rgb_8uc3_orisize.shape[0]), dst=user_rgb_8uc3_orisize.copy(), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_TRANSPARENT) user_res_8uc3_orisize = self.face_enhance.process(hair_gene_color_8uc3_orisize, landmarks_origin_img_1k) # cv2.imshow("user_res_8uc3_orisize", user_res_8uc3_orisize) # cv2.waitKey() user_res_fc32_orisize = user_res_8uc3_orisize.astype(np.float32) / 255 # user_res_fc32_orisize_hsv = cv2.cvtColor(user_res_fc32_orisize, cv2.COLOR_BGR2HSV) ref_rgb_fc32_change_color_768 = ref_rgb_8uc3_change_color_768.astype(np.float32) / 255 ref_matting_fc32_change_color_768 = ref_matting_8uc3_change_color_768.astype(np.float32) / 255 ref_rgb_avg_std = self.getavgstd(ref_rgb_fc32_change_color_768, ref_matting_fc32_change_color_768) user_matting_mask_fc32_orisize = user_matting_mask_8uc3_orisize.astype(np.float32) / 255 user_res_fc32_orisize, src_avg_std = self.reinhard_rgb(user_res_fc32_orisize, user_matting_mask_fc32_orisize, ref_rgb_avg_std, ratio=1.0) user_res_8uc3_orisize_reinhard = (user_res_fc32_orisize * 255).astype(np.uint8) user_res_8uc3_orisize_reinhard_LAB = cv2.cvtColor(user_res_8uc3_orisize_reinhard, cv2.COLOR_BGR2LAB) user_res_8uc3_orisize_LAB = cv2.cvtColor(user_res_8uc3_orisize, cv2.COLOR_BGR2LAB) mix_ratio = 0.1 user_res_8uc3_orisize_reinhard_LAB[:, :, 0] = user_res_8uc3_orisize_LAB[:, :, 0] * mix_ratio + user_res_8uc3_orisize_reinhard_LAB[:, :, 0] * ( 1 - mix_ratio) user_res_8uc3_orisize = cv2.cvtColor(user_res_8uc3_orisize_reinhard_LAB, cv2.COLOR_LAB2BGR) user_res_8uc3_orisize2 = user_rgb_8uc3_orisize * ( 1 - user_matting_mask_8uc3_orisize / 255) + user_res_8uc3_orisize * ( user_matting_mask_8uc3_orisize / 255) user_res_8uc3_orisize2 = (np.clip(user_res_8uc3_orisize2, 0, 255)).astype(np.uint8) # cv2.imwrite('/home/student/Desktop/tmp_color/need/face_base2.png', user_res_8uc3_orisize2) # 调试用，路径不存在会导致崩溃 if return_matting: return user_res_8uc3_orisize2, user_matting_8uc1_bald_orisize, 0 else: return user_res_8uc3_orisize2, 0 def infer_haircolor_0311(self, user_rgb_8uc3_orisize, haircolor_dir): landmarks_origin_img_1k = self.get_landmark.forward(user_rgb_8uc3_orisize) if landmarks_origin_img_1k is None: return None, 10001 _, user_matting_8uc1_bald_orisize = self.process_data.generator_matte.matte_inference(user_rgb_8uc3_orisize, landmarks_origin_img_1k) user_matting_mask_8uc3_orisize = np.repeat(user_matting_8uc1_bald_orisize[:, :, np.newaxis], 3, axis=2) ref_rgb_8uc3_change_color_768 = cv2.imread(os.path.join(haircolor_dir, "ref_rgb_8uc3_color_768.png")) ref_matting_8uc3_change_color_768 = cv2.imread(os.path.join(haircolor_dir, "ref_matting_8uc3_color_768.png")) # user_rgb_8uc3_orisize = self.preprocess_haircolor(user_rgb_8uc3_orisize, user_matting_8uc3_bald_orisize, # ref_rgb_8uc3_change_color_768, ref_matting_8uc3_change_color_768) user_rgb_8uc3_change_color_768, user_matting_8uc3_change_color_768, user_hair_color_M = \ self.process_data.get_prepare_hair_color_user_data(user_rgb_8uc3_orisize, landmarks_origin_img_1k) # cv2.imshow("user_matting_8uc3_change_color_768", user_matting_8uc3_change_color_768) # cv2.waitKey() # start_time = time.time() hair_gene_color_8uc3_768, user_matting_mask_8uc3_768 = self.change_haircolor.Change_Hair_inference( user_rgb_8uc3_change_color_768, user_matting_8uc3_change_color_768, ref_rgb_8uc3_change_color_768, ref_matting_8uc3_change_color_768) # print("Change_Hair_inference cost: ", time.time()-start_time) # cv2.imshow("hair_gene_color_8uc3_768: ", hair_gene_color_8uc3_768) # cv2.imshow("user_matting_mask_8uc3_768: ", user_matting_mask_8uc3_768) # cv2.waitKey() # user_matting_8uc3_orisize = cv2.warpAffine(user_matting_mask_8uc3_768, # cv2.invertAffineTransform(user_hair_color_M), # (user_rgb_8uc3_orisize.shape[1], # user_rgb_8uc3_orisize.shape[0]), flags=cv2.INTER_CUBIC) hair_gene_color_8uc3_orisize = cv2.warpAffine(hair_gene_color_8uc3_768, cv2.invertAffineTransform(user_hair_color_M), (user_rgb_8uc3_orisize.shape[1], user_rgb_8uc3_orisize.shape[0]), dst=user_rgb_8uc3_orisize.copy(), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_TRANSPARENT) # ref_rgb_8uc3_change_color_orisize = cv2.resize(ref_rgb_8uc3_change_color_768, # (user_rgb_8uc3_orisize.shape[1], user_rgb_8uc3_orisize.shape[0]), # interpolation=cv2.INTER_CUBIC) # user_res_8uc3_orisize = user_rgb_8uc3_orisize * ( # 1 - user_matting_8uc3_change_color_orisize / 255) + hair_gene_color_8uc3_orisize * ( # user_matting_8uc3_change_color_orisize / 255) # user_res_8uc3_orisize = (np.clip(user_res_8uc3_orisize, 0, 255)).astype(np.uint8) # cv2.imshow("user_rgb_8uc3_orisize", user_rgb_8uc3_orisize) # cv2.imshow("hair_gene_color_8uc3_orisize", hair_gene_color_8uc3_orisize) # cv2.imshow("user_matting_8uc3_change_color_orisize", user_matting_8uc3_change_color_orisize) # cv2.waitKey() # if self.use_enhance: user_res_8uc3_orisize = self.face_enhance.process(hair_gene_color_8uc3_orisize, landmarks_origin_img_1k) # user_res_8uc3_orisize_hsv = cv2.cvtColor(user_res_8uc3_orisize, cv2.COLOR_BGR2HSV) # user_res_8uc3_orisize_hsv[:, :, 1:2] = np.clip(user_res_8uc3_orisize_hsv[:, :, 1:2] * 1.15, 0, 255) # # user_res_8uc3_orisize_hsv[:, :, 2:3] = np.clip(user_res_8uc3_orisize_hsv[:, :, 2:3], 100, 255) # user_res_8uc3_orisize = cv2.cvtColor(user_res_8uc3_orisize_hsv, cv2.COLOR_HSV2BGR) user_res_fc32_orisize = user_res_8uc3_orisize.astype(np.float32) / 255 user_res_fc32_orisize_hsv = cv2.cvtColor(user_res_fc32_orisize, cv2.COLOR_BGR2HSV) ref_rgb_fc32_change_color_768 = ref_rgb_8uc3_change_color_768.astype(np.float32) / 255 ref_rgb_fc32_change_color_768_hsv = cv2.cvtColor(ref_rgb_fc32_change_color_768, cv2.COLOR_BGR2HSV) loc_index = ref_matting_8uc3_change_color_768[:, :, 0].nonzero() color_val = ref_rgb_fc32_change_color_768_hsv[loc_index] mean_ref_haircolor_hsv = np.mean(color_val, axis=0) print("mean_ref_haircolor_hsv: ", mean_ref_haircolor_hsv) loc_index = user_matting_mask_8uc3_orisize[:, :, 0].nonzero() color_val = user_res_fc32_orisize_hsv[loc_index] mean_user_haircolor_hsv = np.mean(color_val, axis=0) print("mean_user_haircolor_hsv: ", mean_user_haircolor_hsv) user_res_fc32_orisize_hsv[:, :, 0:1] = mean_ref_haircolor_hsv[0] # ratio_h = mean_ref_haircolor_hsv[0] / mean_user_haircolor_hsv[0] # user_res_fc32_orisize_hsv[:, :, 0:1] = np.clip(user_res_fc32_orisize_hsv[:, :, 0:1] + mean_ref_haircolor_hsv[0] - mean_user_haircolor_hsv[0], 0, 360.0) user_res_fc32_orisize_hsv[:, :, 1:2] = np.clip(user_res_fc32_orisize_hsv[:, :, 1:2] * 1.15, 0, 1.0) user_res_fc32_orisize_hsv[:, :, 2:3] = np.clip(user_res_fc32_orisize_hsv[:, :, 2:3], 0, 1.0) user_res_fc32_orisize = cv2.cvtColor(user_res_fc32_orisize_hsv, cv2.COLOR_HSV2BGR) user_res_8uc3_orisize_new = (user_res_fc32_orisize * 255).astype(np.uint8) user_res_8uc3_orisize_new_LAB = cv2.cvtColor(user_res_8uc3_orisize_new, cv2.COLOR_BGR2LAB) user_res_8uc3_orisize_LAB = cv2.cvtColor(user_res_8uc3_orisize, cv2.COLOR_BGR2LAB) user_res_8uc3_orisize_new_LAB[:, :, 0] = user_res_8uc3_orisize_LAB[:, :, 0] user_res_8uc3_orisize = cv2.cvtColor(user_res_8uc3_orisize_new_LAB, cv2.COLOR_LAB2BGR) user_res_8uc3_orisize2 = user_rgb_8uc3_orisize * (1 - user_matting_mask_8uc3_orisize / 255) + user_res_8uc3_orisize * (user_matting_mask_8uc3_orisize / 255) user_res_8uc3_orisize2 = (np.clip(user_res_8uc3_orisize2, 0, 255)).astype(np.uint8) # middle_show = np.concatenate( # (user_rgb_8uc3_orisize, user_matting_mask_8uc3_orisize, user_res_8uc3_orisize, user_res_8uc3_orisize2), axis=1) # ratio = 1536. / max(middle_show.shape[:2]) # middle_show = cv2.resize(middle_show, (0, 0), fx=ratio, fy=ratio) # cv2.imshow('middle_show', middle_show) # cv2.waitKey(0) # user_res_8uc3_orisize_hsv = cv2.cvtColor(user_res_8uc3_orisize, cv2.COLOR_BGR2HSV) # user_res_8uc3_orisize_hsv[:, :, 1:2] = np.clip(user_res_8uc3_orisize_hsv[:, :, 1:2] * 1.1, 0, 255) # user_res_8uc3_orisize = cv2.cvtColor(user_res_8uc3_orisize_hsv, cv2.COLOR_HSV2BGR) # user_res_8uc3_orisize = self.rein_hard_raw(user_res_8uc3_orisize, user_matting_8uc3_bald_orisize, # ref_rgb_8uc3_change_color_768, ref_matting_8uc3_change_color_768) # user_res_8uc3_orisize = self.rein_hard(user_res_8uc3_orisize, user_matting_8uc3_bald_orisize, # ref_rgb_8uc3_change_color_768, ref_matting_8uc3_change_color_768) # user_res_8uc3_orisize = self.rein_hard(user_res_8uc3_orisize, user_matting_8uc3_bald_orisize, # ref_rgb_8uc3_change_color_768, ref_matting_8uc3_change_color_768) # user_res_8uc3_orisize = user_rgb_8uc3_orisize * ( # 1 - user_matting_8uc3_bald_orisize / 255) + user_res_8uc3_orisize * ( # user_matting_8uc3_bald_orisize / 255) # user_res_8uc3_orisize = (np.clip(user_res_8uc3_orisize, 0, 255)).astype(np.uint8) # middle_show = np.concatenate( # (user_rgb_8uc3_orisize, user_matting_8uc3_bald_orisize, user_res_8uc3_orisize), axis=1) # ratio = 1536. / max(middle_show.shape[:2]) # middle_show = cv2.resize(middle_show, (0, 0), fx=ratio, fy=ratio) # cv2.imshow('middle_show', middle_show) # cv2.waitKey(0) # user_res_8uc3_orisize = np.concatenate((user_res_8uc3_orisize, user_matting_8uc3_bald_orisize), axis=1) return user_res_8uc3_orisize2, 0 def getavgstd(self, image, mask): if mask.shape[2] == 1: mask = np.repeat(mask, 3, axis=2) mask_index = np.flatnonzero((mask > 0.1).any(axis=2)) lab_layer = image.reshape(-1, 3)[mask_index] if len(lab_layer) < 100: return None lab_layer = np.float32(lab_layer) avg = [] std = [] image_avg_l = np.mean(lab_layer[:, 0]) image_std_l = np.std(lab_layer[:, 0]) image_avg_a = np.mean(lab_layer[:, 1]) image_std_a = np.std(lab_layer[:, 1]) image_avg_b = np.mean(lab_layer[:, 2]) image_std_b = np.std(lab_layer[:, 2]) avg.append(image_avg_l) avg.append(image_avg_a) avg.append(image_avg_b) std.append(image_std_l) std.append(image_std_a) std.append(image_std_b) return avg, std def reinhard(self, origin_img, mask, tar_avg_std, ratio=0.5): origin_img_lab = cv2.cvtColor(origin_img, cv2.COLOR_BGR2LAB) src_avg_std = self.getavgstd(origin_img_lab, mask) src_avg_std = np.float32(src_avg_std) # origin_img_lab[:, :, 0] = origin_img_lab[:, :, 0] + (tar_avg_std[0][0] - src_avg_std[0][0]) * ratio origin_img_lab[:, :, 1] = origin_img_lab[:, :, 1] + (tar_avg_std[0][1] - src_avg_std[0][1]) * ratio origin_img_lab[:, :, 2] = origin_img_lab[:, :, 2] + (tar_avg_std[0][2] - src_avg_std[0][2]) * ratio img_ret = cv2.cvtColor(origin_img_lab, cv2.COLOR_LAB2BGR) img_ret = np.clip(img_ret, 0.01, 0.99) return img_ret, src_avg_std def reinhard_rgb_old(self, origin_img, mask, tar_avg_std, ratio=0.5): # origin_img_lab = cv2.cvtColor(origin_img, cv2.COLOR_BGR2LAB) origin_img_lab = origin_img src_avg_std = self.getavgstd(origin_img_lab, mask) src_avg_std = np.float32(src_avg_std) origin_img_lab[:, :, 0] = origin_img_lab[:, :, 0] + (tar_avg_std[0][0] - src_avg_std[0][0]) * (tar_avg_std[1][0] / src_avg_std[1][0]) origin_img_lab[:, :, 1] = origin_img_lab[:, :, 1] + (tar_avg_std[0][1] - src_avg_std[0][1]) * (tar_avg_std[1][1] / src_avg_std[1][1]) origin_img_lab[:, :, 2] = origin_img_lab[:, :, 2] + (tar_avg_std[0][2] - src_avg_std[0][2]) * (tar_avg_std[1][2] / src_avg_std[1][2]) # img_ret = cv2.cvtColor(origin_img_lab, cv2.COLOR_LAB2BGR) img_ret = origin_img_lab img_ret = np.clip(img_ret, 0.0, 1.0) return img_ret, src_avg_std def reinhard_rgb(self, origin_img, mask, tar_avg_std, ratio=0.5): # origin_img_lab = cv2.cvtColor(origin_img, cv2.COLOR_BGR2LAB) origin_img_lab = origin_img src_avg_std = self.getavgstd(origin_img_lab, mask) src_avg_std = np.float32(src_avg_std) origin_img_lab[:, :, 0] = origin_img_lab[:, :, 0] + (tar_avg_std[0][0] - src_avg_std[0][0]) * ratio origin_img_lab[:, :, 1] = origin_img_lab[:, :, 1] + (tar_avg_std[0][1] - src_avg_std[0][1]) * ratio origin_img_lab[:, :, 2] = origin_img_lab[:, :, 2] + (tar_avg_std[0][2] - src_avg_std[0][2]) * ratio # img_ret = cv2.cvtColor(origin_img_lab, cv2.COLOR_LAB2BGR) img_ret = origin_img_lab img_ret = np.clip(img_ret, 0.0, 1.0) return img_ret, src_avg_std def infer_haircolor_0313(self, user_rgb_8uc3_orisize, haircolor_dir): landmarks_origin_img_1k = self.get_landmark.forward(user_rgb_8uc3_orisize) if landmarks_origin_img_1k is None: return None, 10001 _, user_matting_8uc1_bald_orisize = self.process_data.generator_matte.matte_inference(user_rgb_8uc3_orisize, landmarks_origin_img_1k) user_matting_mask_8uc3_orisize = np.repeat(user_matting_8uc1_bald_orisize[:, :, np.newaxis], 3, axis=2) ref_rgb_8uc3_change_color_768 = cv2.imread(os.path.join(haircolor_dir, "ref_rgb_8uc3_color_768.png")) ref_matting_8uc3_change_color_768 = cv2.imread(os.path.join(haircolor_dir, "ref_matting_8uc3_color_768.png")) user_rgb_8uc3_change_color_768, user_matting_8uc3_change_color_768, user_hair_color_M = \ self.process_data.get_prepare_hair_color_user_data(user_rgb_8uc3_orisize, landmarks_origin_img_1k) # start_time = time.time() hair_gene_color_8uc3_768, user_matting_mask_8uc3_768 = self.change_haircolor.Change_Hair_inference( user_rgb_8uc3_change_color_768, user_matting_8uc3_change_color_768, ref_rgb_8uc3_change_color_768, ref_matting_8uc3_change_color_768) hair_gene_color_8uc3_orisize = cv2.warpAffine(hair_gene_color_8uc3_768, cv2.invertAffineTransform(user_hair_color_M), (user_rgb_8uc3_orisize.shape[1], user_rgb_8uc3_orisize.shape[0]), dst=user_rgb_8uc3_orisize.copy(), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_TRANSPARENT) # if self.use_enhance: # user_res_8uc3_orisize = self.face_enhance.process(hair_gene_color_8uc3_orisize, landmarks_origin_img_1k) user_res_8uc3_orisize = hair_gene_color_8uc3_orisize.copy() # user_rgb_8uc3_orisize_LAB = cv2.cvtColor(user_rgb_8uc3_orisize, cv2.COLOR_BGR2LAB) # user_res_8uc3_orisize_LAB = cv2.cvtColor(user_res_8uc3_orisize, cv2.COLOR_BGR2LAB) # user_res_8uc3_orisize_LAB[:, :, 0] = user_rgb_8uc3_orisize_LAB[:, :, 0] # user_res_8uc3_orisize = cv2.cvtColor(user_res_8uc3_orisize_LAB, cv2.COLOR_LAB2BGR) # middle_show = np.concatenate((user_rgb_8uc3_orisize, hair_gene_color_8uc3_orisize, user_res_8uc3_orisize, user_res_8uc3_orisize_new), axis=1) # ratio = 1536. / max(middle_show.shape[:2]) # middle_show = cv2.resize(middle_show, (0, 0), fx=ratio, fy=ratio) # cv2.imshow('middle_show', middle_show) # cv2.waitKey(0) # #### version 1 # user_res_fc32_orisize = user_res_8uc3_orisize.astype(np.float32) / 255 # user_res_fc32_orisize_hsv = cv2.cvtColor(user_res_fc32_orisize, cv2.COLOR_BGR2HSV) # ref_rgb_fc32_change_color_768 = ref_rgb_8uc3_change_color_768.astype(np.float32) / 255 # ref_rgb_fc32_change_color_768_hsv = cv2.cvtColor(ref_rgb_fc32_change_color_768, cv2.COLOR_BGR2HSV) # # ref_hair_hsv_mean_std = self.getavgstd(ref_rgb_fc32_change_color_768_hsv, ref_matting_8uc3_change_color_768) # user_hair_hsv_mean_std = self.getavgstd(user_res_fc32_orisize_hsv, user_matting_mask_8uc3_orisize) # print("ref_hair_hsv_mean_std: ", ref_hair_hsv_mean_std) # print("user_hair_hsv_mean_std: ", user_hair_hsv_mean_std) # # loc_index = ref_matting_8uc3_change_color_768[:, :, 0].nonzero() # color_val = ref_rgb_fc32_change_color_768_hsv[loc_index] # mean_ref_haircolor_hsv = np.mean(color_val, axis=0) # print("mean_ref_haircolor_hsv: ", mean_ref_haircolor_hsv) # loc_index = user_matting_mask_8uc3_orisize[:, :, 0].nonzero() # color_val = user_res_fc32_orisize_hsv[loc_index] # mean_user_haircolor_hsv = np.mean(color_val, axis=0) # print("mean_user_haircolor_hsv: ", mean_user_haircolor_hsv) # # # user_res_fc32_orisize_hsv[:, :, 0:1] = user_res_fc32_orisize_hsv[:, :, 0:1] + mean_ref_haircolor_hsv[0] - mean_user_haircolor_hsv[0] # user_res_fc32_orisize_hsv[:, :, 0:1] = user_res_fc32_orisize_hsv[:, :, 0:1] + ref_hair_hsv_mean_std[0][0] - user_hair_hsv_mean_std[0][0] # # user_res_fc32_orisize_hsv[:, :, 0:1] = np.clip(user_res_fc32_orisize_hsv[:, :, 0:1], 0, 360) # # user_res_fc32_orisize_hsv[:, :, 1:2] = user_res_fc32_orisize_hsv[:, :, 1:2] + mean_ref_haircolor_hsv[1] - mean_user_haircolor_hsv[1] # user_res_fc32_orisize_hsv[:, :, 1:2] = user_res_fc32_orisize_hsv[:, :, 1:2] + ref_hair_hsv_mean_std[0][1] - user_hair_hsv_mean_std[0][1] # user_res_fc32_orisize_hsv[:, :, 1:2] = np.clip(user_res_fc32_orisize_hsv[:, :, 1:2], 0, 1.0) # # user_res_fc32_orisize_hsv[:, :, 2:3] = user_res_fc32_orisize_hsv[:, :, 2:3] + mean_ref_haircolor_hsv[2] - mean_user_haircolor_hsv[2] # user_res_fc32_orisize_hsv[:, :, 2:3] = user_res_fc32_orisize_hsv[:, :, 2:3] + ref_hair_hsv_mean_std[0][2] - user_hair_hsv_mean_std[0][2] # user_res_fc32_orisize_hsv[:, :, 2:3] = np.clip(user_res_fc32_orisize_hsv[:, :, 2:3], 0, 1.0) # # # user_res_fc32_orisize_hsv[:, :, 0:1] = mean_ref_haircolor_hsv[0] # # user_res_fc32_orisize_hsv[:, :, 1:2] = np.clip(user_res_fc32_orisize_hsv[:, :, 1:2] * 1.15, 0, 1.0) # # user_res_fc32_orisize_hsv[:, :, 2:3] = np.clip(user_res_fc32_orisize_hsv[:, :, 2:3], 0, 1.0) # user_res_fc32_orisize = cv2.cvtColor(user_res_fc32_orisize_hsv, cv2.COLOR_HSV2BGR) # user_res_8uc3_orisize_new = (user_res_fc32_orisize * 255).astype(np.uint8) # # user_res_8uc3_orisize_new_LAB = cv2.cvtColor(user_res_8uc3_orisize_new, cv2.COLOR_BGR2LAB) # user_res_8uc3_orisize_LAB = cv2.cvtColor(user_res_8uc3_orisize, cv2.COLOR_BGR2LAB) # user_res_8uc3_orisize_new_LAB[:, :, 0] = user_res_8uc3_orisize_LAB[:, :, 0] # user_res_8uc3_orisize = cv2.cvtColor(user_res_8uc3_orisize_new_LAB, cv2.COLOR_LAB2BGR) # # user_res_fc32_orisize = user_res_8uc3_orisize.astype(np.float32) / 255 # user_res_fc32_orisize_hsv = cv2.cvtColor(user_res_fc32_orisize, cv2.COLOR_BGR2HSV) # final_user_hair_hsv_mean_std = self.getavgstd(user_res_fc32_orisize_hsv, user_matting_mask_8uc3_orisize) # print("final_user_hair_hsv_mean_std: ", final_user_hair_hsv_mean_std) # #### version 1 #### version 2 user_res_fc32_orisize = user_res_8uc3_orisize.astype(np.float32) / 255 # user_res_fc32_orisize_hsv = cv2.cvtColor(user_res_fc32_orisize, cv2.COLOR_BGR2HSV) ref_rgb_fc32_change_color_768 = ref_rgb_8uc3_change_color_768.astype(np.float32) / 255 ref_matting_fc32_change_color_768 = ref_matting_8uc3_change_color_768.astype(np.float32) / 255 # ref_rgb_fc32_change_color_768_hsv = cv2.cvtColor(ref_rgb_fc32_change_color_768, cv2.COLOR_BGR2HSV) # ref_rgb_fc32_change_color_768_lab = cv2.cvtColor(ref_rgb_fc32_change_color_768, cv2.COLOR_BGR2LAB) # ref_avg_std = self.getavgstd(ref_rgb_fc32_change_color_768_lab, ref_matting_fc32_change_color_768) # # user_res_fc32_orisize_lab = cv2.cvtColor(user_res_fc32_orisize, cv2.COLOR_BGR2LAB) # # user_avg_std = self.getavgstd(user_res_fc32_orisize_lab, user_matting_mask_8uc3_orisize) # user_matting_mask_fc32_orisize = user_matting_mask_8uc3_orisize.astype(np.float32) / 255 # user_res_fc32_orisize, src_avg_std = self.reinhard(user_res_fc32_orisize, user_matting_mask_fc32_orisize, ref_avg_std, ratio=0.99) ref_rgb_avg_std = self.getavgstd(ref_rgb_fc32_change_color_768, ref_matting_fc32_change_color_768) # user_rgb_avg_std = self.getavgstd(user_res_fc32_orisize, user_matting_mask_8uc3_orisize) user_matting_mask_fc32_orisize = user_matting_mask_8uc3_orisize.astype(np.float32) / 255 user_res_fc32_orisize, src_avg_std = self.reinhard_rgb(user_res_fc32_orisize, user_matting_mask_fc32_orisize, ref_rgb_avg_std, ratio=0.65) user_res_8uc3_orisize_reinhard = (user_res_fc32_orisize * 255).astype(np.uint8) #### version 2 # middle_show = np.concatenate((user_rgb_8uc3_orisize, user_res_8uc3_orisize, user_res_8uc3_orisize_reinhard), axis=1) # ratio = 1536. / max(middle_show.shape[:2]) # middle_show = cv2.resize(middle_show, (0, 0), fx=ratio, fy=ratio) # cv2.imshow('middle_show', middle_show) # cv2.waitKey(0) user_res_8uc3_orisize_reinhard_LAB = cv2.cvtColor(user_res_8uc3_orisize_reinhard, cv2.COLOR_BGR2LAB) user_res_8uc3_orisize_LAB = cv2.cvtColor(user_res_8uc3_orisize, cv2.COLOR_BGR2LAB) mix_ratio = 0.1 user_res_8uc3_orisize_reinhard_LAB[:, :, 0] = user_res_8uc3_orisize_LAB[:, :, 0] * mix_ratio + user_res_8uc3_orisize_reinhard_LAB[:, :, 0] * (1 - mix_ratio) user_res_8uc3_orisize = cv2.cvtColor(user_res_8uc3_orisize_reinhard_LAB, cv2.COLOR_LAB2BGR) user_res_8uc3_orisize2 = user_rgb_8uc3_orisize * (1 - user_matting_mask_8uc3_orisize / 255) + user_res_8uc3_orisize * (user_matting_mask_8uc3_orisize / 255) user_res_8uc3_orisize2 = (np.clip(user_res_8uc3_orisize2, 0, 255)).astype(np.uint8) # middle_show = np.concatenate((user_rgb_8uc3_orisize, user_res_8uc3_orisize, user_matting_mask_8uc3_orisize, user_res_8uc3_orisize2), axis=1) # ratio = 1536. / max(middle_show.shape[:2]) # middle_show = cv2.resize(middle_show, (0, 0), fx=ratio, fy=ratio) # cv2.imshow('middle_show', middle_show) # cv2.waitKey(0) # user_res_8uc3_orisize = np.concatenate((user_res_8uc3_orisize, user_matting_8uc3_bald_orisize), axis=1) return user_res_8uc3_orisize2, 0 def infer_haircolor_0313_basecolor(self, user_rgb_8uc3_orisize, haircolor_dir): landmarks_origin_img_1k = self.get_landmark.forward(user_rgb_8uc3_orisize) if landmarks_origin_img_1k is None: return None, 10001 _, user_matting_8uc1_bald_orisize = self.process_data.generator_matte.matte_inference(user_rgb_8uc3_orisize, landmarks_origin_img_1k) user_matting_mask_8uc3_orisize = np.repeat(user_matting_8uc1_bald_orisize[:, :, np.newaxis], 3, axis=2) ref_rgb_8uc3_change_color_768 = cv2.imread(os.path.join(haircolor_dir, "ref_rgb_8uc3_color_768.png")) ref_matting_8uc3_change_color_768 = cv2.imread(os.path.join(haircolor_dir, "ref_matting_8uc3_color_768.png")) user_rgb_8uc3_change_color_768, user_matting_8uc3_change_color_768, user_hair_color_M = \ self.process_data.get_prepare_hair_color_user_data(user_rgb_8uc3_orisize, landmarks_origin_img_1k) haircolor_dir_basecolor = "/home/yangchaojie/Desktop/hairstyle/hairstyle_infer/data/ref_haircolor/HDR10_443322" ref_rgb_8uc3_change_color_768_basecolor = cv2.imread(os.path.join(haircolor_dir_basecolor, "ref_rgb_8uc3_color_768.png")) ref_matting_8uc3_change_color_768_basecolor = cv2.imread(os.path.join(haircolor_dir_basecolor, "ref_matting_8uc3_color_768.png")) # start_time = time.time() hair_gene_color_8uc3_768, user_matting_mask_8uc3_768 = self.change_haircolor.Change_Hair_inference( user_rgb_8uc3_change_color_768, user_matting_8uc3_change_color_768, ref_rgb_8uc3_change_color_768_basecolor, ref_matting_8uc3_change_color_768_basecolor) hair_gene_color_8uc3_orisize = cv2.warpAffine(hair_gene_color_8uc3_768, cv2.invertAffineTransform(user_hair_color_M), (user_rgb_8uc3_orisize.shape[1], user_rgb_8uc3_orisize.shape[0]), dst=user_rgb_8uc3_orisize.copy(), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_TRANSPARENT) # if self.use_enhance: user_res_8uc3_orisize = self.face_enhance.process(hair_gene_color_8uc3_orisize, landmarks_origin_img_1k) #### version 2 user_res_fc32_orisize = user_res_8uc3_orisize.astype(np.float32) / 255 # user_res_fc32_orisize_hsv = cv2.cvtColor(user_res_fc32_orisize, cv2.COLOR_BGR2HSV) ref_rgb_fc32_change_color_768 = ref_rgb_8uc3_change_color_768.astype(np.float32) / 255 ref_matting_fc32_change_color_768 = ref_matting_8uc3_change_color_768.astype(np.float32) / 255 # ref_rgb_fc32_change_color_768_hsv = cv2.cvtColor(ref_rgb_fc32_change_color_768, cv2.COLOR_BGR2HSV) # ref_rgb_fc32_change_color_768_lab = cv2.cvtColor(ref_rgb_fc32_change_color_768, cv2.COLOR_BGR2LAB) # ref_avg_std = self.getavgstd(ref_rgb_fc32_change_color_768_lab, ref_matting_fc32_change_color_768) # # user_res_fc32_orisize_lab = cv2.cvtColor(user_res_fc32_orisize, cv2.COLOR_BGR2LAB) # # user_avg_std = self.getavgstd(user_res_fc32_orisize_lab, user_matting_mask_8uc3_orisize) # user_matting_mask_fc32_orisize = user_matting_mask_8uc3_orisize.astype(np.float32) / 255 # user_res_fc32_orisize, src_avg_std = self.reinhard(user_res_fc32_orisize, user_matting_mask_fc32_orisize, ref_avg_std, ratio=0.99) ref_rgb_avg_std = self.getavgstd(ref_rgb_fc32_change_color_768, ref_matting_fc32_change_color_768) # user_rgb_avg_std = self.getavgstd(user_res_fc32_orisize, user_matting_mask_8uc3_orisize) user_matting_mask_fc32_orisize = user_matting_mask_8uc3_orisize.astype(np.float32) / 255 user_res_fc32_orisize, src_avg_std = self.reinhard_rgb(user_res_fc32_orisize, user_matting_mask_fc32_orisize, ref_rgb_avg_std, ratio=0.95) user_res_8uc3_orisize = (user_res_fc32_orisize * 255).astype(np.uint8) #### version 2 user_res_8uc3_orisize2 = user_rgb_8uc3_orisize * (1 - user_matting_mask_8uc3_orisize / 255) + user_res_8uc3_orisize * (user_matting_mask_8uc3_orisize / 255) user_res_8uc3_orisize2 = (np.clip(user_res_8uc3_orisize2, 0, 255)).astype(np.uint8) # middle_show = np.concatenate((user_rgb_8uc3_orisize, user_res_8uc3_orisize, user_matting_mask_8uc3_orisize, user_res_8uc3_orisize2), axis=1) # ratio = 1536. / max(middle_show.shape[:2]) # middle_show = cv2.resize(middle_show, (0, 0), fx=ratio, fy=ratio) # cv2.imshow('middle_show', middle_show) # cv2.waitKey(0) # user_res_8uc3_orisize = np.concatenate((user_res_8uc3_orisize, user_matting_8uc3_bald_orisize), axis=1) return user_res_8uc3_orisize2, 0 def infer_haircolor_0313_rgb(self, user_rgb_8uc3_orisize, haircolor_dir): landmarks_origin_img_1k = self.get_landmark.forward(user_rgb_8uc3_orisize) if landmarks_origin_img_1k is None: return None, 10001 _, user_matting_8uc1_bald_orisize = self.process_data.generator_matte.matte_inference(user_rgb_8uc3_orisize, landmarks_origin_img_1k) user_matting_mask_8uc3_orisize = np.repeat(user_matting_8uc1_bald_orisize[:, :, np.newaxis], 3, axis=2) ref_rgb_8uc3_change_color_768 = cv2.imread(os.path.join(haircolor_dir, "ref_rgb_8uc3_color_768.png")) ref_matting_8uc3_change_color_768 = cv2.imread(os.path.join(haircolor_dir, "ref_matting_8uc3_color_768.png")) # if self.use_enhance: user_res_8uc3_orisize = user_rgb_8uc3_orisize.copy() #### version 2 user_res_fc32_orisize = user_res_8uc3_orisize.astype(np.float32) / 255 # user_res_fc32_orisize_hsv = cv2.cvtColor(user_res_fc32_orisize, cv2.COLOR_BGR2HSV) ref_rgb_fc32_change_color_768 = ref_rgb_8uc3_change_color_768.astype(np.float32) / 255 ref_matting_fc32_change_color_768 = ref_matting_8uc3_change_color_768.astype(np.float32) / 255 # ref_rgb_fc32_change_color_768_hsv = cv2.cvtColor(ref_rgb_fc32_change_color_768, cv2.COLOR_BGR2HSV) # ref_rgb_fc32_change_color_768_lab = cv2.cvtColor(ref_rgb_fc32_change_color_768, cv2.COLOR_BGR2LAB) # ref_avg_std = self.getavgstd(ref_rgb_fc32_change_color_768_lab, ref_matting_fc32_change_color_768) # # user_res_fc32_orisize_lab = cv2.cvtColor(user_res_fc32_orisize, cv2.COLOR_BGR2LAB) # # user_avg_std = self.getavgstd(user_res_fc32_orisize_lab, user_matting_mask_8uc3_orisize) # user_matting_mask_fc32_orisize = user_matting_mask_8uc3_orisize.astype(np.float32) / 255 # user_res_fc32_orisize, src_avg_std = self.reinhard(user_res_fc32_orisize, user_matting_mask_fc32_orisize, ref_avg_std, ratio=0.99) ref_rgb_avg_std = self.getavgstd(ref_rgb_fc32_change_color_768, ref_matting_fc32_change_color_768) # user_rgb_avg_std = self.getavgstd(user_res_fc32_orisize, user_matting_mask_8uc3_orisize) user_matting_mask_fc32_orisize = user_matting_mask_8uc3_orisize.astype(np.float32) / 255 user_res_fc32_orisize, src_avg_std = self.reinhard_rgb(user_res_fc32_orisize, user_matting_mask_fc32_orisize, ref_rgb_avg_std, ratio=1.0) user_res_8uc3_orisize = (user_res_fc32_orisize * 255).astype(np.uint8) #### version 2 # user_res_8uc3_orisize_LAB = cv2.cvtColor(user_res_8uc3_orisize, cv2.COLOR_BGR2LAB) # user_rgb_8uc3_orisize_LAB = cv2.cvtColor(user_rgb_8uc3_orisize, cv2.COLOR_BGR2LAB) # mix_ratio = 0.5 # user_res_8uc3_orisize_LAB[:, :, 0] = user_rgb_8uc3_orisize_LAB[:, :, 0] * mix_ratio + user_res_8uc3_orisize_LAB[:, :, 0] * (1 - mix_ratio) # user_res_8uc3_orisize = cv2.cvtColor(user_res_8uc3_orisize_LAB, cv2.COLOR_LAB2BGR) user_res_8uc3_orisize2 = user_rgb_8uc3_orisize * (1 - user_matting_mask_8uc3_orisize / 255) + user_res_8uc3_orisize * (user_matting_mask_8uc3_orisize / 255) user_res_8uc3_orisize2 = (np.clip(user_res_8uc3_orisize2, 0, 255)).astype(np.uint8) # middle_show = np.concatenate((user_rgb_8uc3_orisize, user_res_8uc3_orisize, user_matting_mask_8uc3_orisize, user_res_8uc3_orisize2), axis=1) # ratio = 1536. / max(middle_show.shape[:2]) # middle_show = cv2.resize(middle_show, (0, 0), fx=ratio, fy=ratio) # cv2.imshow('middle_show', middle_show) # cv2.waitKey(0) # user_res_8uc3_orisize = np.concatenate((user_res_8uc3_orisize, user_matting_8uc3_bald_orisize), axis=1) return user_res_8uc3_orisize2, 0 def infer_bald(self, user_rgb_8uc3_orisize): landmarks_origin_img_1k = self.get_landmark.forward(user_rgb_8uc3_orisize) if landmarks_origin_img_1k is None: return None, 10001 # 返回光头 user_res_8uc3_orisize, user_inter_res_8uc3_orisize, user_matting_8uc3_bald_768, user_baldseg_8uc3_bald_768 = self.process_data.get_user_blad( user_rgb_8uc3_orisize, landmarks_origin_img_1k) if self.use_enhance: user_res_8uc3_orisize = self.face_enhance.process(user_res_8uc3_orisize, landmarks_origin_img_1k) return user_res_8uc3_orisize, 0 def buff(self, img, img_skin, value1, value2): img = img.astype(np.float32) dx = value1 * 5 fc = value1 * 12.5 p = 80 temp1 = cv2.bilateralFilter(img, dx, fc, fc) temp2 = (temp1 - img + 128) temp2 = np.clip(temp2, 0, 255) temp3 = cv2.GaussianBlur(temp2, (2 * value2 - 1, 2 * value2 - 1), 0, 0) temp4 = img + 2 * temp3 - 255 temp4 = np.clip(temp4, 0, 255) dst = img * ((100 - p) / 100) + temp4 * (p / 100) img_skin_c = 1-img_skin dst = dst * img_skin + img * img_skin_c # mid_show = np.concatenate((img, img_skin, dst), axis=1) # ratio = 1224. / max(mid_show.shape[:2]) # mid_show = cv2.resize(mid_show, (0, 0), fx=ratio, fy=ratio) # cv2.imshow('mid_show', mid_show) # cv2.waitKey(0) return dst.astype(np.uint8) def whitening(self, img, img_skin, value): midtones_add = np.zeros(256) for i in range(256): midtones_add[i] = 0.667 * (1 - ((i - 127) / 127) * ((i - 127) / 127)) lookup = np.zeros(256, dtype='uint8') for i in range(256): red = i red += value * midtones_add[red] red = max(0, red) lookup[i] = np.uint(red) w, h, c = img.shape img_skin = img_skin[:,:,-1] index = np.where(img_skin == 1) for i in range(index[0].shape[0]): img[index[0][i], index[1][i], 0] = lookup[int(img[index[0][i], index[1][i], 0])] img[index[0][i], index[1][i], 1] = lookup[int(img[index[0][i], index[1][i], 1])] img[index[0][i], index[1][i], 2] = lookup[int(img[index[0][i], index[1][i], 2])] # for i in range(w): # for j in range(h): # if img_skin[i, j, 0] == 1: # img[i, j, 0] = lookup[img[i, j, 0]] # img[i, j, 1] = lookup[img[i, j, 1]] # img[i, j, 2] = lookup[img[i, j, 2]] return img def infer_face(self, origin_img): landmarks_origin_img_1k = self.get_landmark.forward(origin_img) if landmarks_origin_img_1k is None: return None landmark_137 = landmark_processor.pts_1k_to_137(landmarks_origin_img_1k).astype(np.int32) face_mask_ori = self.face_seg.inference(origin_img, landmarks_origin_img_1k) # if get_facecolor: h, w, _ = origin_img.shape inpaint_mask = np.zeros((h, w), dtype=np.uint8) cv2.fillConvexPoly(inpaint_mask, cv2.convexHull(landmark_137[129:137]), (255,)) cv2.fillConvexPoly(inpaint_mask, cv2.convexHull(landmark_137[121:129]), (255,)) cv2.fillConvexPoly(inpaint_mask, cv2.convexHull(landmark_137[22:48]), (255,)) cv2.fillConvexPoly(inpaint_mask, cv2.convexHull(landmark_137[88:104]), (255,)) cv2.fillConvexPoly(inpaint_mask, cv2.convexHull(landmark_137[105:121]), (255,)) inpaint_mask = cv2.cvtColor(inpaint_mask, cv2.COLOR_GRAY2BGR) face_mask_del = np.clip(face_mask_ori - (inpaint_mask / 255).astype(np.float32), 0, 1) # face_image_del = (origin_img*face_mask_del).astype(np.uint8) # face_mask_del = np.around(face_mask_del) dst_img_white = self.whitening(origin_img, face_mask_del, 10) # cv2.imwrite("/home/yangchaojie/Desktop/hairstyle/hairstyle_infer/test_data/test_face_white.jpg", dst_img_white) # whiten_show = np.concatenate((origin_img, dst_img_white), axis=1) # ratio = 1024. / max(whiten_show.shape[:2]) # whiten_show = cv2.resize(whiten_show, (0, 0), fx=ratio, fy=ratio) # cv2.imshow('whiten_show', whiten_show) # cv2.waitKey(0) # bounding_boxs = bounding_boxs.astype(np.int32) face_left = landmark_137[16] face_right = landmark_137[6] face_top = landmark_137[11] face_bottom = landmark_137[0] x1, y1, x2, y2 = max(face_left[0] - 20, 0), max(face_top[1] - 20, 0), min(face_right[0] + 20, w), min( face_bottom[1] + 20, h) img_beauty = dst_img_white.copy() dst_img_buff = self.buff(dst_img_white[y1:y2, x1:x2], face_mask_del[y1:y2, x1:x2], 4, 3) img_beauty[y1:y2, x1:x2] = dst_img_buff # cv2.imwrite("/home/yangchaojie/Desktop/hairstyle/hairstyle_infer/test_data/test_face_buff.jpg", img_beauty) # buff_show = np.concatenate((origin_img, dst_img_white, img_beauty), axis=1) # ratio = 1224. / max(buff_show.shape[:2]) # buff_show = cv2.resize(buff_show, (0, 0), fx=ratio, fy=ratio) # cv2.imshow('buff_show', buff_show) # cv2.waitKey(0) return img_beauty if __name__=='__main__': hf = HairStyle_Model_Infer() # user_img = cv2.imread() # ref_color = # hf.infer_haircolor_new()