xsl/change_hair
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
0926c61bd0cedd6165d70e237c20d898613a9b09
change_hair/project/hair_service_sd/hairstyle_model_infer.py
T
xslandCursor 0926c61bd0 fix: 适配 ubuntu 路径并增强换发色/训练流程稳定性 
将配置与训练脚本从 /home/xsl 切到本机 /home/ubuntu；换发色在 webui 增强失败或缺色板时降级返回，训练结束后自动重启 hair 服务再回调。

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-07-19 00:49:16 +08:00

1677 lines
99 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
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)
if status != 0:
return user_rgb_8uc3_orisize, None, status
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, 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"))
if os.path.exists(haircolor_dir_tmp):
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
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()
Reference in New Issue View Git Blame Copy Permalink