import os import time import torchvision import pickle import sys import torch from torch.nn import functional as F from torch import nn import numpy as np import cv2 import math # import onnxruntime from utils import landmark_processor # import torchvision from models.MomocvFaceAlignment1K import MomocvFaceAlignment1K from mtcnn.detector import MTCNNFaceDetector from utils import util from matting.networks import generators from seg.hairseg_single_model import Evaluator from models.Generator_BaldSeg import Generator_BaldSeg_5c from bodyseg.msc_distilling import DeepLab from utils import model_io modelRoot = "./weights" class Pose_KeypointsProcessorV2(object): def __init__(self, gpu_id=0): import torch sys.path.append("..") from keypoints.lib.config import cfg, update_config from keypoints.lib.models.pose_hrnet import get_pose_net sys.path.remove("..") class Namespace: def __init__(self, **kwargs): self.__dict__.update(kwargs) args = Namespace(cfg=os.path.join(modelRoot, 'keypoints/experiments/coco/hrnet/coco25_384x288_adam_lr1e-3.yaml'), opts=['TEST.MODEL_FILE', os.path.join(modelRoot, 'kpnts_detect_lyq_20200720.pth'), 'TEST.USE_GT_BBOX', 'False'], dataDir='', logDir='', modelDir='', prevModelDir='') update_config(cfg, args) # print(cfg) self.model = get_pose_net(cfg, is_train=False) self.model.eval() if cfg.TEST.MODEL_FILE: if gpu_id == -1: self.model.load_state_dict(torch.load(cfg.TEST.MODEL_FILE, map_location='cpu'), strict=False) else: # self.model.load_state_dict(cp['best_state_dict']) self.model.load_state_dict(torch.load(cfg.TEST.MODEL_FILE, map_location=lambda storage, loc: storage), strict=False) print('load keypoints model weights') if gpu_id != -1: self.model.cuda(gpu_id) self.gpu_id = gpu_id self.image_size = [288, 384] self.cfg = cfg self.compare_table = [[0, 0], [1, 16], [2, 15], [3, 18], [4, 17], \ [5, 5], [6, 2], [7, 6], [8, 3], [9, 7], \ [10, 4], [11, 12], [12, 9], [13, 13], [14, 10], \ [15, 14], [16, 11]] self.ref_lds = [[142.50674, 55.46675], [142.58651, 93.27755], [115.88326, 95.49268], [107.85087, 135.16175], [103.58789, 169.31675], [169.52114, 89.90296], [186.71601, 131.36675], [197.37347, 167.41925], [146.29747, 175.92391], [125.37664, 174.99589], [141.95472, 241.42175], [150.48068, 311.62925], [163.45916, 176.62695], [161.13813, 249.01175], [156.87515, 296.80147], [135.56025, 49.77425], [150.48068, 47.87675], [124.90280, 55.46675], [156.87515, 51.67175], [150.12675, 322.52566], [155.76106, 320.27593], [158.85028, 295.46357], [154.74366, 336.29675], [146.21770, 336.29675], [147.24197, 316.33891]] self.ref_lds = np.array(self.ref_lds) self.compare_table = np.array(self.compare_table) self.ref_lds = self.ref_lds[self.compare_table[:, 1]] # @staticmethod def get_refimage(self, srcImg, landmark, refLandmark): def setInvM(M): # return inv M D = M[0, 0] * M[1, 1] - M[0, 1] * M[1, 0] D = 1 / D if D != 0 else 0.0 invM = np.zeros((2, 3)) invM[0, 0] = M[1, 1] * D invM[0, 1] = M[0, 1] * D * (-1) invM[1, 0] = M[1, 0] * D * (-1) invM[1, 1] = M[0, 0] * D invM[0, 2] = -invM[0, 0] * M[0, 2] - invM[0, 1] * M[1, 2] invM[1, 2] = -invM[1, 0] * M[0, 2] - invM[1, 1] * M[1, 2] return invM def setTransMatrix(srcPts, dstPts): ms_x = np.mean(srcPts[:, 0]) ms_y = np.mean(srcPts[:, 1]) srcPts -= np.array([ms_x, ms_y]) x = np.concatenate((srcPts[:, 0], srcPts[:, 1])) y = np.concatenate((dstPts[:, 0], dstPts[:, 1])) a = y.dot(x) / x.dot(x) offset = int(len(x) / 2) b = 0 for i in range(offset): b += x[i] * y[offset + i] - x[offset + i] * y[i] b /= x.dot(x) md_x = np.mean(dstPts[:, 0]) md_y = np.mean(dstPts[:, 1]) M = np.zeros((2, 3)) M[0, 0] = a M[0, 1] = -b M[1, 0] = b M[1, 1] = a M[0, 2] = md_x - M[0, 0] * ms_x - M[0, 1] * ms_y M[1, 2] = md_y - M[1, 0] * ms_x - M[1, 1] * ms_y invM = setInvM(M) return invM, M def transImage(srcImg, M, invM, back=False): T = M if back else invM dstImg = cv2.warpAffine(srcImg, T, (288, 384)) return dstImg def transShape(srcShape, M, invM, forward=False): T = invM if forward else M x = T[0, 0] * srcShape[:, 0] + T[0, 1] * srcShape[:, 1] + T[0, 2] y = T[1, 0] * srcShape[:, 1] + T[1, 1] * srcShape[:, 1] + T[1, 2] # print(x, y) tmp = zip(x, y) return tmp landmark_tmp = landmark.copy() M, invM = setTransMatrix(landmark_tmp, refLandmark) dstImg = transImage(srcImg, M, invM) return dstImg, invM # @staticmethod def transform_points(self, points, mat, invert=False): if invert: mat = cv2.invertAffineTransform(mat) points = np.expand_dims(points, axis=1) points = cv2.transform(points, mat, points.shape) points = np.squeeze(points) return points # @staticmethod def affine_trans(self, img, joints): img1 = img.copy() save_pts = joints.copy() mask = [] obj_lds = np.array(joints[:, :2]).copy() tag = 0 for i in joints[:, 2]: if i > 0.6: mask.append(True) tag += 1 else: mask.append(False) if tag < 8: mask = [False, True, True, False, False, True, True, False, False, False, False, False, False, False, False, False, False] ref_lds = self.ref_lds.copy() ref_lds = np.array(ref_lds) mask = np.array(mask) ref_lds = ref_lds[mask] obj_lds = obj_lds[mask] img, M = self.get_refimage(img, obj_lds, ref_lds) return img, M def get_max_preds(self, batch_heatmaps): ''' get predictions from score maps heatmaps: numpy.ndarray([batch_size, num_joints, height, width]) ''' assert isinstance(batch_heatmaps, np.ndarray), \ 'batch_heatmaps should be numpy.ndarray' assert batch_heatmaps.ndim == 4, 'batch_images should be 4-ndim' batch_size = batch_heatmaps.shape[0] num_joints = batch_heatmaps.shape[1] width = batch_heatmaps.shape[3] heatmaps_reshaped = batch_heatmaps.reshape((batch_size, num_joints, -1)) idx = np.argmax(heatmaps_reshaped, 2) maxvals = np.amax(heatmaps_reshaped, 2) maxvals = maxvals.reshape((batch_size, num_joints, 1)) idx = idx.reshape((batch_size, num_joints, 1)) preds = np.tile(idx, (1, 1, 2)).astype(np.float32) preds[:, :, 0] = (preds[:, :, 0]) % width preds[:, :, 1] = np.floor((preds[:, :, 1]) / width) pred_mask = np.tile(np.greater(maxvals, 0.0), (1, 1, 2)) pred_mask = pred_mask.astype(np.float32) preds *= pred_mask return preds, maxvals def preprocess(self, img, pre_pts): input, M = self.affine_trans(img, pre_pts) # print(input1.shape) input = cv2.cvtColor(input, cv2.COLOR_BGR2RGB) input = np.divide(np.subtract(input.astype(np.float32) / 255, [0.406, 0.456, 0.485]), [0.225, 0.224, 0.229]) input = np.expand_dims(input.astype(np.float32).transpose((2, 0, 1)), axis=0) input = torch.from_numpy(input) return input, M def postprocess(self, output, M): output = output.detach().cpu().numpy() lds, maxvals = self.get_max_preds(output) lds *= 4 # print(lds, lds.shape) lds = lds[0] lds = self.transform_points(lds[:, :2], M, True) res = np.concatenate((lds, maxvals[0]), axis=1) # print(res) return res def forward(self, img, pre_lds): ss = img.copy() img, M = self.preprocess(img, pre_lds) if self.gpu_id != -1: img = img.cuda(self.gpu_id) output = self.model(img) pts = self.postprocess(output, M) return pts def pt_conv_25_to_17(pt25): index_map = [ 0, # 0 16, # 1 15, # 2 18, # 3 17, # 4 5, # 5 2, # 6 6, # 7 3, # 8 7, # 9 4, # 10 12, # 11 9, # 12 13, # 13 10, # 14 14, # 15 11, # 16 ] pt_17 = np.zeros((17, 3), dtype=np.float32) for idx17, idx25 in enumerate(index_map): pt_17[idx17, :] = pt25[idx25] return pt_17 class Human_Keypoints: def __init__(self, gpu=True, device_id=0): if gpu: self.gpu_id = device_id else: self.gpu_id = -1 with torch.no_grad(): # self.pose_keypoints_processor = Pose_KeypointsProcessor(gpu_id=self.gpu_id) self.pose_keypoints_processor2 = Pose_KeypointsProcessorV2(gpu_id=self.gpu_id) def inference(self, input_img, human_rect, kpnts17, hand_seg=None): # pose_keypoints = self.pose_keypoints_processor.forward(input_img, human_rect) pose_keypoints = self.pose_keypoints_processor2.forward(input_img, kpnts17) # new version lyq # # check output 25 kpnts # main_image_show = input_img.copy() # for i in range(pose_keypoints.shape[0]): # cv2.circle(main_image_show, (int(pose_keypoints[i][0]), int(pose_keypoints[i][1])), 5, (0, 255, 0), 5) # old green # cv2.circle(main_image_show, (int(pose_keypoints2[i][0]), int(pose_keypoints2[i][1])), 5, (220, 0, 220), 5) # new red # cv2.imshow("kpnts compare", main_image_show) # cv2.waitKey() if len(pose_keypoints) == 0: pose_keypoints = np.zeros((25, 3), dtype=np.float32) return pose_keypoints def convert_pose_points(self, pose_keypoints, confidence): if (pose_keypoints[11, 0] - pose_keypoints[10, 0]) ** 2 + (pose_keypoints[11, 1] - pose_keypoints[10, 1]) ** 2 < \ 0.09 * ((pose_keypoints[10, 0] - pose_keypoints[9, 0]) ** 2 + (pose_keypoints[10, 1] - pose_keypoints[9, 1]) ** 2): pose_keypoints[11, 2] = 0 if (pose_keypoints[14, 0] - pose_keypoints[13, 0]) ** 2 + (pose_keypoints[14, 1] - pose_keypoints[13, 1]) ** 2 < \ 0.09 * ((pose_keypoints[13, 0] - pose_keypoints[12, 0]) ** 2 + (pose_keypoints[13, 1] - pose_keypoints[12, 1]) ** 2): pose_keypoints[14, 2] = 0 if pose_keypoints[10, 2] <= confidence: pose_keypoints[11, 2] = 0 if pose_keypoints[13, 2] <= confidence: pose_keypoints[14, 2] = 0 ### heel if (pose_keypoints[24, 0] - pose_keypoints[11, 0]) ** 2 + (pose_keypoints[24, 1] - pose_keypoints[11, 1]) ** 2 > \ 0.2 * ((pose_keypoints[11, 0] - pose_keypoints[10, 0]) ** 2 + (pose_keypoints[11, 1] - pose_keypoints[10, 1]) ** 2): pose_keypoints[24, 2] = 0 if (pose_keypoints[21, 0] - pose_keypoints[14, 0]) ** 2 + (pose_keypoints[21, 1] - pose_keypoints[14, 1]) ** 2 > \ 0.2 * ((pose_keypoints[14, 0] - pose_keypoints[13, 0]) ** 2 + (pose_keypoints[14, 1] - pose_keypoints[13, 1]) ** 2): pose_keypoints[21, 2] = 0 ### foot keypoints if pose_keypoints[11, 2] <= confidence: pose_keypoints[22, 2] = 0 pose_keypoints[23, 2] = 0 pose_keypoints[24, 2] = 0 if pose_keypoints[14, 2] <= confidence: pose_keypoints[19, 2] = 0 pose_keypoints[20, 2] = 0 pose_keypoints[21, 2] = 0 return pose_keypoints def draw_pose_keypoints(self, img, pose_keypoints, confidence): if len(pose_keypoints) == 0: return img # self.convert_pose_points(pose_keypoints, confidence) connect_body = [[0, 15], [0, 16], [15, 17], [16, 18], [0, 1], [1, 2], [1, 5], [2, 3], [3, 4], [5, 6], [6, 7], [1, 8], [8, 9], [9, 10], [10, 11], [8, 12], [12, 13], [13, 14], [14, 21], [14, 19], [19, 20], [11, 24], [11, 22], [22, 23]] color_body = [(128, 245, 223), (245, 223, 113), (35, 243, 46), (67, 187, 255), (222, 33, 245), (192, 133, 45)] for i in range(len(pose_keypoints)): if pose_keypoints[i, 2] > confidence: cv2.circle(img, (int(pose_keypoints[i, 0]), int(pose_keypoints[i, 1])), 3, (0, 255, 0), 4) cv2.putText(img, str(i), (int(pose_keypoints[i, 0]), int(pose_keypoints[i, 1])), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 0, 0), 1) for i, con in enumerate(connect_body): pt_0 = con[0] pt_1 = con[1] if pose_keypoints[pt_0, 2] > confidence and pose_keypoints[pt_1, 2] > confidence: cv2.line(img, (int(pose_keypoints[pt_0, 0]), int(pose_keypoints[pt_0, 1])), (int(pose_keypoints[pt_1, 0]), int(pose_keypoints[pt_1, 1])), color_body[int(i / 4)], 3) return img def draw_single_hand_keypoints(self, img, pts, confidence): connect_hand = [[0, 1], [1, 2], [2, 3], [3, 4], [0, 5], [5, 6], [6, 7], [7, 8], [0, 9], [9, 10], [10, 11], [11, 12], [0, 13], [13, 14], [14, 15], [15, 16], [0, 17], [17, 18], [18, 19], [19, 20]] color_hand = [(128, 245, 223), (245, 223, 113), (35, 243, 46), (67, 187, 255), (222, 33, 245)] for i in range(len(pts)): if pts[i, 2] > confidence: cv2.circle(img, (int(pts[i, 0]), int(pts[i, 1])), 1, (0, 255, 0), 2) # cv2.putText(img, str(i), (int(pts[i, 0]), int(pts[i, 1])), cv2.FONT_HERSHEY_SIMPLEX, 0.3, (255, 0, 0), 1) for i, con in enumerate(connect_hand): pt_0 = con[0] pt_1 = con[1] if pts[pt_0, 2] > confidence and pts[pt_1, 2] > confidence: cv2.line(img, (int(pts[pt_0, 0]), int(pts[pt_0, 1])), (int(pts[pt_1, 0]), int(pts[pt_1, 1])), color_hand[int(i / 4)], 2) return img def draw_hand_keypoints(self, img, hands_list, hand_keypoints, confidence): if len(hand_keypoints) == 0: return img connect_hand = [[0, 1], [1, 2], [2, 3], [3, 4], [0, 5], [5, 6], [6, 7], [7, 8], [0, 9], [9, 10], [10, 11], [11, 12], [0, 13], [13, 14], [14, 15], [15, 16], [0, 17], [17, 18], [18, 19], [19, 20]] color_hand = [(128, 245, 223), (245, 223, 113), (35, 243, 46), (67, 187, 255), (222, 33, 245)] for single_hand_keypoints in hand_keypoints: pts = single_hand_keypoints["hand_keypoints"] for i in range(len(pts)): if pts[i, 2] > confidence: cv2.circle(img, (int(pts[i, 0]), int(pts[i, 1])), 1, (0, 255, 0), 2) # cv2.putText(img, str(i), (int(pts[i, 0]), int(pts[i, 1])), cv2.FONT_HERSHEY_SIMPLEX, 0.3, (255, 0, 0), 1) for i, con in enumerate(connect_hand): pt_0 = con[0] pt_1 = con[1] if pts[pt_0, 2] > confidence and pts[pt_1, 2] > confidence: cv2.line(img, (int(pts[pt_0, 0]), int(pts[pt_0, 1])), (int(pts[pt_1, 0]), int(pts[pt_1, 1])), color_hand[int(i / 4)], 2) for x, y, w, is_left in hands_list: cv2.rectangle(img, (x, y), (x+w, y+w), (255, 0, 255), 1) return img def draw_face_keypoints(self, img, face_keypoints): if len(face_keypoints) == 0: return img pts = face_keypoints for i in range(len(pts)): cv2.circle(img, (int(pts[i, 0]), int(pts[i, 1])), 1, (0, 255, 0), 1) ### draw face_outline for i in range(0, 16): pt_0 = i pt_1 = i + 1 cv2.line(img, (int(pts[pt_0, 0]), int(pts[pt_0, 1])), (int(pts[pt_1, 0]), int(pts[pt_1, 1])), (0, 255, 0), 1) ### draw left eye_brow for i in range(17, 21): pt_0 = i pt_1 = i + 1 cv2.line(img, (int(pts[pt_0, 0]), int(pts[pt_0, 1])), (int(pts[pt_1, 0]), int(pts[pt_1, 1])), (0, 255, 0), 1) ### draw right eye_brow for i in range(22, 26): pt_0 = i pt_1 = i + 1 cv2.line(img, (int(pts[pt_0, 0]), int(pts[pt_0, 1])), (int(pts[pt_1, 0]), int(pts[pt_1, 1])), (0, 255, 0), 1) ### draw right eye for i in range(42, 47): pt_0 = i pt_1 = i + 1 cv2.line(img, (int(pts[pt_0, 0]), int(pts[pt_0, 1])), (int(pts[pt_1, 0]), int(pts[pt_1, 1])), (0, 255, 0), 1) ### draw left eye for i in range(42, 47): pt_0 = i pt_1 = i + 1 cv2.line(img, (int(pts[pt_0, 0]), int(pts[pt_0, 1])), (int(pts[pt_1, 0]), int(pts[pt_1, 1])), (0, 255, 0), 1) ### draw left eye for i in range(36, 41): pt_0 = i pt_1 = i + 1 cv2.line(img, (int(pts[pt_0, 0]), int(pts[pt_0, 1])), (int(pts[pt_1, 0]), int(pts[pt_1, 1])), (0, 255, 0), 1) ### draw nose line for i in range(27, 30): pt_0 = i pt_1 = i + 1 cv2.line(img, (int(pts[pt_0, 0]), int(pts[pt_0, 1])), (int(pts[pt_1, 0]), int(pts[pt_1, 1])), (0, 255, 0), 1) ### draw nose hole for i in range(31, 35): pt_0 = i pt_1 = i + 1 cv2.line(img, (int(pts[pt_0, 0]), int(pts[pt_0, 1])), (int(pts[pt_1, 0]), int(pts[pt_1, 1])), (0, 255, 0), 1) ### draw big mouth for i in range(48, 59): pt_0 = i pt_1 = i + 1 cv2.line(img, (int(pts[pt_0, 0]), int(pts[pt_0, 1])), (int(pts[pt_1, 0]), int(pts[pt_1, 1])), (0, 255, 0), 1) ### draw small mouth for i in range(60, 69): pt_0 = i pt_1 = i + 1 cv2.line(img, (int(pts[pt_0, 0]), int(pts[pt_0, 1])), (int(pts[pt_1, 0]), int(pts[pt_1, 1])), (0, 255, 0), 1) # cv2.rectangle(img, (face_rect[0], face_rect[1]), (face_rect[2], face_rect[3]), (0, 255, 255), 1) return img class Get_Landmark(object): def __init__(self, gpu_id=0): self.img_size = 640 self.face_alignmenter_1k = MomocvFaceAlignment1K(gpu_id=gpu_id) self.face_detector = MTCNNFaceDetector(gpu_id=gpu_id) def get_max_rect(self, bounding_boxes): max_area = 0 index = 0 for i, box in enumerate(bounding_boxes): width = box[2] - box[0] height = box[3] - box[1] if width * height > max_area: index = i max_area = width * height return index def forward(self, img): # bounding_boxes, landmarks = self.face_detector.forward(img, min_face_size=50, # thresholds=[0.6, 0.7, 0.9]) bounding_boxes, landmarks = self.face_detector.forward(img, min_face_size=30, thresholds=[0.6, 0.6, 0.6]) if len(bounding_boxes) == 0: return None if len(bounding_boxes) > 0: box_index = self.get_max_rect(bounding_boxes) pts5 = landmarks[box_index] else: return None landmarks1k = self.face_alignmenter_1k.detect_according_5pts(img, pts5) return landmarks1k def forward_infer(self, img): bounding_boxes, landmarks = self.face_detector.forward(img, min_face_size=50) if len(bounding_boxes) == 0: return None, None, None if len(bounding_boxes) >= 1: box_index = self.get_max_rect(bounding_boxes) pts5 = landmarks[box_index] else: return None, None, None landmarks1k = self.face_alignmenter_1k.detect_according_5pts(img, pts5) return landmarks1k def forward_color(self, img): bounding_boxes, landmarks = self.face_detector.forward(img, min_face_size=50) if len(bounding_boxes) == 0: return None if len(bounding_boxes) >= 1: box_index = self.get_max_rect(bounding_boxes) pts5 = landmarks[box_index] else: return None landmarks1k = self.face_alignmenter_1k.detect_according_5pts(img, pts5) return landmarks1k def forward_infer(self, img): bounding_boxes, landmarks = self.face_detector.forward(img, min_face_size=50) if len(bounding_boxes) == 0: return None, None, None if len(bounding_boxes) >= 1: box_index = self.get_max_rect(bounding_boxes) pts5 = landmarks[box_index] else: return None, None, None landmarks1k = self.face_alignmenter_1k.detect_according_5pts(img, pts5) return landmarks1k def forward_diy(self, img): # bounding_boxes, landmarks = self.face_detector.forward(img, min_face_size=100,thresholds=[0.6, 0.8, 0.8]) # mtcnn face detect input bounding_boxes, landmarks = self.face_detector.forward(img, min_face_size=50) if len(bounding_boxes) == 0: return None, None, None if len(bounding_boxes) >=1 : box_index = self.get_max_rect(bounding_boxes) pts5 = landmarks[box_index] else: return None, None, None # for i in range(5): # cv2.circle(img, (int(pts5[i*2]), int(pts5[i*2+1])), 1, (255, 0,0), 1) # cv2.imshow('img', img) # cv2.waitKey() landmarks1k = self.face_alignmenter_1k.detect_according_5pts(img, pts5) # pts137 = landmark_processor.pts_1k_to_137(landmarks1k) # movie_params = self.model_3d.detect([img], [pts137])[0] # pitch, yaw, roll = movie_params[1:4] ret_info = [landmarks1k, bounding_boxes[box_index], None] return ret_info def forward_v2(self, img): # bounding_boxes, landmarks = self.face_detector.forward(img, min_face_size=50, # thresholds=[0.6, 0.7, 0.9]) bounding_boxes, landmarks = self.face_detector.forward_v2(img, min_face_size=30, thresholds=[0.6, 0.6, 0.6]) if len(bounding_boxes) == 0: return None if len(bounding_boxes) > 0: box_index = self.get_max_rect(bounding_boxes) pts5 = landmarks[box_index] else: return None landmarks1k = self.face_alignmenter_1k.detect_according_5pts(img, pts5) return landmarks1k def get_face_shape(self, xiaba_type, landmark137): face_shape = ['chang', 'fang', 'yuan', 'tuoyuan', 'xin'] forhead_idxs = [12, 13, 14] mid_idxs = [15, 16, 17] lowwer_idxs = [18, 19, 20] face_width = max(landmark137[:,0]) - min(landmark137[:,0]) face_widthest_idx = list(landmark137[:,0]).index(min(landmark137[:,0])) face_scale_lw = 1.4 face_scale_qx_top = 1.4 face_scale_qx_bot = 1.2 face_height = landmark137[0][1] - landmark137[11][1] face_width_qiane = landmark137[9][0] - landmark137[13][0] face_width_xiahe = landmark137[3][0] - landmark137[19][0] check_scale = face_width_qiane / face_width_xiahe print('face_height', face_height) print('face_width', face_width) print('face_scale_lw', face_scale_lw) print('face_width_qiane', face_width_qiane) print('face_width_xiahe', face_width_xiahe) print('face_scale_qx', check_scale) if face_widthest_idx in forhead_idxs: if face_height / face_width >= face_scale_lw: return face_shape[0] else: if check_scale > face_scale_qx_bot and face_scale_qx_bot < face_scale_qx_top: if xiaba_type == 'jian': return face_shape[0] elif xiaba_type == 'fang': return face_shape[1] else: return face_shape[1] elif check_scale > face_scale_qx_top: if xiaba_type == 'jian': return face_shape[4] elif xiaba_type == 'fang': return face_shape[4] else: return face_shape[1] else: return face_shape[1] elif face_widthest_idx in mid_idxs: if face_height / face_width >= face_scale_lw: return face_shape[0] else: if check_scale > face_scale_qx_bot and face_scale_qx_bot < face_scale_qx_top: if xiaba_type == 'jian': return face_shape[3] elif xiaba_type == 'fang': return face_shape[1] else: return face_shape[2] elif check_scale > face_scale_qx_top: if xiaba_type == 'jian': return face_shape[3] elif xiaba_type == 'fang': return face_shape[3] else: return face_shape[2] else: return face_shape[1] else: if face_height / face_width >= face_scale_lw: return face_shape[0] else: if check_scale > face_scale_qx_bot and face_scale_qx_bot < face_scale_qx_top: if xiaba_type == 'jian': return face_shape[0] elif xiaba_type == 'fang': return face_shape[1] else: return face_shape[1] else: return face_shape[1] class GenTrimap(object): def __init__(self): self.erosion_kernels = [None] + [cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (size, size)) for size in range(1,30)] def __call__(self, alpha): fg_mask = np.zeros_like(alpha) bg_mask = np.zeros_like(alpha) fg_mask[alpha == 255] = 1 bg_mask[alpha == 0] = 1 fg_mask = fg_mask.astype(np.int32).astype(np.uint8) bg_mask = bg_mask.astype(np.int32).astype(np.uint8) fg_mask = cv2.erode(fg_mask, self.erosion_kernels[15]) bg_mask = cv2.erode(bg_mask, self.erosion_kernels[29]) trimap = np.ones_like(alpha, dtype=np.uint8) * 128 trimap[fg_mask == 1] = 255 trimap[bg_mask == 1] = 0 return trimap def single_inference(model, image_dict, device, return_offset=False): with torch.no_grad(): image, trimap = image_dict['image'], image_dict['trimap'] alpha_shape = image_dict['alpha_shape'] image = image.to(device) trimap = trimap.to(device) alpha_pred, info_dict = model(image, trimap) fg_pred = alpha_pred[:, :-1, :, :] alpha_pred = alpha_pred[:, -1, :, :].unsqueeze(1) trimap_argmax = trimap.argmax(dim=1, keepdim=True) alpha_pred[trimap_argmax == 2] = 1 alpha_pred[trimap_argmax == 0] = 0 h, w = alpha_shape test_fg_pred = fg_pred[0].detach().cpu().numpy().transpose(1, 2, 0)[:, :, ::-1] * 255 test_fg_pred = test_fg_pred.astype(np.uint8) test_fg_pred = test_fg_pred[32:h+32, 32:w+32] test_pred = alpha_pred[0, 0, ...].detach().cpu().numpy() * 255 test_pred = test_pred.astype(np.uint8) test_pred = test_pred[32:h+32, 32:w+32] # cv2.imshow('test_fg_pred', test_fg_pred) # cv2.imshow('test_pred', test_pred) # cv2.waitKey() if return_offset: short_side = h if h < w else w ratio = 512 / short_side offset_1 = util.flow_to_image(info_dict['offset_1'][0][0, ...].data.cpu().numpy()).astype(np.uint8) # write softmax_scale to offset image scale = info_dict['offset_1'][1].cpu() offset_1 = cv2.resize(offset_1, (int(w * ratio), int(h * ratio)), interpolation=cv2.INTER_NEAREST) text = 'unknown: {:.2f}, known: {:.2f}'.format(scale[-1, 0].item(), scale[-1,1].item()) offset_1 = cv2.putText(offset_1, text, (10, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.8, 0, thickness=2) offset_2 = util.flow_to_image(info_dict['offset_2'][0][0, ...].data.cpu().numpy()).astype(np.uint8) # write softmax_scale to offset image scale = info_dict['offset_2'][1].cpu() offset_2 = cv2.resize(offset_2, (int(w * ratio), int(h * ratio)), interpolation=cv2.INTER_NEAREST) text = 'unknown: {:.2f}, known: {:.2f}'.format(scale[-1, 0].item(), scale[-1,1].item()) offset_2 = cv2.putText(offset_2, text, (10, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.8, 0, thickness=2) return test_fg_pred, test_pred, (offset_1, offset_2) else: return test_fg_pred, test_pred, None def generator_tensor_dict(image, trimap): sample = {'image': image, 'trimap': trimap, 'alpha_shape': trimap.shape} # reshape h, w = sample["alpha_shape"] if h % 32 == 0 and w % 32 == 0: padded_image = np.pad(sample['image'], ((32, 32), (32, 32), (0, 0)), mode="reflect") padded_trimap = np.pad(sample['trimap'], ((32, 32), (32, 32)), mode="reflect") sample['image'] = padded_image sample['trimap'] = padded_trimap else: target_h = 32 * ((h - 1) // 32 + 1) target_w = 32 * ((w - 1) // 32 + 1) pad_h = target_h - h pad_w = target_w - w padded_image = np.pad(sample['image'], ((32, pad_h+32), (32, pad_w+32), (0, 0)), mode="reflect") padded_trimap = np.pad(sample['trimap'], ((32, pad_h+32), (32, pad_w+32)), mode="reflect") sample['image'] = padded_image sample['trimap'] = padded_trimap # ImageNet mean & std mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1) std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1) # convert GBR images to RGB image, trimap = sample['image'][:, :, ::-1], sample['trimap'] # swap color axis image = image.transpose((2, 0, 1)).astype(np.float32) trimap[trimap < 85] = 0 trimap[trimap >= 170] = 2 trimap[trimap >= 85] = 1 # normalize image image /= 255. # to tensor sample['image'], sample['trimap'] = torch.from_numpy(image), torch.from_numpy(trimap).to(torch.long) sample['image'] = sample['image'].sub_(mean).div_(std) sample['trimap'] = F.one_hot(sample['trimap'], num_classes=3).permute(2, 0, 1).float() # add first channel sample['image'], sample['trimap'] = sample['image'][None, ...], sample['trimap'][None, ...] return sample class Generator_Matte(object): def __init__(self, gpu, device_id): 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.output_img_size = 512 triseg_model_path = os.path.join(modelRoot, 'deeplabv3_hair512_360_0520_wl.pth') self.triseg_model = Evaluator(gpu_id=device_id, output_img_size=self.output_img_size, nclass=3, seg_model_path=triseg_model_path) hair_matte_model_path = os.path.join(modelRoot, 'gca-dist-fg-0430-latest_model.pth') # gca-dist-fg-0203-latest_model gca-dist-fg-0430-latest_model self.matte_model = self.load_hair_matte_model(hair_matte_model_path) self.gen_trimap = GenTrimap() def load_hair_matte_model(self, hair_matte_model_path): # build model model = generators.get_generator(encoder="resnet_gca_encoder_29", decoder="res_gca_decoder_22", num_class=4) # load checkpoint checkpoint = torch.load(hair_matte_model_path, map_location=lambda storage, loc: storage) model.load_state_dict(util.remove_prefix_state_dict(checkpoint['state_dict']), strict=True) model.to(self.device) # print("matte_model: ", model) # inference model = model.eval() return model def matte_inference(self, image, landmark1k): trimap = self.triseg_model.eval(image, landmark1k) ori_h, ori_w, _ = image.shape limit_size = 1600 if ori_h > limit_size or ori_w > limit_size: if ori_h > ori_w: new_tri_h = limit_size new_tri_w = int(ori_w * limit_size / ori_h) else: new_tri_w = limit_size new_tri_h = int(ori_h * limit_size / ori_w) image_resize = cv2.resize(image, (new_tri_w, new_tri_h), interpolation=cv2.INTER_CUBIC) trimap_resize = cv2.resize(trimap, (new_tri_w, new_tri_h), interpolation=cv2.INTER_NEAREST) else: image_resize = image trimap_resize = trimap trimap = self.gen_trimap(trimap_resize[:, :, 0]) # cv2.imwrite(os.path.join(args.output, image_name.replace(ext, "_trimap.png")), trimap) image_dict = generator_tensor_dict(image_resize, trimap) pred_fg, pred, offset = single_inference(self.matte_model, image_dict, device=self.device) pred_fg[trimap == 1] = image_resize[trimap == 1] if pred.shape[1] != image.shape[1] or pred.shape[0] != image.shape[0]: pred = cv2.resize(pred, (image.shape[1], image.shape[0])) return pred_fg, pred class Generator_Bald(object): def __init__(self, gpu, device_id): 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.model_dir = modelRoot generator_bald_model_path = os.path.join(self.model_dir, 'zxm_v7_gen_bald_add_changed_hair_pairdata_5seg_0611.pt') print("load generator_bald_model_path", generator_bald_model_path) self.load_generator_bald768_model(generator_bald_model_path) self.output_size = 768 def load_generator_bald768_model(self, generator_bald_model_path): self.generator_bald_model = torch.jit.load(generator_bald_model_path, map_location='cpu').to(self.device) def Geneator_Bald_inference(self, user_rgb_8uc3_bald_768, user_matting_8uc3_bald_768, user_baldseg_8uc3_bald_768, user_landmark_f1k2_bald_768): """ input： 图像尺寸基于： 人脸 512， 图像均为3通道 user_rgb_8uc3_bald_512: 输入图， size 512， uint8 （0-255） user_matting_8uc3_bald_512： 输入图 matting alpha 值， uint8 （0-255） user_baldseg_8uc3_bald_512： 输入图 光头分割， uint8 （0-255） user_landmark_f1k2_bald_512：输入图 关键点 1k*2 float32 output： bald_gene_8uc3_bald_512: 生成 光头图， uint8 （0-255） """ # cv2.imshow("user_rgb_8uc3_bald_768 ori", user_rgb_8uc3_bald_768) # cv2.imshow("user_matting_8uc3_bald_768", user_matting_8uc3_bald_768) # cv2.imshow("user_baldseg_8uc3_bald_768", user_baldseg_8uc3_bald_768) # user_baldseg_8uc3_bald_768_cp = user_baldseg_8uc3_bald_768.copy() # cood_y = int((user_landmark_f1k2_bald_768[214, 1] + user_landmark_f1k2_bald_768[99, 1]) // 2) # # random_int = 30 # random.randint(0, 30) # kernel = np.ones((random_int, 1), np.uint8) # user_mask_erode = cv2.erode(user_baldseg_8uc3_bald_768, kernel, iterations=1) # user_baldseg_8uc3_bald_768[:cood_y, :, :] = user_mask_erode[:cood_y, :, :] # bald_con = np.concatenate((user_baldseg_8uc3_bald_768_cp, user_baldseg_8uc3_bald_768), axis=1) # cv2.imshow("user_rgb_8uc3_bald_768", user_rgb_8uc3_bald_768) # cv2.imshow("user_matting_8uc3_bald_768", user_matting_8uc3_bald_768) # user_matting_8uc3_bald_768_nonzero = np.zeros_like(user_matting_8uc3_bald_768) # user_matting_8uc3_bald_768_nonzero[(user_matting_8uc3_bald_768 > 50).all(axis=2)] = 255 # # cv2.imshow("user_matting_8uc3_bald_768 >0", user_matting_8uc3_bald_768_nonzero) # cv2.imshow("bald_con", bald_con) # cv2.waitKey() # cv2.imwrite("/media/DATA_4T/project/hairstyle/分割测试/debug_line_res/user_rgb_8uc3_bald_768.png", user_rgb_8uc3_bald_768) # cv2.imwrite("/media/DATA_4T/project/hairstyle/分割测试/debug_line_res/user_matting_8uc3_bald_768.png", user_matting_8uc3_bald_768) # cv2.imwrite("/media/DATA_4T/project/hairstyle/分割测试/debug_line_res/user_baldseg_8uc3_bald_768.png", user_baldseg_8uc3_bald_768) # np.savetxt("/media/DATA_4T/project/hairstyle/分割测试/debug_line_res/user_landmark_f1k2_bald_768.txt", user_landmark_f1k2_bald_768) real_res = user_rgb_8uc3_bald_768.copy() user_landmark_f1k2_bald_768 = user_landmark_f1k2_bald_768.astype(np.int32) # kernel = np.ones((8, 8), np.uint8) kernel_1 = np.ones((20, 20), np.uint8) kernel_2 = np.ones((40, 40), np.uint8) user_matting_8uc3_bald_768_dilate_1 = cv2.dilate(user_matting_8uc3_bald_768[:, :, 0], kernel_1, iterations=1) user_matting_8uc3_bald_768_dilate_2 = cv2.dilate(user_matting_8uc3_bald_768[:, :, 0], kernel_2, iterations=1) user_matting_8uc3_bald_768_dilate_2[user_matting_8uc3_bald_768_dilate_2 > 0] = 255 b = user_baldseg_8uc3_bald_768[:, :, 0] > 125 g = user_baldseg_8uc3_bald_768[:, :, 1] > 0 r = user_baldseg_8uc3_bald_768[:, :, 2] > 125 user_matting_8uc3_bald_768_blur = user_matting_8uc3_bald_768_dilate_2.copy() # cv2.imshow("user_matting_8uc3_bald_768_blur noface", user_matting_8uc3_bald_768_blur) # TODO change noface mask # user_matting_8uc3_bald_768_blur[g] = 255 user_matting_8uc3_bald_768_blur = cv2.blur(user_matting_8uc3_bald_768_blur, (30, 30)) # user_matting_8uc3_bald_768_blur_2 = cv2.blur(user_matting_8uc3_bald_768_blur, (20, 20)) # user_matting_8uc3_bald_768_blur[(~b) & (~g) & (~r)] = user_matting_8uc3_bald_768_blur_2[(~b) & (~g) & (~r)] # cv2.imshow("user_matting_8uc3_bald_768_blur", user_matting_8uc3_bald_768_blur) # user_matting_8uc3_bald_768_dilate_1[(~b) * (~g)] = user_matting_8uc3_bald_768[:, :, 0][(~b) * (~g)] user_matting_8uc3_bald_768_dilate_1[(~b) * (~g) * (~r)] = user_matting_8uc3_bald_768[:, :, 0][(~b) * (~g) * (~r)] # user_matting_8uc3_bald_768_dilate_1[(~b) * g * r] = user_matting_8uc3_bald_768[:, :, 0][(~b) * g * r] user_rgb_8uc3_bald_768[user_matting_8uc3_bald_768_dilate_1.astype(bool)] = 255 ######################################## # cv2.fillPoly(user_baldseg_8uc3_bald_512, user_pts1k[:311][np.newaxis, :, :], (200, 175, 0)) # face cv2.fillPoly(user_baldseg_8uc3_bald_768, user_landmark_f1k2_bald_768[928:999][np.newaxis, :, :], (0, 125, 0)) # left eyebrow cv2.fillPoly(user_baldseg_8uc3_bald_768, user_landmark_f1k2_bald_768[856:927][np.newaxis, :, :], (125, 0, 0)) # right eyebrow cv2.fillPoly(user_baldseg_8uc3_bald_768, user_landmark_f1k2_bald_768[691:754][np.newaxis, :, :], (0, 0, 125)) # left eye cv2.fillPoly(user_baldseg_8uc3_bald_768, user_landmark_f1k2_bald_768[792:855][np.newaxis, :, :], (125, 125, 0)) # right eye cv2.fillPoly(user_baldseg_8uc3_bald_768, user_landmark_f1k2_bald_768[548:616][np.newaxis, :, :], (0, 125, 125)) # rose cv2.fillPoly(user_baldseg_8uc3_bald_768, user_landmark_f1k2_bald_768[312:547][np.newaxis, :, :], (125, 125, 125)) # mouth cv2.fillPoly(user_baldseg_8uc3_bald_768, user_landmark_f1k2_bald_768[655:690][np.newaxis, :, :], (0, 0, 175)) # left black_eyeball cv2.fillPoly(user_baldseg_8uc3_bald_768, user_landmark_f1k2_bald_768[756:791][np.newaxis, :, :], (125, 0, 125)) # right black_eyeball eyes_mask = np.zeros(user_baldseg_8uc3_bald_768[:, :, 0].shape) cv2.fillPoly(eyes_mask, user_landmark_f1k2_bald_768[691:754][np.newaxis, :, :], 1) # left eye cv2.fillPoly(eyes_mask, user_landmark_f1k2_bald_768[792:855][np.newaxis, :, :], 1) # right eye kernel = np.ones((30, 30), np.uint8) eyes_mask = cv2.dilate(eyes_mask, kernel, iterations=1) cv2.fillPoly(eyes_mask, user_landmark_f1k2_bald_768[312:467][np.newaxis, :, :], 1) # mouth user_rgb_8uc3_bald_768[eyes_mask.astype(bool)] = real_res[eyes_mask.astype(bool)] ######################### add by zxm # face_erode_mask = np.zeros(user_baldseg_8uc3_bald_768[:, :, 0].shape) # cv2.fillPoly(face_erode_mask, user_landmark_f1k2_bald_768[:311][np.newaxis, :, :], 1) # face # kernel = np.ones((60, 60), np.uint8) # face_erode_mask = cv2.erode(face_erode_mask, kernel, iterations=1) # user_rgb_8uc3_bald_768[face_erode_mask.astype(bool)] = real_res[face_erode_mask.astype(bool)] # user_rgb_8uc3_bald_768[user_matting_8uc3_bald_768_dilate_copy.astype(bool)] = 255 ######################### add by zxm # repeat # user_rgb_8uc3_bald_512[eyes_mask.astype(bool)] = real_res[eyes_mask.astype(bool)] # cv2.imshow("user_baldseg_8uc3_bald_768_condition", user_baldseg_8uc3_bald_768) # cv2.imshow("user_rgb_8uc3_bald_768_input_paf", user_rgb_8uc3_bald_768) # cv2.waitKey() # cv2.imwrite( # "/media/liyang/DATA1/test_hair/换状后台数据_测试/badcase_0205/光头badcase/bald_input_0128/bald_condition.png", # user_baldseg_8uc3_bald_768) # cv2.imwrite( # "/media/liyang/DATA1/test_hair/换状后台数据_测试/badcase_0205/光头badcase/bald_input_0128/inpu_paf.png", # user_rgb_8uc3_bald_768) # cv2.imshow("user_baldseg_8uc3_bald_768", user_baldseg_8uc3_bald_768) # cv2.imshow("user_rgb_8uc3_bald_768", user_rgb_8uc3_bald_768) # cv2.waitKey() condition = user_baldseg_8uc3_bald_768 # np.concatenate((user_baldseg_8uc3_bald_512, user_hair_mask), axis=2) condition = (condition.astype(np.float32) / 255).transpose((2, 0, 1))[np.newaxis, :, :, :] input_paf = (user_rgb_8uc3_bald_768.astype(np.float32) / 255).transpose((2, 0, 1))[np.newaxis, :, :, :] condition = torch.from_numpy(condition).to(self.device) input_paf = torch.from_numpy(input_paf).to(self.device) # ******************* forward *******************$ # test_fake, _, _ = model.preview(input_paf, condition) with torch.no_grad(): test_fake = self.generator_bald_model(input_paf, condition) test_res = (test_fake.detach()[0].to('cpu').numpy().transpose((1, 2, 0)) * 255).astype(np.uint8).copy() # show_concat = np.concatenate((user_rgb_8uc3_bald_768, user_baldseg_8uc3_bald_768, test_res), axis=1) # cv2.imshow("bald_concat", show_concat) # cv2.waitKey() # kernel = np.ones((25, 25), np.uint8) # user_hair_mask_dilate = cv2.dilate(user_matting_8uc3_bald_768[:, :, 0], kernel) # user_hair_mask_dilate[user_hair_mask_dilate > 0] = 255 # user_hair_mask_blur = cv2.blur(user_hair_mask_dilate, (20, 20)) test_res_clear = real_res * (1 - user_matting_8uc3_bald_768_blur[:, :, np.newaxis] / 255) + test_res * ( user_matting_8uc3_bald_768_blur[:, :, np.newaxis] / 255) bald_gene_8uc3_bald_768 = (np.clip(test_res_clear, 0, 255)).astype(np.uint8) # cv2.imshow("bald_gene_8uc3_bald_768", bald_gene_8uc3_bald_768) # cv2.imshow("user_matting_8uc3_bald_768_blur", user_matting_8uc3_bald_768_blur) # cv2.imshow("test_res", test_res) # cv2.waitKey() # cv2.imwrite("/media/DATA_4T/project/hairstyle/分割测试/debug_line_res/test_res_768.png", test_res) # cv2.imwrite("/media/DATA_4T/project/hairstyle/分割测试/debug_line_res/bald_gene_8uc3_bald_768.png", bald_gene_8uc3_bald_768) # cv2.imwrite("/media/DATA_4T/project/hairstyle/分割测试/debug_line_res/user_matting_8uc3_bald_768_blur.png", user_matting_8uc3_bald_768_blur) return bald_gene_8uc3_bald_768, user_matting_8uc3_bald_768_blur class Process_Data(object): def __init__(self, gpu, device_id, save_name=None): 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.generator_matte = Generator_Matte(gpu, device_id) self.generator_baldseg = Generator_BaldSeg_5c(gpu, device_id) self.generator_bald = Generator_Bald(gpu, device_id) self.hair_size = 768 self.bald_output_size = 768 self.color_output_size = 768 def get_hair_M(self, landmark1k): hairstyle_M = landmark_processor.get_transform_mat_hair_ratio(landmark1k, 512, 0.5) return hairstyle_M def get_hair_M_girl_v1(self, landmark1k): hairstyle_M = landmark_processor.get_transform_mat_hair_ratio_v1(landmark1k, 768, ratio=0.5, h_offset=0.45) # default 0.5 0.45 ratio=0.35, h_offset=0.32 return hairstyle_M def get_hair_M_girl_v2(self, landmark1k, long=False): hairstyle_M = landmark_processor.get_transform_mat_hair_ratio_v1(landmark1k, 768, ratio=0.35, h_offset=0.32) # default 0.5 0.45 ratio=0.35, h_offset=0.32 if long: print("------long process-----------") hairstyle_M = landmark_processor.get_transform_mat_hair_ratio_v1(landmark1k, 768, ratio=0.27, h_offset=0.28) # default 0.5 0.45 ratio=0.35, h_offset=0.32 return hairstyle_M def get_hair_M_girl_v3(self, landmark1k): print("get_hair_M_girl_v3 start......") hairstyle_M = landmark_processor.get_transform_mat_hair_ratio_v1(landmark1k, 768, ratio=0.33, h_offset=0.30) # default 0.5 0.45 ratio=0.35, h_offset=0.32 print("get_hair_M_girl_v3 end......") return hairstyle_M def get_hair_M_boy_v1(self, landmark1k): hairstyle_M = landmark_processor.get_transform_mat_hair_ratio_v1(landmark1k, 768, ratio=0.6, h_offset=0.65) return hairstyle_M def get_color_hair_M(self, landmark1k): # color_hair_M = landmark_processor.get_transform_mat_full_face_ratio_stylegan(landmark1k, self.color_output_size, 0.35) # 0.4 color_hair_M = landmark_processor.get_transform_mat_hair_ratio_v1(landmark1k, self.color_output_size, ratio=0.35, h_offset=0.35) return color_hair_M def get_prepare_data_bald(self, user_rgb_8uc3_orisize, user_baldseg_8uc3_orisize, user_landmark_1k2_f_orisize): user_pts1k = user_landmark_1k2_f_orisize.astype(np.int32) user_bald_M = landmark_processor.get_transform_mat_full_face_ratio_stylegan(user_pts1k, self.bald_output_size, 0.4) user_rgb_8uc3_bald_512 = cv2.warpAffine(user_rgb_8uc3_orisize, user_bald_M, (self.bald_output_size, self.bald_output_size)) user_baldseg_8uc3_bald_512 = cv2.warpAffine(user_baldseg_8uc3_orisize, user_bald_M, (self.bald_output_size, self.bald_output_size), flags=cv2.INTER_NEAREST) user_landmark_f1k2_bald_512 = landmark_processor.transform_points(user_pts1k, user_bald_M) return user_rgb_8uc3_bald_512, user_baldseg_8uc3_bald_512, user_landmark_f1k2_bald_512, user_bald_M def get_prepare_data_bald_768(self, user_baldseg_8uc3_orisize, user_landmark_1k2_f_orisize, user_color_M): user_pts1k = user_landmark_1k2_f_orisize.astype(np.int32) user_bald_M = landmark_processor.get_transform_mat_full_face_ratio_stylegan(user_pts1k, self.hair_size, 0.4) # user_rgb_8uc3_bald_768 = cv2.warpAffine(user_rgb_8uc3_orisize, user_color_M, (self.bald_output_size, self.bald_output_size)) user_baldseg_8uc3_bald_768 = cv2.warpAffine(user_baldseg_8uc3_orisize, user_color_M, (self.bald_output_size, self.bald_output_size), flags=cv2.INTER_NEAREST) user_baldseg_8uc3_bald_512 = cv2.warpAffine(user_baldseg_8uc3_orisize, user_bald_M, (self.hair_size, self.hair_size), flags=cv2.INTER_NEAREST) # user_baldseg_8uc3_bald_768_cp = user_baldseg_8uc3_bald_768.copy() # # user_pts1k_768 = landmark_processor.transform_points(user_pts1k, user_color_M) # cood_y = int((user_pts1k_768[214, 1] + user_pts1k_768[99, 1]) // 2) # random_int = 30 # random.randint(0, 30) # kernel = np.ones((random_int, 1), np.uint8) # user_mask_erode = cv2.erode(user_baldseg_8uc3_bald_768, kernel, iterations=1) # user_baldseg_8uc3_bald_768[:cood_y, :, :] = user_mask_erode[:cood_y, :, :] # # inter_res = np.concatenate((user_baldseg_8uc3_bald_768_cp, user_baldseg_8uc3_bald_768), axis=1) # cv2.imshow("inter_res", inter_res) # cv2.waitKey() # user_baldseg_8uc3_orisize = cv2.warpAffine(user_baldseg_8uc3_bald_768, cv2.invertAffineTransform(user_color_M), # dsize=(user_baldseg_8uc3_orisize.shape[1], user_baldseg_8uc3_orisize.shape[0]), flags=cv2.INTER_NEAREST) return user_baldseg_8uc3_bald_768, user_baldseg_8uc3_bald_512, user_bald_M # , user_baldseg_8uc3_orisize def get_user_blad(self, user_rgb_8uc3_orisize, user_landmark_1k2_f_orisize): """ input： user_rgb_8uc3_orisize: 用户图 原始尺寸 user_landmark_1k2_f_orisize: 用户图 1k 点 output: 图像尺寸基于： 人脸 512， 图像均为3通道 user_bald_8uc3_512： 用户图 光头， uint8 （0-255） """ # 得到512 小图 M 矩阵 user_color_M = self.get_color_hair_M(user_landmark_1k2_f_orisize) # 得到512 小图 M 矩阵 # user_hairstyle_M = self.get_hair_M(user_landmark_1k2_f_orisize) # 1k 点 转换到 512 尺寸 # user_landmark_f1k2_512 = landmark_processor.transform_points(user_landmark_1k2_f_orisize, user_hairstyle_M) user_landmark_f1k2_bald_768 = landmark_processor.transform_points(user_landmark_1k2_f_orisize, user_color_M) # # 用户图 头发毛躁 warpaffine usr_color_ratio = np.sqrt((user_color_M[0][0] * user_color_M[0][0]) + (user_color_M[1][0] * user_color_M[1][0])) user_color_M_tmp = user_color_M / usr_color_ratio user_rgb_8uc3_bald_768 = cv2.warpAffine(user_rgb_8uc3_orisize, user_color_M_tmp, ( int(self.bald_output_size / usr_color_ratio), int(self.bald_output_size / usr_color_ratio)), borderMode=cv2.BORDER_CONSTANT, borderValue=[255, 255, 255]) user_rgb_8uc3_bald_768 = cv2.resize(user_rgb_8uc3_bald_768, (self.bald_output_size, self.bald_output_size), fx=usr_color_ratio, fy=usr_color_ratio, interpolation=cv2.INTER_AREA) user_rgb_8uc3_bald_768_bl_bg = cv2.warpAffine(user_rgb_8uc3_orisize, user_color_M_tmp, ( int(self.bald_output_size / usr_color_ratio), int(self.bald_output_size / usr_color_ratio))) user_rgb_8uc3_bald_768_bl_bg = cv2.resize(user_rgb_8uc3_bald_768_bl_bg, (self.bald_output_size, self.bald_output_size), fx=usr_color_ratio, fy=usr_color_ratio, interpolation=cv2.INTER_AREA) # 用户图得到 头发 matting 小图 _, user_matting_8uc1_bald_768 = self.generator_matte.matte_inference(user_rgb_8uc3_bald_768, user_landmark_f1k2_bald_768) # cv2.imshow('user_rgb_8uc3_bald_768', user_rgb_8uc3_bald_768) # cv2.imshow('user_matting_8uc1_bald_768', user_matting_8uc1_bald_768) # cv2.waitKey() if user_matting_8uc1_bald_768.shape[0] != 768 or user_matting_8uc1_bald_768.shape[1] != 768: user_matting_8uc1_bald_768 = cv2.resize(user_matting_8uc1_bald_768, (user_rgb_8uc3_bald_768.shape[1], user_rgb_8uc3_bald_768.shape[0]), interpolation=cv2.INTER_CUBIC) user_matting_8uc3_bald_768 = np.repeat(user_matting_8uc1_bald_768[:, :, np.newaxis], 3, axis=2) user_matting_8uc3_bald_orisize = cv2.warpAffine(user_matting_8uc1_bald_768, cv2.invertAffineTransform(user_color_M), ( user_rgb_8uc3_orisize.shape[1], user_rgb_8uc3_orisize.shape[0]), flags=cv2.INTER_CUBIC) # 光头分割 user_baldseg_8uc3_orisize = self.generator_baldseg.forward(user_rgb_8uc3_orisize, user_matting_8uc3_bald_orisize, user_landmark_1k2_f_orisize) # 用户图 生成光头 user_baldseg_8uc3_bald_768, user_baldseg_8uc3_bald_512, user_bald_M \ = self.get_prepare_data_bald_768(user_baldseg_8uc3_orisize, user_landmark_1k2_f_orisize, user_color_M) orig_h, orig_w, _ = user_rgb_8uc3_orisize.shape # TODO: 测试直接采用M 矩阵warp 到 user_hairstyle_M 后的512 尺寸 # bald_gene_8uc3_bald_768, user_hair_mask_blur_768 = self.generator_bald.Geneator_Bald_inference(user_rgb_8uc3_bald_768.copy(), # user_matting_8uc3_bald_768, # user_baldseg_8uc3_bald_768, # user_landmark_f1k2_bald_768) bald_gene_8uc3_bald_768, user_hair_mask_blur_768 = self.generator_bald.Geneator_Bald_inference( user_rgb_8uc3_bald_768_bl_bg.copy(), user_matting_8uc3_bald_768, user_baldseg_8uc3_bald_768, user_landmark_f1k2_bald_768) inv_M_user = cv2.invertAffineTransform(user_color_M) bald_gene_8uc3_orisize = user_rgb_8uc3_orisize.copy() user_bald_gene_8uc3_orisize = cv2.warpAffine(bald_gene_8uc3_bald_768, inv_M_user, (orig_w, orig_h), dst=bald_gene_8uc3_orisize.copy(), borderMode=cv2.BORDER_TRANSPARENT) user_hair_mask_blur_orisize = cv2.warpAffine(user_hair_mask_blur_768, inv_M_user, (orig_w, orig_h), flags=cv2.INTER_CUBIC) user_bald_res_8uc3_orisize = ( (1 - user_hair_mask_blur_orisize[:, :, np.newaxis] / 255) * user_rgb_8uc3_orisize + \ (user_hair_mask_blur_orisize[:, :, np.newaxis] / 255) * user_bald_gene_8uc3_orisize).astype( np.uint8) return user_bald_res_8uc3_orisize, user_bald_gene_8uc3_orisize, user_matting_8uc3_bald_768, user_baldseg_8uc3_bald_768 def hair512_to_768(self, user_color_M, user_hairstyle_M): M_ori = np.zeros((3, 3), dtype=np.float32) M_ori[:2, :] = cv2.invertAffineTransform(user_hairstyle_M) M_ori[2:, :] = [0, 0, 1] matAffine_ori = np.zeros((3, 3), dtype=np.float32) matAffine_ori[:2, :] = user_color_M matAffine_ori[2:, :] = [0, 0, 1] new_mat = matAffine_ori.dot(M_ori) return new_mat[:2, :] def get_prepare_data(self, user_rgb_8uc3_orisize, user_landmark_1k2_f_orisize, ref_rgb_8uc3_orisize, ref_landmark_1k2_f_orisize): """ input： user_rgb_8uc3_orisize: 用户图 原始尺寸 user_landmark_1k2_f_orisize: 用户图 1k 点 ref_rgb_8uc3_orisize: 参考图 原始尺寸 ref_landmark_1k2_f_orisize: 参考图 1k 点 output: 图像尺寸基于： 人脸 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 user_baldseg_8uc3_512：用户图 光头分割 mask， uint8 （0-255） user_bald_8uc3_512： 用户图 光头， uint8 （0-255） user_landmark_f1k2_512： 用户图 关键点 1k*2 float32 user_hairstyle_M: 用户图 原图到512小图 M 矩阵 """ # 得到512 小图 M 矩阵 user_hairstyle_M = self.get_hair_M(user_landmark_1k2_f_orisize) ref_hairstyle_M = self.get_hair_M(ref_landmark_1k2_f_orisize) # 1k 点 转换到 512 尺寸 user_landmark_f1k2_512 = landmark_processor.transform_points(user_landmark_1k2_f_orisize, user_hairstyle_M) ref_landmark_f1k2_512 = landmark_processor.transform_points(ref_landmark_1k2_f_orisize, ref_hairstyle_M) # 用户图 头发毛躁 warpaffine usr_ratio = np.sqrt( (user_hairstyle_M[0][0] * user_hairstyle_M[0][0]) + (user_hairstyle_M[1][0] * user_hairstyle_M[1][0])) user_hairstyle_M_tmp = user_hairstyle_M / usr_ratio user_rgb_8uc3_512 = cv2.warpAffine(user_rgb_8uc3_orisize, user_hairstyle_M_tmp, (int(self.hair_size / usr_ratio), int(self.hair_size / usr_ratio)), borderMode=cv2.BORDER_CONSTANT, borderValue=[255, 255, 255]) user_rgb_8uc3_512 = cv2.resize(user_rgb_8uc3_512, (self.hair_size, self.hair_size), fx=usr_ratio, fy=usr_ratio, interpolation=cv2.INTER_AREA) # 参考图 头发毛躁 warpaffine ratio = np.sqrt( (ref_hairstyle_M[0][0] * ref_hairstyle_M[0][0]) + (ref_hairstyle_M[1][0] * ref_hairstyle_M[1][0])) ref_hairstyle_M_tmp = ref_hairstyle_M / ratio ref_rgb_8uc3_512 = cv2.warpAffine(ref_rgb_8uc3_orisize, ref_hairstyle_M_tmp, (int(self.hair_size / ratio), int(self.hair_size / ratio)), borderMode=cv2.BORDER_CONSTANT, borderValue=[255, 255, 255]) ref_rgb_8uc3_512 = cv2.resize(ref_rgb_8uc3_512, (self.hair_size, self.hair_size), fx=ratio, fy=ratio, interpolation=cv2.INTER_AREA) # ref_rgb_8uc3_512 = cv2.warpAffine(ref_rgb_8uc3_orisize, ref_hairstyle_M, (self.hair_size, self.hair_size), borderValue=[255, 255, 255]) # 参考图直接得到 头发 matting 小图 _, ref_matte_pred_8uc1_512 = self.generator_matte.matte_inference(ref_rgb_8uc3_512, ref_landmark_f1k2_512) torch.cuda.empty_cache() ref_matting_8uc1_512 = cv2.resize(ref_matte_pred_8uc1_512, (ref_rgb_8uc3_512.shape[1], ref_rgb_8uc3_512.shape[0]), interpolation=cv2.INTER_CUBIC) ref_matting_8uc3_512 = np.repeat(ref_matting_8uc1_512[:, :, np.newaxis], 3, axis=2) # 光头分割 user_baldseg_8uc3_orisize = self.generator_baldseg.forward(user_rgb_8uc3_orisize, user_landmark_1k2_f_orisize) ref_baldseg_8uc3_orisize = self.generator_baldseg.forward(ref_rgb_8uc3_orisize, ref_landmark_1k2_f_orisize) # 用户图 生成光头 user_rgb_8uc3_bald_512, user_baldseg_8uc3_bald_512, user_landmark_f1k2_bald_512, \ user_bald_M = self.get_prepare_data_bald(user_rgb_8uc3_orisize, user_baldseg_8uc3_orisize, user_landmark_1k2_f_orisize) # 用户图得到 头发 matting 小图 for blad _, user_matting_8uc1_bald_512 = self.generator_matte.matte_inference(user_rgb_8uc3_bald_512, user_landmark_f1k2_bald_512) if user_matting_8uc1_bald_512.shape[0] != 512 or user_matting_8uc1_bald_512.shape[1] != 512: user_matting_8uc1_bald_512 = cv2.resize(user_matting_8uc1_bald_512, (user_rgb_8uc3_bald_512.shape[1], user_rgb_8uc3_bald_512.shape[0]), interpolation=cv2.INTER_CUBIC) user_matting_8uc3_bald_512 = np.repeat(user_matting_8uc1_bald_512[:, :, np.newaxis], 3, axis=2) orig_h, orig_w, _ = user_rgb_8uc3_orisize.shape # TODO: 测试直接采用M 矩阵warp 到 user_hairstyle_M 后的512 尺寸 bald_gene_8uc3_bald_512 = self.generator_bald.Geneator_Bald_inference(user_rgb_8uc3_bald_512, user_matting_8uc3_bald_512, user_baldseg_8uc3_bald_512, user_landmark_f1k2_bald_512) inv_M_user = cv2.invertAffineTransform(user_bald_M) bald_gene_8uc3_orisize = user_rgb_8uc3_orisize.copy() user_bald_gene_8uc3_orisize = cv2.warpAffine(bald_gene_8uc3_bald_512, inv_M_user, (orig_w, orig_h), dst=bald_gene_8uc3_orisize, borderMode=cv2.BORDER_TRANSPARENT) user_baldseg_8uc3_512 = cv2.warpAffine(user_baldseg_8uc3_orisize, user_hairstyle_M, (self.hair_size, self.hair_size)) user_bald_8uc3_512 = cv2.warpAffine(user_bald_gene_8uc3_orisize, user_hairstyle_M, (self.hair_size, self.hair_size), borderValue=[255, 255, 255]) ref_baldseg_8uc3_512 = cv2.warpAffine(ref_baldseg_8uc3_orisize, ref_hairstyle_M, (self.hair_size, self.hair_size)) return user_rgb_8uc3_512, user_baldseg_8uc3_512, user_bald_8uc3_512, user_landmark_f1k2_512, user_hairstyle_M, ref_rgb_8uc3_512, ref_matting_8uc3_512, ref_baldseg_8uc3_512, ref_landmark_f1k2_512 def get_prepare_user_data(self, user_rgb_8uc3_orisize, user_landmark_1k2_f_orisize): """ input：F user_rgb_8uc3_orisize: 用户图 原始尺寸 user_landmark_1k2_f_orisize: 用户图 1k 点 output: 图像尺寸基于： 人脸 512， 图像均为3通道 user_baldseg_8uc3_512：用户图 光头分割 mask， uint8 （0-255） user_bald_8uc3_512： 用户图 光头， uint8 （0-255） user_landmark_f1k2_512： 用户图 关键点 1k*2 float32 user_hairstyle_M: 用户图 原图到512小图 M 矩阵 """ # 得到512 小图 M 矩阵 user_hairstyle_M = self.get_hair_M(user_landmark_1k2_f_orisize) # 1k 点 转换到 512 尺寸 user_landmark_f1k2_512 = landmark_processor.transform_points(user_landmark_1k2_f_orisize, user_hairstyle_M) # 用户图 头发毛躁 warpaffine usr_ratio = np.sqrt( (user_hairstyle_M[0][0] * user_hairstyle_M[0][0]) + (user_hairstyle_M[1][0] * user_hairstyle_M[1][0])) user_hairstyle_M_tmp = user_hairstyle_M / usr_ratio user_rgb_8uc3_512 = cv2.warpAffine(user_rgb_8uc3_orisize, user_hairstyle_M_tmp, (int(self.hair_size / usr_ratio), int(self.hair_size / usr_ratio)), borderMode=cv2.BORDER_CONSTANT, borderValue=[255, 255, 255]) user_rgb_8uc3_512 = cv2.resize(user_rgb_8uc3_512, (self.hair_size, self.hair_size), fx=usr_ratio, fy=usr_ratio, interpolation=cv2.INTER_AREA) # 光头分割 user_baldseg_8uc3_orisize = self.generator_baldseg.forward(user_rgb_8uc3_orisize, user_landmark_1k2_f_orisize) # 用户图 生成光头 user_rgb_8uc3_bald_512, user_baldseg_8uc3_bald_512, user_landmark_f1k2_bald_512, \ user_bald_M = self.get_prepare_data_bald(user_rgb_8uc3_orisize, user_baldseg_8uc3_orisize, user_landmark_1k2_f_orisize) # 用户图得到 头发 matting 小图 for blad _, user_matting_8uc1_bald_512 = self.generator_matte.matte_inference(user_rgb_8uc3_bald_512, user_landmark_f1k2_bald_512) if user_matting_8uc1_bald_512.shape[0] != 512 or user_matting_8uc1_bald_512.shape[1] != 512: user_matting_8uc1_bald_512 = cv2.resize(user_matting_8uc1_bald_512, (user_rgb_8uc3_bald_512.shape[1], user_rgb_8uc3_bald_512.shape[0]), interpolation=cv2.INTER_CUBIC) user_matting_8uc3_bald_512 = np.repeat(user_matting_8uc1_bald_512[:, :, np.newaxis], 3, axis=2) orig_h, orig_w, _ = user_rgb_8uc3_orisize.shape # TODO: 测试直接采用M 矩阵warp 到 user_hairstyle_M 后的512 尺寸 bald_gene_8uc3_bald_512 = self.generator_bald.Geneator_Bald_inference( user_rgb_8uc3_bald_512, user_matting_8uc3_bald_512, user_baldseg_8uc3_bald_512, user_landmark_f1k2_bald_512) inv_M_user = cv2.invertAffineTransform(user_bald_M) bald_gene_8uc3_orisize = user_rgb_8uc3_orisize.copy() user_bald_gene_8uc3_orisize = cv2.warpAffine(bald_gene_8uc3_bald_512, inv_M_user, (orig_w, orig_h), dst=bald_gene_8uc3_orisize, borderMode=cv2.BORDER_TRANSPARENT) user_baldseg_8uc3_512 = cv2.warpAffine(user_baldseg_8uc3_orisize, user_hairstyle_M, (self.hair_size, self.hair_size)) user_bald_8uc3_512 = cv2.warpAffine(user_bald_gene_8uc3_orisize, user_hairstyle_M, (self.hair_size, self.hair_size), borderValue=[255, 255, 255]) return user_rgb_8uc3_512, user_baldseg_8uc3_512, user_bald_8uc3_512, user_landmark_f1k2_512, user_hairstyle_M def get_prepare_user_768_data(self, user_rgb_8uc3_orisize, user_landmark_1k2_f_orisize, ratio=1): """ input： user_rgb_8uc3_orisize: 用户图 原始尺寸 user_landmark_1k2_f_orisize: 用户图 1k 点 output: 图像尺寸基于： 人脸 512， 图像均为3通道 user_baldseg_8uc3_768：用户图 光头分割 mask， uint8 （0-255） user_baldseg_8uc3_512：用户图 光头分割 mask， uint8 （0-255） user_bald_8uc3_512： 用户图 光头， uint8 （0-255） user_landmark_f1k2_768： 用户图 关键点 1k*2 float32 user_landmark_f1k2_512： 用户图 关键点 1k*2 float32 user_color_M: 用户图 原图到768小图 M 矩阵 user_hairstyle_M: 用户图 原图到512小图 M 矩阵 """ # 得到768 小图 M 矩阵 user_color_M = self.get_color_hair_M(user_landmark_1k2_f_orisize) # 得到768 小图 M 矩阵 发型 if ratio == 0: user_hairstyle_M = self.get_hair_M_boy_v1(user_landmark_1k2_f_orisize) elif ratio == 1: user_hairstyle_M = self.get_hair_M_girl_v1(user_landmark_1k2_f_orisize) elif ratio == 2: user_hairstyle_M = self.get_hair_M_girl_v2(user_landmark_1k2_f_orisize) elif ratio == 3: user_hairstyle_M = self.get_hair_M_girl_v3(user_landmark_1k2_f_orisize) else: user_hairstyle_M = self.get_hair_M_girl_v1(user_landmark_1k2_f_orisize) # np.save("/home/yangchaojie/Desktop/diffusion_model/lora/data_for_liveme/tmp/hair_swap/person/1706086560_1706086559583_ST0023_W0_045_0.npy", user_hairstyle_M) # user_color_M = user_hairstyle_M # 1k 点 转换到 768 尺寸 user_landmark_f1k2_bald_768 = landmark_processor.transform_points(user_landmark_1k2_f_orisize, user_color_M) # 1k 点 转换到 768 尺寸 发型 user_landmark_f1k2_768 = landmark_processor.transform_points(user_landmark_1k2_f_orisize, user_hairstyle_M) # 用户图 头发毛躁 warpaffine 768 usr_color_ratio = np.sqrt( (user_color_M[0][0] * user_color_M[0][0]) + (user_color_M[1][0] * user_color_M[1][0])) user_color_M_tmp = user_color_M / usr_color_ratio user_rgb_8uc3_768 = cv2.warpAffine(user_rgb_8uc3_orisize, user_color_M_tmp, (int(self.bald_output_size / usr_color_ratio), int(self.bald_output_size / usr_color_ratio)), borderMode=cv2.BORDER_CONSTANT, borderValue=[255, 255, 255]) user_rgb_8uc3_768 = cv2.resize(user_rgb_8uc3_768, (self.bald_output_size, self.bald_output_size), fx=usr_color_ratio, fy=usr_color_ratio, interpolation=cv2.INTER_AREA) user_rgb_8uc3_768_bald = cv2.warpAffine(user_rgb_8uc3_orisize, user_color_M_tmp, (int(self.bald_output_size / usr_color_ratio), int(self.bald_output_size / usr_color_ratio))) user_rgb_8uc3_768_bald = cv2.resize(user_rgb_8uc3_768_bald, (self.bald_output_size, self.bald_output_size), fx=usr_color_ratio, fy=usr_color_ratio, interpolation=cv2.INTER_AREA) # 用户图得到 头发 matting 小图 for blad user_matting_fg_8uc3_bald_768, user_matting_8uc1_bald_768 = self.generator_matte.matte_inference( user_rgb_8uc3_768, user_landmark_f1k2_bald_768) if user_matting_8uc1_bald_768.shape[0] != 768 or user_matting_8uc1_bald_768.shape[1] != 768: user_matting_8uc1_bald_768 = cv2.resize(user_matting_8uc1_bald_768, (user_rgb_8uc3_768.shape[1], user_rgb_8uc3_768.shape[0]), interpolation=cv2.INTER_CUBIC) # user_matting_fg_8uc3_bald_768 = cv2.resize(user_matting_fg_8uc3_bald_768, # (user_rgb_8uc3_768.shape[1], user_rgb_8uc3_768.shape[0]), # interpolation=cv2.INTER_CUBIC) user_matting_8uc3_bald_768 = np.repeat(user_matting_8uc1_bald_768[:, :, np.newaxis], 3, axis=2) user_matting_8uc3_bald_orisize = cv2.warpAffine(user_matting_8uc1_bald_768, cv2.invertAffineTransform(user_color_M), ( user_rgb_8uc3_orisize.shape[1], user_rgb_8uc3_orisize.shape[0]), flags=cv2.INTER_CUBIC) # 光头分割 user_baldseg_8uc3_orisize = self.generator_baldseg.forward(user_rgb_8uc3_orisize, user_matting_8uc3_bald_orisize, user_landmark_1k2_f_orisize) # 用户图 生成光头 需准备 768 尺寸； 512 尺寸各一 user_baldseg_8uc3_bald_768, user_baldseg_8uc3_bald_512, user_bald_M \ = self.get_prepare_data_bald_768(user_baldseg_8uc3_orisize, user_landmark_1k2_f_orisize, user_color_M) orig_h, orig_w, _ = user_rgb_8uc3_orisize.shape # cv2.imshow("user_matting_8uc3_bald_768",user_matting_8uc3_bald_768) bald_gene_8uc3_bald_768, user_hair_mask_blur_768 = self.generator_bald.Geneator_Bald_inference( user_rgb_8uc3_768_bald.copy(), user_matting_8uc3_bald_768, user_baldseg_8uc3_bald_768, user_landmark_f1k2_bald_768) # cv2.imshow("user_hair_mask_blur_768", user_hair_mask_blur_768) # cv2.waitKey() inv_M_user = cv2.invertAffineTransform(user_color_M) bald_gene_8uc3_orisize = user_rgb_8uc3_orisize.copy() user_bald_gene_8uc3_orisize = cv2.warpAffine(bald_gene_8uc3_bald_768, inv_M_user, (orig_w, orig_h), dst=bald_gene_8uc3_orisize.copy(), borderMode=cv2.BORDER_TRANSPARENT) user_hair_mask_blur_orisize = cv2.warpAffine(user_hair_mask_blur_768, inv_M_user, (orig_w, orig_h), flags=cv2.INTER_CUBIC) user_bald_res_8uc3_orisize = ( (1 - user_hair_mask_blur_orisize[:, :, np.newaxis] / 255) * user_rgb_8uc3_orisize + \ (user_hair_mask_blur_orisize[:, :, np.newaxis] / 255) * user_bald_gene_8uc3_orisize).astype( np.uint8) user_baldseg_8uc3_768 = cv2.warpAffine(user_baldseg_8uc3_orisize, user_hairstyle_M, (self.hair_size, self.hair_size)) # , borderValue=[255, 255, 255] user_bald_8uc3_768 = cv2.warpAffine(user_bald_res_8uc3_orisize, user_hairstyle_M, (self.hair_size, self.hair_size)) # , borderValue=[255, 255, 255] return 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 def get_prepare_ref_data(self, ref_rgb_8uc3_orisize, ref_landmark_1k2_f_orisize): """ input： ref_rgb_8uc3_orisize: 参考图 原始尺寸 ref_landmark_1k2_f_orisize: 参考图 1k 点 output: 图像尺寸基于： 人脸 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 """ # 得到512 小图 M 矩阵 ref_hairstyle_M = self.get_hair_M(ref_landmark_1k2_f_orisize) # ref_color_hair_M = self.get_color_hair_M(ref_landmark_1k2_f_orisize) # 1k 点 转换到 512 尺寸 ref_landmark_f1k2_512 = landmark_processor.transform_points(ref_landmark_1k2_f_orisize, ref_hairstyle_M) # ref_landmark_color_f1k2_512 = landmark_processor.transform_points(ref_landmark_1k2_f_orisize, ref_color_hair_M) # 参考图 头发毛躁 warpaffine ratio = np.sqrt( (ref_hairstyle_M[0][0] * ref_hairstyle_M[0][0]) + (ref_hairstyle_M[1][0] * ref_hairstyle_M[1][0])) ref_hairstyle_M_tmp = ref_hairstyle_M / ratio ref_rgb_8uc3_512 = cv2.warpAffine(ref_rgb_8uc3_orisize, ref_hairstyle_M_tmp, (int(self.hair_size / ratio), int(self.hair_size / ratio)), borderMode=cv2.BORDER_CONSTANT, borderValue=[255, 255, 255]) ref_rgb_8uc3_512 = cv2.resize(ref_rgb_8uc3_512, (self.hair_size, self.hair_size), fx=ratio, fy=ratio, interpolation=cv2.INTER_AREA) # # 参考图 头发毛躁 warpaffine # ratio_color = np.sqrt( # (ref_color_hair_M[0][0] * ref_color_hair_M[0][0]) + (ref_color_hair_M[1][0] * ref_color_hair_M[1][0])) # ref_color_hair_M_tmp = ref_color_hair_M / ratio_color # ref_rgb_8uc3_768 = cv2.warpAffine(ref_rgb_8uc3_orisize, ref_color_hair_M_tmp, # (int(self.bald_output_size / ratio_color), int(self.bald_output_size / ratio_color)), # borderMode=cv2.BORDER_CONSTANT, borderValue=[255, 255, 255]) # ref_rgb_8uc3_768 = cv2.resize(ref_rgb_8uc3_768, (self.bald_output_size, self.bald_output_size), fx=ratio_color, fy=ratio_color, # interpolation=cv2.INTER_AREA) # 参考图直接得到 头发 matting 小图 _, ref_matte_pred_8uc1_512 = self.generator_matte.matte_inference(ref_rgb_8uc3_512, ref_landmark_f1k2_512) torch.cuda.empty_cache() ref_matting_8uc1_512 = cv2.resize(ref_matte_pred_8uc1_512, (ref_rgb_8uc3_512.shape[1], ref_rgb_8uc3_512.shape[0]), interpolation=cv2.INTER_CUBIC) ref_matting_8uc3_512 = np.repeat(ref_matting_8uc1_512[:, :, np.newaxis], 3, axis=2) # 光头分割 ref_baldseg_8uc3_orisize = self.generator_baldseg.forward(ref_rgb_8uc3_orisize, ref_landmark_1k2_f_orisize) ref_baldseg_8uc3_512 = cv2.warpAffine(ref_baldseg_8uc3_orisize, ref_hairstyle_M, (self.hair_size, self.hair_size), borderValue=[255, 255, 255]) return ref_rgb_8uc3_512, ref_matting_8uc3_512, ref_baldseg_8uc3_512, ref_landmark_f1k2_512 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 get_prepare_ref_768_data(self, ref_rgb_8uc3_orisize, ref_landmark_1k2_f_orisize, ratio=1, long=False): """ input： ref_rgb_8uc3_orisize: 参考图 原始尺寸 ref_landmark_1k2_f_orisize: 参考图 1k 点 output: 图像尺寸基于： 人脸 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 """ # 得到512 小图 M 矩阵 if ratio == 0: ref_hairstyle_M = self.get_hair_M_boy_v1(ref_landmark_1k2_f_orisize) elif ratio == 1: ref_hairstyle_M = self.get_hair_M_girl_v1(ref_landmark_1k2_f_orisize) elif ratio == 2: ref_hairstyle_M = self.get_hair_M_girl_v2(ref_landmark_1k2_f_orisize, long=long) else: ref_hairstyle_M = self.get_hair_M_girl_v1(ref_landmark_1k2_f_orisize) # ref_color_hair_M = self.get_color_hair_M(ref_landmark_1k2_f_orisize) # print("ref_hairstyle_M: ", ref_hairstyle_M) # 1k 点 转换到 512 尺寸 ref_landmark_f1k2_768 = landmark_processor.transform_points(ref_landmark_1k2_f_orisize, ref_hairstyle_M) # 参考图 头发毛躁 warpaffine ratio = np.sqrt( (ref_hairstyle_M[0][0] * ref_hairstyle_M[0][0]) + (ref_hairstyle_M[1][0] * ref_hairstyle_M[1][0])) ref_hairstyle_M_tmp = ref_hairstyle_M / ratio ref_rgb_8uc3_768 = cv2.warpAffine(ref_rgb_8uc3_orisize, ref_hairstyle_M_tmp, (int(self.bald_output_size / ratio), int(self.bald_output_size / ratio)), borderMode=cv2.BORDER_CONSTANT) # , borderValue=[255, 255, 255] ref_rgb_8uc3_768 = cv2.resize(ref_rgb_8uc3_768, (self.bald_output_size, self.bald_output_size), fx=ratio, fy=ratio, interpolation=cv2.INTER_AREA) # 参考图直接得到 头发 matting 小图 ref_matte_fg_8uc3_768, ref_matte_pred_8uc1_768 = self.generator_matte.matte_inference(ref_rgb_8uc3_768, ref_landmark_f1k2_768) torch.cuda.empty_cache() ref_matting_fg_8uc3_768 = cv2.resize(ref_matte_fg_8uc3_768, (ref_rgb_8uc3_768.shape[1], ref_rgb_8uc3_768.shape[0]), interpolation=cv2.INTER_CUBIC) ref_matting_8uc1_768 = cv2.resize(ref_matte_pred_8uc1_768, (ref_rgb_8uc3_768.shape[1], ref_rgb_8uc3_768.shape[0]), interpolation=cv2.INTER_CUBIC) ref_matting_8uc3_768 = np.repeat(ref_matting_8uc1_768[:, :, np.newaxis], 3, axis=2) ref_matting_8uc1_orisize = cv2.warpAffine(ref_matting_8uc1_768, cv2.invertAffineTransform(ref_hairstyle_M), (ref_rgb_8uc3_orisize.shape[1], ref_rgb_8uc3_orisize.shape[0]), flags=cv2.INTER_CUBIC) # 光头分割 ref_baldseg_8uc3_orisize = self.generator_baldseg.forward(ref_rgb_8uc3_orisize, ref_matting_8uc1_orisize, ref_landmark_1k2_f_orisize) ref_baldseg_8uc3_768 = cv2.warpAffine(ref_baldseg_8uc3_orisize, ref_hairstyle_M, (self.bald_output_size, self.bald_output_size), flags=cv2.INTER_NEAREST) # , borderValue=[255, 255, 255] return ref_rgb_8uc3_768, ref_matting_fg_8uc3_768, ref_matting_8uc3_768, ref_baldseg_8uc3_768, ref_landmark_f1k2_768 def get_prepare_ref_768_bald_data(self, ref_rgb_8uc3_orisize, ref_landmark_1k2_f_orisize): """ input： ref_rgb_8uc3_orisize: 参考图 原始尺寸 ref_landmark_1k2_f_orisize: 参考图 1k 点 output: 图像尺寸基于： 人脸 512， 图像均为3通道 ref_rgb_8uc3_768: 参考图， size 512， uint8 （0-255） ref_matting_8uc3_768：参考图 matting alpha 值， uint8 （0-255） ref_baldseg_8uc3_768: 参考图 光头分割， uint8 （0-255） ref_landmark_f1k2_768： 参考图 关键点 1k*2 float32 """ # 得到512 小图 M 矩阵 # ref_hairstyle_M = self.get_hair_M(ref_landmark_1k2_f_orisize) ref_color_hair_M = self.get_color_hair_M(ref_landmark_1k2_f_orisize) # 1k 点 转换到 512 尺寸 # ref_landmark_f1k2_512 = landmark_processor.transform_points(ref_landmark_1k2_f_orisize, ref_hairstyle_M) ref_landmark_color_f1k2_768 = landmark_processor.transform_points(ref_landmark_1k2_f_orisize, ref_color_hair_M) # 参考图 头发毛躁 warpaffine ratio = np.sqrt( (ref_color_hair_M[0][0] * ref_color_hair_M[0][0]) + (ref_color_hair_M[1][0] * ref_color_hair_M[1][0])) ref_color_hair_M_tmp = ref_color_hair_M / ratio ref_rgb_8uc3_768 = cv2.warpAffine(ref_rgb_8uc3_orisize, ref_color_hair_M_tmp, (int(self.color_output_size / ratio), int(self.color_output_size / ratio)), borderMode=cv2.BORDER_CONSTANT) # , borderValue=[255, 255, 255] ref_rgb_8uc3_768 = cv2.resize(ref_rgb_8uc3_768, (self.color_output_size, self.color_output_size), fx=ratio, fy=ratio, interpolation=cv2.INTER_AREA) # 参考图直接得到 头发 matting 小图 _, ref_matte_pred_8uc1_768 = self.generator_matte.matte_inference(ref_rgb_8uc3_768, ref_landmark_color_f1k2_768) torch.cuda.empty_cache() ref_matting_8uc1_768 = cv2.resize(ref_matte_pred_8uc1_768, (ref_rgb_8uc3_768.shape[1], ref_rgb_8uc3_768.shape[0]), interpolation=cv2.INTER_CUBIC) ref_matting_8uc3_768 = np.repeat(ref_matting_8uc1_768[:, :, np.newaxis], 3, axis=2) ref_matting_8uc1_orisize = cv2.warpAffine(ref_matting_8uc1_768, cv2.invertAffineTransform(ref_color_hair_M), (ref_rgb_8uc3_orisize.shape[1], ref_rgb_8uc3_orisize.shape[0]), flags=cv2.INTER_CUBIC) # 光头分割 ref_baldseg_8uc3_orisize = self.generator_baldseg.forward(ref_rgb_8uc3_orisize, ref_matting_8uc1_orisize, ref_landmark_1k2_f_orisize) ref_baldseg_8uc3_768 = cv2.warpAffine(ref_baldseg_8uc3_orisize, ref_color_hair_M, (self.color_output_size, self.color_output_size)) # , borderValue=[255, 255, 255] return ref_rgb_8uc3_768, ref_matting_8uc3_768, ref_baldseg_8uc3_768, ref_landmark_color_f1k2_768 def get_prepare_hair_color_data(self, user_rgb_8uc3_orisize, user_landmark_1k2_f_orisize, ref_rgb_8uc3_orisize, ref_landmark_1k2_f_orisize): """ input： user_rgb_8uc3_orisize: 用户图 原始尺寸 user_landmark_1k2_f_orisize: 用户图 1k 点 ref_rgb_8uc3_orisize: 参考图 原始尺寸 ref_landmark_1k2_f_orisize: 参考图 1k 点 output: 图像尺寸基于： 人脸 768， 图像均为3通道 user_rgb_8uc3_change_color_768, user_matting_8uc3_change_color_768, ref_rgb_8uc3_change_color_768, ref_matting_8uc3_change_color_768 user_rgb_8uc3_change_color_768: 用户图 align， size: 768 uint8 （0-255） user_matting_8uc3_change_color_768: 用户图 matting， size: 768 uint8 （0-255） ref_rgb_8uc3_change_color_768: 参考图， align, size 512， uint8 （0-255） ref_matting_8uc3_change_color_768: 参考图 matting alpha 值， uint8 （0-255） """ # 得到768 小图 M 矩阵 user_hair_color_M = self.get_color_hair_M(user_landmark_1k2_f_orisize) ref_hair_color_M = self.get_color_hair_M(ref_landmark_1k2_f_orisize) # 1k 点 转换到 768 尺寸 user_landmark_f1k2_768 = landmark_processor.transform_points(user_landmark_1k2_f_orisize, user_hair_color_M) ref_landmark_f1k2_768 = landmark_processor.transform_points(ref_landmark_1k2_f_orisize, ref_hair_color_M) # 用户图 头发毛躁 warpaffine usr_ratio = np.sqrt( (user_hair_color_M[0][0] * user_hair_color_M[0][0]) + (user_hair_color_M[1][0] * user_hair_color_M[1][0])) user_hair_color_M_tmp = user_hair_color_M / usr_ratio user_rgb_8uc3_768 = cv2.warpAffine(user_rgb_8uc3_orisize, user_hair_color_M_tmp, (int(self.color_output_size / usr_ratio), int(self.color_output_size / usr_ratio))) user_rgb_8uc3_768 = cv2.resize(user_rgb_8uc3_768, (self.color_output_size, self.color_output_size), fx=usr_ratio, fy=usr_ratio, interpolation=cv2.INTER_AREA) # 参考图 头发毛躁 warpaffine ratio = np.sqrt( (ref_hair_color_M[0][0] * ref_hair_color_M[0][0]) + (ref_hair_color_M[1][0] * ref_hair_color_M[1][0])) ref_hair_color_M_tmp = ref_hair_color_M / ratio ref_rgb_8uc3_768 = cv2.warpAffine(ref_rgb_8uc3_orisize, ref_hair_color_M_tmp, (int(self.color_output_size / ratio), int(self.color_output_size / ratio))) ref_rgb_8uc3_768 = cv2.resize(ref_rgb_8uc3_768, (self.color_output_size, self.color_output_size), fx=ratio, fy=ratio, interpolation=cv2.INTER_AREA) # 参考图直接得到 头发 matting 小图 _, ref_matte_pred_8uc1_768 = self.generator_matte.matte_inference(ref_rgb_8uc3_768, ref_landmark_f1k2_768) torch.cuda.empty_cache() ref_matting_8uc3_768 = np.repeat(ref_matte_pred_8uc1_768[:, :, np.newaxis], 3, axis=2) # 用户图得到 头发 matting 小图 _, user_matting_8uc1_bald_768 = self.generator_matte.matte_inference(user_rgb_8uc3_768, user_landmark_f1k2_768) user_matting_8uc3_768 = np.repeat(user_matting_8uc1_bald_768[:, :, np.newaxis], 3, axis=2) return user_rgb_8uc3_768, user_matting_8uc3_768, ref_rgb_8uc3_768, ref_matting_8uc3_768 def get_prepare_hair_color_user_data(self, user_rgb_8uc3_orisize, user_landmark_1k2_f_orisize): """ input： user_rgb_8uc3_orisize: 用户图 原始尺寸 user_landmark_1k2_f_orisize: 用户图 1k 点 output: 图像尺寸基于： 人脸 768， 图像均为3通道 user_rgb_8uc3_change_color_768, user_matting_8uc3_change_color_768, ref_rgb_8uc3_change_color_768, ref_matting_8uc3_change_color_768 user_rgb_8uc3_change_color_768: 用户图 align， size: 768 uint8 （0-255） user_matting_8uc3_change_color_768: 用户图 matting， size: 768 uint8 （0-255） """ # 得到768 小图 M 矩阵 user_hair_color_M = self.get_color_hair_M(user_landmark_1k2_f_orisize) # 1k 点 转换到 768 尺寸 user_landmark_f1k2_768 = landmark_processor.transform_points(user_landmark_1k2_f_orisize, user_hair_color_M) # # 用户图 头发毛躁 warpaffine # usr_ratio = np.sqrt( # (user_hair_color_M[0][0] * user_hair_color_M[0][0]) + (user_hair_color_M[1][0] * user_hair_color_M[1][0])) # user_hair_color_M_tmp = user_hair_color_M / usr_ratio # user_rgb_8uc3_768 = cv2.warpAffine(user_rgb_8uc3_orisize, user_hair_color_M_tmp, # (int(self.color_output_size / usr_ratio), int(self.color_output_size / usr_ratio))) # user_rgb_8uc3_768 = cv2.resize(user_rgb_8uc3_768, (self.color_output_size, self.color_output_size), fx=usr_ratio, fy=usr_ratio, # interpolation=cv2.INTER_AREA) user_rgb_8uc3_768 = cv2.warpAffine(user_rgb_8uc3_orisize, user_hair_color_M, (self.color_output_size, self.color_output_size)) # # usr_ratio = np.sqrt( # (user_hair_color_M[0][0] * user_hair_color_M[0][0]) + (user_hair_color_M[1][0] * user_hair_color_M[1][0])) # user_hair_color_M_tmp = user_hair_color_M / usr_ratio # user_rgb_8uc3_bald_768 = cv2.warpAffine(user_rgb_8uc3_orisize, user_hair_color_M_tmp, # (int(self.color_output_size / usr_ratio), int(self.color_output_size / usr_ratio)), # borderMode=cv2.BORDER_CONSTANT, borderValue=[255, 255, 255]) # user_rgb_8uc3_bald_768 = cv2.resize(user_rgb_8uc3_bald_768, (self.color_output_size, self.color_output_size), fx=usr_ratio, fy=usr_ratio, # interpolation=cv2.INTER_AREA) # 用户图得到 头发 matting 小图 _, user_matting_8uc1_bald_768 = self.generator_matte.matte_inference(user_rgb_8uc3_768, user_landmark_f1k2_768) user_matting_8uc3_768 = np.repeat(user_matting_8uc1_bald_768[:, :, np.newaxis], 3, axis=2) return user_rgb_8uc3_768, user_matting_8uc3_768, user_hair_color_M def get_matte_img(self, img, landmark1k): matte_fg, matte_img = self.generator_matte.matte_inference(img, landmark1k) return matte_fg, matte_img def get_max_countour(self, contours): index_contour = 0 max_num = 0 for i, contour in enumerate(contours): if len(contour) > max_num: max_num = len(contour) index_contour = i return index_contour def judge_hair_pos(self, user_matting_8uc3_bald_768, user_baldseg_8uc3_bald_768, retData): # cv2.imshow("user_matting_8uc3_bald_768", user_matting_8uc3_bald_768) # cv2.imshow("user_baldseg_8uc3_bald_768", user_baldseg_8uc3_bald_768) # cv2.waitKey() cloth_mask = np.zeros_like(user_baldseg_8uc3_bald_768, dtype=np.uint8) cloth_pos = (user_baldseg_8uc3_bald_768[:, :, 0] < 10) & (user_baldseg_8uc3_bald_768[:, :, 1] < 10) & ( user_baldseg_8uc3_bald_768[:, :, 2] > 250) cloth_mask[cloth_pos] = 255 cloth_withhair_mask = (cloth_mask * user_matting_8uc3_bald_768.astype(np.float32) / 255).astype(np.uint8) cloth_withhair_count = len(cloth_withhair_mask[cloth_withhair_mask[:, :, 0] > 0]) # cv2.imshow("cloth_mask", cloth_mask) # cv2.imshow("cloth_withhair_mask", cloth_withhair_mask) # print("cloth_withhair_count: ", cloth_withhair_count / (cloth_withhair_mask.shape[0] * cloth_withhair_mask.shape[1])) # cv2.waitKey() retData["cloth_withhair_count"] = str(cloth_withhair_count) def get_fusion_res(self, user_baldseg_8uc3_512, user_bald_8uc3_512, hair_gene_8uc3_512, user_landmark_f1k2_512): """ input： 图像尺寸基于： 人脸 512， 图像均为3通道 user_baldseg_8uc3_512：用户图 光头分割 mask， uint8 （0-255） user_bald_8uc3_512： 用户图 光头， uint8 （0-255） hair_gene_8uc3_512: 用户图 换发型 贴合 光头， uint8 （0-255） user_landmark_f1k2_512： 用户图 关键点 1k*2 float32 output： hair_gene_fusion_8uc3_512: 生成 发型图， uint8 （0-255） """ # "user_baldseg_8uc3_512, user_bald_8uc3_512, hair_gene_8uc3_512" # cv2.imshow("user_baldseg_8uc3_512", user_baldseg_8uc3_512) # cv2.imshow("user_bald_8uc3_512", user_bald_8uc3_512) # cv2.imshow("hair_gene_8uc3_512", hair_gene_8uc3_512) # 换发型后的 matting 图 hair_gene_matte_fg_8uc3_512, hair_gene_matte_8uc1_512 = self.get_matte_img(hair_gene_8uc3_512, user_landmark_f1k2_512) hair_gene_matte_32fc1_512 = hair_gene_matte_8uc1_512.astype(np.float32) / 255 hair_gene_matte_32fc3_512 = np.repeat(hair_gene_matte_32fc1_512[:, :, np.newaxis], 3, axis=2) hair_bg = (user_bald_8uc3_512 * (1 - hair_gene_matte_32fc3_512)).astype(np.uint8) hair_bg_face = np.zeros_like(user_bald_8uc3_512) face_index = (user_baldseg_8uc3_512[:, :, 1] == 255) & (user_baldseg_8uc3_512[:, :, 0] == 0) & ( user_baldseg_8uc3_512[:, :, 2] == 0) hair_bg_face[face_index] = user_bald_8uc3_512[face_index] hair_face_bg = np.zeros_like(user_bald_8uc3_512) hair_face_bg[face_index] = hair_gene_8uc3_512[face_index] hair_face_mask = np.zeros_like(user_bald_8uc3_512, dtype=np.uint8) hair_face_mask[face_index] = 255 kernel_size = 7 kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size)) hair_face_mask = cv2.erode(hair_face_mask, kernel, iterations=1).astype(np.uint8) hair_face_mask_blur = cv2.blur(hair_face_mask, (5, 5)) # cv2.imshow("hair_face_bg ", hair_face_bg) # cv2.imshow("hair_face_mask ", (hair_face_mask).astype(np.uint8)) # cv2.imshow("hair_face_mask_blur ", hair_face_mask_blur) # cv2.imshow("user_bald_8uc3_512 ", user_bald_8uc3_512) # cv2.imshow("hair_gene_matte_fg_8uc3_512 ", hair_gene_matte_fg_8uc3_512) # paste_hair = np.zeros_like(user_bald_8uc3_512) hair_gene_fusion_8uc3_512 = (user_bald_8uc3_512 * ( 1 - hair_gene_matte_32fc3_512) + hair_gene_matte_fg_8uc3_512 * hair_gene_matte_32fc3_512).astype( np.uint8) # cv2.imshow("hair_gene_fusion_8uc3_512 nouse mask ", hair_gene_fusion_8uc3_512) hair_gene_fusion_8uc3_512 = (hair_gene_fusion_8uc3_512 * ( 1 - hair_face_mask_blur.astype(np.float32) / 255) + hair_face_bg * hair_face_mask_blur.astype( np.float32) / 255).astype(np.uint8) # hair_gene_fusion_8uc3_512[face_index] = hair_face_bg[face_index] # cv2.waitKey() return hair_gene_fusion_8uc3_512 def get_fusion_res_hairpaste(self, user_bald_8uc3_orisize, hair_gene_8uc3_768, user_landmark_f1k2_768, user_hairstyle_M): """ input： 图像尺寸基于： 人脸 512， 图像均为3通道 user_baldseg_8uc3_512：用户图 光头分割 mask， uint8 （0-255） user_bald_8uc3_512： 用户图 光头， uint8 （0-255） hair_gene_8uc3_512: 用户图 换发型 贴合 光头， uint8 （0-255） user_landmark_f1k2_512： 用户图 关键点 1k*2 float32 output： hair_gene_fusion_8uc3_512: 生成 发型图， uint8 （0-255） """ hair_gene_matte_fg_8uc3_768, hair_gene_matte_8uc1_768 = self.get_matte_img(hair_gene_8uc3_768, user_landmark_f1k2_768) # cv2.imshow("hair_gene_8uc3_768", hair_gene_8uc3_768) # cv2.imshow("hair_gene_matte_8uc1_768", hair_gene_matte_8uc1_768) # cv2.imshow("hair_gene_matte_fg_8uc3_768", hair_gene_matte_fg_8uc3_768) # cv2.waitKey() # hair_gene_8uc3_orisize = cv2.warpAffine(hair_gene_8uc3_768, cv2.invertAffineTransform(user_hairstyle_M), # (user_bald_8uc3_orisize.shape[1], # user_bald_8uc3_orisize.shape[0]), # dst=user_bald_8uc3_orisize.copy(), # borderMode=cv2.BORDER_TRANSPARENT) hair_gene_matte_8uc1_orisize = cv2.warpAffine(hair_gene_matte_8uc1_768, cv2.invertAffineTransform(user_hairstyle_M), (user_bald_8uc3_orisize.shape[1], user_bald_8uc3_orisize.shape[0]), flags=cv2.INTER_CUBIC) hair_gene_matte_fg_8uc3_orisize = cv2.warpAffine(hair_gene_matte_fg_8uc3_768, cv2.invertAffineTransform(user_hairstyle_M), (user_bald_8uc3_orisize.shape[1], user_bald_8uc3_orisize.shape[0]), flags=cv2.INTER_CUBIC) hair_gene_matte_8uc3_orisize = np.repeat(hair_gene_matte_8uc1_orisize[:, :, np.newaxis], 3, axis=2) hair_gene_matte_32fc1_orisize = hair_gene_matte_8uc1_orisize.astype(np.float32) / 255 hair_gene_matte_32fc3_orisize = np.repeat(hair_gene_matte_32fc1_orisize[:, :, np.newaxis], 3, axis=2) hair_gene_fusion_8uc3_orisize = (user_bald_8uc3_orisize * ( 1 - hair_gene_matte_32fc3_orisize) + hair_gene_matte_fg_8uc3_orisize * hair_gene_matte_32fc3_orisize).astype( np.uint8) return hair_gene_fusion_8uc3_orisize, hair_gene_matte_8uc3_orisize def get_fusion_res_hairblur(self, user_baldseg_8uc3_orisize, user_bald_8uc3_orisize, hair_gene_8uc3_768, user_landmark_f1k2_768, user_hairstyle_M, ref_hairstyle_dir): """ input： 图像尺寸基于： 人脸 512， 图像均为3通道 user_baldseg_8uc3_512：用户图 光头分割 mask， uint8 （0-255） user_bald_8uc3_512： 用户图 光头， uint8 （0-255） hair_gene_8uc3_512: 用户图 换发型 贴合 光头， uint8 （0-255） user_landmark_f1k2_512： 用户图 关键点 1k*2 float32 output： hair_gene_fusion_8uc3_512: 生成 发型图， uint8 （0-255） """ # cv2.imshow("user_baldseg_8uc3_512", user_baldseg_8uc3_512) # cv2.imshow("user_bald_8uc3_512", user_bald_8uc3_512) # cv2.imshow("hair_gene_8uc3_512", hair_gene_8uc3_512) # 换发型后的 matting 图 hair_gene_matte_fg_8uc3_768, hair_gene_matte_8uc1_768 = self.get_matte_img(hair_gene_8uc3_768, user_landmark_f1k2_768) # cv2.imshow("hair_gene_matte_8uc1_768 before", hair_gene_matte_8uc1_768) # cv2.imshow("hair_gene_matte_fg_8uc3_768 before", hair_gene_matte_fg_8uc3_768) hair_gene_8uc3_orisize = cv2.warpAffine(hair_gene_8uc3_768, cv2.invertAffineTransform(user_hairstyle_M), (user_bald_8uc3_orisize.shape[1], user_bald_8uc3_orisize.shape[0]), dst=user_bald_8uc3_orisize.copy(), borderMode=cv2.BORDER_TRANSPARENT) default_hair_tail_mask = cv2.imread(os.path.join(ref_hairstyle_dir, "default_hair_tail_mask.png")) hair_gene_matte_8uc1_768 = ( hair_gene_matte_8uc1_768 * default_hair_tail_mask[:, :, 0].astype(np.float32) / 255).astype( np.uint8) # cv2.imshow("hair_gene_matte_8uc1_768 after", hair_gene_matte_8uc1_768) hair_gene_matte_8uc1_orisize = cv2.warpAffine(hair_gene_matte_8uc1_768, cv2.invertAffineTransform(user_hairstyle_M), (user_bald_8uc3_orisize.shape[1], user_bald_8uc3_orisize.shape[0]), flags=cv2.INTER_CUBIC) hair_gene_matte_fg_8uc3_768 = ( hair_gene_matte_fg_8uc3_768 * default_hair_tail_mask.astype(np.float32) / 255).astype(np.uint8) # cv2.imshow("hair_gene_matte_fg_8uc3_768 after", hair_gene_matte_fg_8uc3_768) # cv2.waitKey() hair_gene_matte_fg_8uc3_orisize = cv2.warpAffine(hair_gene_matte_fg_8uc3_768, cv2.invertAffineTransform(user_hairstyle_M), (user_bald_8uc3_orisize.shape[1], user_bald_8uc3_orisize.shape[0]), flags=cv2.INTER_CUBIC) hair_gene_matte_32fc1_orisize = hair_gene_matte_8uc1_orisize.astype(np.float32) / 255 hair_gene_matte_32fc3_orisize = np.repeat(hair_gene_matte_32fc1_orisize[:, :, np.newaxis], 3, axis=2) # cv2.imshow("hair_gene_8uc3_512", hair_gene_8uc3_512) # cv2.imshow("hair_gene_matte_8uc1_512", hair_gene_matte_8uc1_512) # cv2.imshow("hair_gene_matte_32fc3_512", (hair_gene_matte_32fc3_512*255).astype(np.uint8)) # cv2.waitKey() hair_bg_face = np.zeros_like(user_bald_8uc3_orisize) face_index = (user_baldseg_8uc3_orisize[:, :, 1] == 255) & (user_baldseg_8uc3_orisize[:, :, 0] == 0) & ( user_baldseg_8uc3_orisize[:, :, 2] == 0) hair_bg_face[face_index] = user_bald_8uc3_orisize[face_index] hair_face_bg = np.zeros_like(user_bald_8uc3_orisize) hair_face_bg[face_index] = hair_gene_8uc3_orisize[face_index] hair_face_mask = np.zeros_like(user_bald_8uc3_orisize, dtype=np.uint8) hair_face_mask[face_index] = 255 hair_face_rec_mask = (((hair_gene_matte_8uc1_orisize > 0) & (hair_face_mask[:, :, 0] > 0)).astype( np.float32) * 255).astype(np.uint8) # cv2.imshow("hair_face_rec_mask no resize", hair_face_rec_mask) hair_face_rec_mask = cv2.resize(hair_face_rec_mask, (768, 768), interpolation=cv2.INTER_CUBIC) # cv2.imshow("hair_face_rec_mask ", np.repeat(hair_face_rec_mask[:, :, np.newaxis], 3, axis=2) ) kernel_size = 11 kernel_dilate = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size)) hair_face_rec_mask_dilate = cv2.erode(hair_face_rec_mask, kernel_dilate, iterations=1) hair_face_rec_mask_blur = cv2.blur(hair_face_rec_mask_dilate, (31, 31)) hair_face_rec_mask_blur = cv2.resize(hair_face_rec_mask_blur, (user_bald_8uc3_orisize.shape[1], user_bald_8uc3_orisize.shape[0]), interpolation=cv2.INTER_CUBIC) # cv2.imshow("hair_face_rec_mask ", (hair_face_rec_mask * 255).astype(np.uint8)) # cv2.imshow("hair_face_rec_mask_blur ", (np.repeat(hair_face_rec_mask_blur[:, :, np.newaxis], 3, axis=2)).astype(np.uint8)) # cv2.waitKey() kernel_size_2 = 7 kernel_erode_2 = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size_2, kernel_size_2)) hair_face_mask_erode_2 = cv2.erode(hair_face_mask, kernel_erode_2, iterations=1).astype(np.uint8) hair_face_mask_blur = cv2.blur(hair_face_mask_erode_2, (5, 5)) hair_face_mask_blur = cv2.resize(hair_face_mask_blur, (user_bald_8uc3_orisize.shape[1], user_bald_8uc3_orisize.shape[0]), interpolation=cv2.INTER_CUBIC) # cv2.imshow("hair_face_bg ", hair_face_bg) # cv2.imshow("hair_face_mask ", (hair_face_mask).astype(np.uint8)) # cv2.imshow("hair_face_mask_rolling ", (hair_face_mask_rolling).astype(np.uint8)) # cv2.imshow("hair_face_mask_use_blur ", (hair_face_mask_use_blur).astype(np.uint8)) # cv2.imshow("hair_face_mask_blur ", hair_face_mask_blur) # cv2.imshow("hair_matte ", (hair_gene_matte_32fc3_512 * 255).astype(np.uint8)) # cv2.imshow("user_bald_8uc3_512 ", user_bald_8uc3_512) # cv2.imshow("hair_gene_matte_fg_8uc3_512 ", hair_gene_matte_fg_8uc3_512) # cv2.imshow("hair_matte blur", (hair_gene_matte_32fc3_512_blur * 255).astype(np.uint8)) # cv2.imshow("user_bald_8uc3_fg ", user_bald_8uc3_fg) # paste_hair = np.zeros_like(user_bald_8uc3_512) hair_gene_fusion_8uc3_orisize = (user_bald_8uc3_orisize * ( 1 - hair_gene_matte_32fc3_orisize) + hair_gene_matte_fg_8uc3_orisize * hair_gene_matte_32fc3_orisize).astype( np.uint8) # cv2.imshow("hair_gene_fusion_8uc3_512 nouse mask ", hair_gene_fusion_8uc3_512) hair_gene_fusion_8uc3_orisize_gene = hair_gene_fusion_8uc3_orisize.copy() hair_gene_fusion_8uc3_orisize_gene = (hair_gene_fusion_8uc3_orisize_gene * ( 1 - hair_face_mask_blur.astype(np.float32) / 255) + hair_face_bg * hair_face_mask_blur.astype( np.float32) / 255).astype(np.uint8) # cv2.imshow("hair_gene_fusion_8uc3_512_gene", hair_gene_fusion_8uc3_512_gene) hair_gene_fusion_8uc3_512 = (hair_gene_fusion_8uc3_orisize_gene * ( hair_face_rec_mask_blur[:, :, np.newaxis].astype(np.float32) / 255) + hair_gene_fusion_8uc3_orisize * ( 1 - hair_face_rec_mask_blur[:, :, np.newaxis].astype( np.float32) / 255)).astype(np.uint8) # cv2.imshow("hair_gene_fusion_8uc3_512 use mask ", hair_gene_fusion_8uc3_512) # cv2.waitKey() return hair_gene_fusion_8uc3_512, hair_gene_8uc3_orisize def get_fusion_res_onlyhair(self, user_baldseg_8uc3_512, user_bald_8uc3_512, hair_gene_8uc3_512, user_landmark_f1k2_512): """ input： 图像尺寸基于： 人脸 512， 图像均为3通道 user_baldseg_8uc3_512：用户图 光头分割 mask， uint8 （0-255） user_bald_8uc3_512： 用户图 光头， uint8 （0-255） hair_gene_8uc3_512: 用户图 换发型 贴合 光头， uint8 （0-255） user_landmark_f1k2_512： 用户图 关键点 1k*2 float32 output： hair_gene_fusion_8uc3_512: 生成 发型图， uint8 （0-255） """ # "user_baldseg_8uc3_512, user_bald_8uc3_512, hair_gene_8uc3_512" cv2.imshow("user_baldseg_8uc3_512", user_baldseg_8uc3_512) # cv2.imshow("user_bald_8uc3_512", user_bald_8uc3_512) cv2.imshow("hair_gene_8uc3_512", hair_gene_8uc3_512) # 换发型后的 matting 图 hair_gene_matte_fg_8uc3_512, hair_gene_matte_8uc1_512 = self.get_matte_img(hair_gene_8uc3_512, user_landmark_f1k2_512) hair_gene_matte_32fc1_512 = hair_gene_matte_8uc1_512.astype(np.float32) / 255 hair_gene_matte_32fc3_512 = np.repeat(hair_gene_matte_32fc1_512[:, :, np.newaxis], 3, axis=2) hair_bg = (user_bald_8uc3_512 * (1 - hair_gene_matte_32fc3_512)).astype(np.uint8) hair_bg_face = np.zeros_like(user_bald_8uc3_512) face_index = (user_baldseg_8uc3_512[:, :, 1] == 255) & (user_baldseg_8uc3_512[:, :, 0] == 0) & ( user_baldseg_8uc3_512[:, :, 2] == 0) hair_bg_face[face_index] = user_bald_8uc3_512[face_index] hair_face_bg = np.zeros_like(user_bald_8uc3_512) hair_face_bg[face_index] = hair_gene_8uc3_512[face_index] hair_face_mask = np.zeros_like(user_bald_8uc3_512, dtype=np.uint8) hair_face_mask[face_index] = 255 kernel_size_1 = 15 kernel_size_2 = 7 kernel_erode_1 = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size_1, kernel_size_1)) kernel_erode_2 = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size_2, kernel_size_2)) hair_face_mask_erode_1 = cv2.erode(hair_face_mask, kernel_erode_1, iterations=1).astype(np.uint8) hair_face_mask_erode_2 = cv2.erode(hair_face_mask, kernel_erode_2, iterations=1).astype(np.uint8) # hair_face_mask_dilate = cv2.dilate(hair_face_mask, kernel, iterations=1).astype(np.uint8) hair_face_mask_rolling = hair_face_mask_erode_2 - hair_face_mask_erode_1 hair_face_mask_use = 255 - hair_face_mask_rolling hair_face_mask_blur = cv2.blur(hair_face_mask_erode_2, (5, 5)) hair_face_mask_use_blur = cv2.blur(hair_face_mask_use, (11, 11)) cv2.imshow("hair_face_bg ", hair_face_bg) cv2.imshow("hair_face_mask ", (hair_face_mask).astype(np.uint8)) cv2.imshow("hair_face_mask_rolling ", (hair_face_mask_rolling).astype(np.uint8)) cv2.imshow("hair_face_mask_use_blur ", (hair_face_mask_use_blur).astype(np.uint8)) cv2.imshow("hair_face_mask_blur ", hair_face_mask_blur) cv2.imshow("hair_matte ", (hair_gene_matte_32fc3_512 * 255).astype(np.uint8)) cv2.imshow("user_bald_8uc3_512 ", user_bald_8uc3_512) cv2.imshow("hair_gene_matte_fg_8uc3_512 ", hair_gene_matte_fg_8uc3_512) hair_gene_matte_32fc3_512_blur = cv2.blur(hair_gene_matte_32fc3_512, (3, 3)) cv2.imshow("hair_matte blur", (hair_gene_matte_32fc3_512_blur * 255).astype(np.uint8)) user_bald_8uc3_white = np.full_like(user_bald_8uc3_512, (187, 202, 240), dtype=(np.uint8)) user_bald_8uc3_fg = (user_bald_8uc3_white * ( 1 - hair_gene_matte_32fc3_512) + hair_gene_matte_fg_8uc3_512 * hair_gene_matte_32fc3_512).astype( np.uint8) cv2.imshow("user_bald_8uc3_fg ", user_bald_8uc3_fg) # paste_hair = np.zeros_like(user_bald_8uc3_512) hair_gene_fusion_8uc3_512 = (user_bald_8uc3_512 * ( 1 - hair_gene_matte_32fc3_512) + hair_gene_matte_fg_8uc3_512 * hair_gene_matte_32fc3_512).astype( np.uint8) cv2.imshow("hair_gene_fusion_8uc3_512 nouse mask ", hair_gene_fusion_8uc3_512) hair_gene_fusion_8uc3_512_gene = hair_gene_fusion_8uc3_512.copy() hair_gene_fusion_8uc3_512_gene = (hair_gene_fusion_8uc3_512_gene * ( 1 - hair_face_mask_blur.astype(np.float32) / 255) + hair_face_bg * hair_face_mask_blur.astype( np.float32) / 255).astype(np.uint8) cv2.imshow("hair_gene_fusion_8uc3_512_gene", hair_gene_fusion_8uc3_512_gene) hair_gene_fusion_8uc3_512 = (hair_gene_fusion_8uc3_512_gene * (1 - hair_face_mask_use_blur.astype( np.float32) / 255) + hair_gene_fusion_8uc3_512 * hair_face_mask_use_blur.astype(np.float32) / 255).astype( np.uint8) cv2.imshow("hair_gene_fusion_8uc3_512 use mask ", hair_gene_fusion_8uc3_512) cv2.waitKey() return hair_gene_fusion_8uc3_512 def get_fusion_res_forehead(self, user_baldseg_8uc3_512, user_bald_8uc3_512, hair_gene_8uc3_512, user_landmark_f1k2_512): """ input： 图像尺寸基于： 人脸 512， 图像均为3通道 user_baldseg_8uc3_512：用户图 光头分割 mask， uint8 （0-255） user_bald_8uc3_512： 用户图 光头， uint8 （0-255） hair_gene_8uc3_512: 用户图 换发型 贴合 光头， uint8 （0-255） user_landmark_f1k2_512： 用户图 关键点 1k*2 float32 output： hair_gene_fusion_8uc3_512: 生成 发型图， uint8 （0-255） """ # 换发型后的 matting 图 hair_gene_matte_fg_8uc3_512, hair_gene_matte_8uc1_512 = self.get_matte_img(hair_gene_8uc3_512, user_landmark_f1k2_512) hair_gene_matte_32fc1_512 = hair_gene_matte_8uc1_512.astype(np.float32) / 255 hair_gene_matte_32fc3_512 = np.repeat(hair_gene_matte_32fc1_512[:, :, np.newaxis], 3, axis=2) hair_bg = (user_bald_8uc3_512 * (1 - hair_gene_matte_32fc3_512)).astype(np.uint8) hair_bg_face = np.zeros_like(user_bald_8uc3_512) face_index = (user_baldseg_8uc3_512[:, :, 1] == 255) & (user_baldseg_8uc3_512[:, :, 0] == 0) & ( user_baldseg_8uc3_512[:, :, 2] == 0) hair_bg_face[face_index] = user_bald_8uc3_512[face_index] hair_face_bg = np.zeros_like(user_bald_8uc3_512) hair_face_bg[face_index] = hair_gene_8uc3_512[face_index] hair_face_mask = np.zeros_like(user_bald_8uc3_512, dtype=np.uint8) hair_face_mask[face_index] = 255 user_landmark_f137_512 = landmark_processor.pts_1k_to_137(user_landmark_f1k2_512) pts_leye_up = user_landmark_f137_512[89:96, :] pts_reye_up = user_landmark_f137_512[106:113, :] pts_leye_up_low = pts_leye_up.min(axis=0)[1] pts_reye_up_low = pts_reye_up.min(axis=0)[1] pts_eye_up_low = min(pts_leye_up_low, pts_reye_up_low) hair_face_mask[int(pts_eye_up_low):, :, :] = 0 kernel_size = 7 kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size)) hair_face_mask = cv2.erode(hair_face_mask, kernel, iterations=1).astype(np.uint8) hair_face_mask_blur = cv2.blur(hair_face_mask, (5, 5)) # cv2.imshow("hair_face_bg ", hair_face_bg) # cv2.imshow("hair_face_mask ", (hair_face_mask).astype(np.uint8)) # cv2.imshow("hair_face_mask_blur ", hair_face_mask_blur) # cv2.imshow("user_bald_8uc3_512 ", user_bald_8uc3_512) # cv2.imshow("hair_gene_matte_fg_8uc3_512 ", hair_gene_matte_fg_8uc3_512) # paste_hair = np.zeros_like(user_bald_8uc3_512) hair_gene_fusion_8uc3_512 = (user_bald_8uc3_512 * ( 1 - hair_gene_matte_32fc3_512) + hair_gene_matte_fg_8uc3_512 * hair_gene_matte_32fc3_512).astype( np.uint8) # cv2.imshow("hair_gene_fusion_8uc3_512 nouse mask ", hair_gene_fusion_8uc3_512) hair_gene_fusion_8uc3_512 = (hair_gene_fusion_8uc3_512 * ( 1 - hair_face_mask_blur.astype(np.float32) / 255) + hair_face_bg * hair_face_mask_blur.astype( np.float32) / 255).astype(np.uint8) # hair_gene_fusion_8uc3_512[face_index] = hair_face_bg[face_index] # cv2.imshow("hair_gene_fusion_8uc3_512 mask ", hair_gene_fusion_8uc3_512) cv2.waitKey() return hair_gene_fusion_8uc3_512 class PersonProcessor_yolov5(object): def __init__(self, gpu_id=0): self.img_size = 640 self.gpu_id = gpu_id self.device = torch.device("cuda:%d" % gpu_id) self.confidence_threshold = 0.5 self.nms_threshold = 0.8 self.model = torch.load(os.path.join(modelRoot, "yolov5l.pt"), map_location=self.device)['model'].float() # load to FP32 # torch.save(torch.load(weights, map_location=device), weights) # update model if SourceChangeWarning # model.fuse() # 修复 Upsample 层缺失属性的问题 for m in self.model.modules(): if isinstance(m, torch.nn.Upsample): if not hasattr(m, 'recompute_scale_factor'): m.recompute_scale_factor = None # 或 False，根据需求 self.model.to(self.device).eval() print("init yolov5 model succeed!!!") def forward(self, image): # Padded resize img = self.letterbox(image, new_shape=self.img_size)[0] # Convert img = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416 img = np.ascontiguousarray(img) ### image is float32 img = torch.from_numpy(img).to(self.device) img = img.float() # uint8 to fp16/32 img /= 255.0 if img.ndimension() == 3: img = img.unsqueeze(0) # Inference pred = self.model(img)[0] # Apply NMS pred = self.non_max_suppression(pred, self.confidence_threshold, self.nms_threshold, classes=None, agnostic=False) boxes = [] scores = [] classes = [] # Process detections det = pred[0] if det is not None and len(det): # Rescale boxes from img_size to im0 size det[:, :4] = self.scale_coords(img.shape[2:], det[:, :4], image.shape).round() det = det.detach().cpu().numpy() for i in range(len(det)): if det[i, 5] == 0: boxes.append([det[i, 0], det[i, 1], det[i, 2], det[i, 3]]) scores.append(det[i, 4]) classes.append(det[i, 5]) # print("det: ", det) boxes = np.array(boxes).astype(np.int32) return { 'boxes': boxes, 'scores': scores, 'classes': classes, } def process(self, in_frame): """ 返回的人体检测框的list，具体rect的格式是[x1, y1, x2, y2] """ res = self.forward(in_frame) if len(res['boxes']) == 0: return None 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 / in_frame.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 None filter_boxes = [item[0] for item in filter_boxes] return filter_boxes def letterbox(self, img, new_shape=(416, 416), color=(114, 114, 114), auto=False, scaleFill=False, scaleup=True): # Resize image to a 32-pixel-multiple rectangle https://github.com/ultralytics/yolov3/issues/232 shape = img.shape[:2] # current shape [height, width] if isinstance(new_shape, int): new_shape = (new_shape, new_shape) # Scale ratio (new / old) r = min(new_shape[0] / shape[0], new_shape[1] / shape[1]) if not scaleup: # only scale down, do not scale up (for better test mAP) r = min(r, 1.0) # Compute padding ratio = r, r # width, height ratios new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r)) dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] # wh padding if auto: # minimum rectangle dw, dh = np.mod(dw, 64), np.mod(dh, 64) # wh padding elif scaleFill: # stretch dw, dh = 0.0, 0.0 new_unpad = new_shape ratio = new_shape[0] / shape[1], new_shape[1] / shape[0] # width, height ratios dw /= 2 # divide padding into 2 sides dh /= 2 if shape[::-1] != new_unpad: # resize img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR) top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1)) left, right = int(round(dw - 0.1)), int(round(dw + 0.1)) img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) # add border return img, ratio, (dw, dh) def scale_coords(self, img1_shape, coords, img0_shape, ratio_pad=None): # Rescale coords (xyxy) from img1_shape to img0_shape if ratio_pad is None: # calculate from img0_shape gain = max(img1_shape) / max(img0_shape) # gain = old / new pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2 # wh padding else: gain = ratio_pad[0][0] pad = ratio_pad[1] coords[:, [0, 2]] -= pad[0] # x padding coords[:, [1, 3]] -= pad[1] # y padding coords[:, :4] /= gain self.clip_coords(coords, img0_shape) return coords def clip_coords(self, boxes, img_shape): # Clip bounding xyxy bounding boxes to image shape (height, width) boxes[:, 0].clamp_(0, img_shape[1]) # x1 boxes[:, 1].clamp_(0, img_shape[0]) # y1 boxes[:, 2].clamp_(0, img_shape[1]) # x2 boxes[:, 3].clamp_(0, img_shape[0]) # y2 def non_max_suppression(self, prediction, conf_thres=0.1, iou_thres=0.6, merge=False, classes=None, agnostic=False): """Performs Non-Maximum Suppression (NMS) on inference results Returns: detections with shape: nx6 (x1, y1, x2, y2, conf, cls) """ if prediction.dtype is torch.float16: prediction = prediction.float() # to FP32 nc = prediction[0].shape[1] - 5 # number of classes xc = prediction[..., 4] > conf_thres # candidates # Settings min_wh, max_wh = 2, 4096 # (pixels) minimum and maximum box width and height max_det = 300 # maximum number of detections per image time_limit = 10.0 # seconds to quit after redundant = True # require redundant detections multi_label = nc > 1 # multiple labels per box (adds 0.5ms/img) t = time.time() output = [None] * prediction.shape[0] for xi, x in enumerate(prediction): # image index, image inference # Apply constraints # x[((x[..., 2:4] < min_wh) | (x[..., 2:4] > max_wh)).any(1), 4] = 0 # width-height x = x[xc[xi]] # confidence # If none remain process next image if not x.shape[0]: continue # Compute conf x[:, 5:] *= x[:, 4:5] # conf = obj_conf * cls_conf # Box (center x, center y, width, height) to (x1, y1, x2, y2) box = self.xywh2xyxy(x[:, :4]) # Detections matrix nx6 (xyxy, conf, cls) if multi_label: i, j = (x[:, 5:] > conf_thres).nonzero().t() x = torch.cat((box[i], x[i, j + 5, None], j[:, None].float()), 1) else: # best class only conf, j = x[:, 5:].max(1, keepdim=True) x = torch.cat((box, conf, j.float()), 1)[conf.view(-1) > conf_thres] # Filter by class if classes: x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)] # Apply finite constraint # if not torch.isfinite(x).all(): # x = x[torch.isfinite(x).all(1)] # If none remain process next image n = x.shape[0] # number of boxes if not n: continue # Sort by confidence # x = x[x[:, 4].argsort(descending=True)] # Batched NMS c = x[:, 5:6] * (0 if agnostic else max_wh) # classes boxes, scores = x[:, :4] + c, x[:, 4] # boxes (offset by class), scores i = torchvision.ops.boxes.nms(boxes, scores, iou_thres) if i.shape[0] > max_det: # limit detections i = i[:max_det] if merge and (1 < n < 3E3): # Merge NMS (boxes merged using weighted mean) try: # update boxes as boxes(i,4) = weights(i,n) * boxes(n,4) iou = self.box_iou(boxes[i], boxes) > iou_thres # iou matrix weights = iou * scores[None] # box weights x[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(1, keepdim=True) # merged boxes if redundant: i = i[iou.sum(1) > 1] # require redundancy except: # possible CUDA error https://github.com/ultralytics/yolov3/issues/1139 print(x, i, x.shape, i.shape) pass output[xi] = x[i] if (time.time() - t) > time_limit: break # time limit exceeded return output def box_iou(self, box1, box2): # https://github.com/pytorch/vision/blob/master/torchvision/ops/boxes.py """ Return intersection-over-union (Jaccard index) of boxes. Both sets of boxes are expected to be in (x1, y1, x2, y2) format. Arguments: box1 (Tensor[N, 4]) box2 (Tensor[M, 4]) Returns: iou (Tensor[N, M]): the NxM matrix containing the pairwise IoU values for every element in boxes1 and boxes2 """ def box_area(box): # box = 4xn return (box[2] - box[0]) * (box[3] - box[1]) area1 = box_area(box1.t()) area2 = box_area(box2.t()) # inter(N,M) = (rb(N,M,2) - lt(N,M,2)).clamp(0).prod(2) inter = (torch.min(box1[:, None, 2:], box2[:, 2:]) - torch.max(box1[:, None, :2], box2[:, :2])).clamp(0).prod(2) return inter / (area1[:, None] + area2 - inter) # iou = inter / (area1 + area2 - inter) def xywh2xyxy(self, x): # Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right y = torch.zeros_like(x) if isinstance(x, torch.Tensor) else np.zeros_like(x) y[:, 0] = x[:, 0] - x[:, 2] / 2 # top left x y[:, 1] = x[:, 1] - x[:, 3] / 2 # top left y y[:, 2] = x[:, 0] + x[:, 2] / 2 # bottom right x y[:, 3] = x[:, 1] + x[:, 3] / 2 # bottom right y return y class KeypointsProcessor(object): def __init__(self, gpu_id=0): import torch from keypoints.lib.config import cfg, update_config from keypoints.lib.models.pose_hrnet import get_pose_net class Namespace: def __init__(self, **kwargs): self.__dict__.update(kwargs) args = Namespace(cfg=os.path.join(modelRoot, 'keypoints/experiments/coco/hrnet/w48_384x288_adam_lr1e-3.yaml'), opts=['TEST.MODEL_FILE', os.path.join(modelRoot, 'pose_hrnet_w48_384x288.pth'), 'TEST.USE_GT_BBOX', 'False'], dataDir='', logDir='', modelDir='', prevModelDir='') update_config(cfg, args) # print(cfg) self.model = get_pose_net(cfg, is_train=False) self.model.eval() if gpu_id != -1: self.model.cuda(gpu_id) if cfg.TEST.MODEL_FILE: if gpu_id == -1: self.model.load_state_dict(torch.load(cfg.TEST.MODEL_FILE, map_location='cpu'), strict=False) else: self.model.load_state_dict(torch.load(cfg.TEST.MODEL_FILE, map_location=torch.device(gpu_id)), strict=False) # print('load keypoints model weights') self.gpu_id = gpu_id self.image_size = [288, 384] self.cfg = cfg @staticmethod def _xywh2cs(x, y, w, h, aspect_ratio=288.0 / 384.0, pixel_std=200): center = np.zeros((2), dtype=np.float32) center[0] = x + w * 0.5 center[1] = y + h * 0.5 if w > aspect_ratio * h: h = w * 1.0 / aspect_ratio elif w < aspect_ratio * h: w = h * aspect_ratio scale = np.array( [w * 1.0 / pixel_std, h * 1.0 / pixel_std], dtype=np.float32) if center[0] != -1: scale = scale * 1.25 return center, scale @staticmethod def get_affine_transform( center, scale, rot, output_size, shift=np.array([0, 0], dtype=np.float32), inv=0 ): def get_dir(src_point, rot_rad): sn, cs = np.sin(rot_rad), np.cos(rot_rad) src_result = [0, 0] src_result[0] = src_point[0] * cs - src_point[1] * sn src_result[1] = src_point[0] * sn + src_point[1] * cs return src_result def get_3rd_point(a, b): direct = a - b return b + np.array([-direct[1], direct[0]], dtype=np.float32) if not isinstance(scale, np.ndarray) and not isinstance(scale, list): scale = np.array([scale, scale]) scale_tmp = scale * 200.0 src_w = scale_tmp[0] dst_w = output_size[0] dst_h = output_size[1] rot_rad = np.pi * rot / 180 src_dir = get_dir([0, src_w * -0.5], rot_rad) dst_dir = np.array([0, dst_w * -0.5], np.float32) src = np.zeros((3, 2), dtype=np.float32) dst = np.zeros((3, 2), dtype=np.float32) src[0, :] = center + scale_tmp * shift src[1, :] = center + src_dir + scale_tmp * shift dst[0, :] = [dst_w * 0.5, dst_h * 0.5] dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5]) + dst_dir src[2:, :] = get_3rd_point(src[0, :], src[1, :]) dst[2:, :] = get_3rd_point(dst[0, :], dst[1, :]) if inv: trans = cv2.getAffineTransform(np.float32(dst), np.float32(src)) else: trans = cv2.getAffineTransform(np.float32(src), np.float32(dst)) return trans def forward(self, image, boxes): from keypoints.lib.core.inference import get_final_preds all_preds = [] for ix, box in enumerate(boxes): x1, y1, x2, y2 = box x, y, w, h, = x1, y1, x2 - x1, y2 - y1 c, s = self._xywh2cs(x, y, w, h) r = 0 trans = self.get_affine_transform(c, s, r, self.image_size) input = cv2.warpAffine( image, trans, (int(self.image_size[0]), int(self.image_size[1])), flags=cv2.INTER_LINEAR) # cv2.imshow('keypoint_input', input) input_rgb = cv2.cvtColor(input, cv2.COLOR_BGR2RGB) input_rgb = np.divide(np.subtract(input_rgb.astype(np.float32) / 255, [0.406, 0.456, 0.485]), [0.225, 0.224, 0.229]) img = np.expand_dims(input_rgb.astype(np.float32).transpose((2, 0, 1)), axis=0) img = torch.from_numpy(img) if self.gpu_id != -1: img = img.cuda(self.gpu_id) output = self.model(img) output_numpy = output.detach().cpu().numpy() c = np.array([c]) s = np.array([s]) preds, maxvals = get_final_preds(self.cfg, output_numpy, c, s) all_preds.append(np.concatenate((preds, maxvals), axis=2)) if len(all_preds) > 0: res = np.concatenate(all_preds, axis=0) else: res = [] return res def forward_keypoints(self, image, keypoints): from keypoints.lib.core.inference import get_final_preds all_preds = [] for ix, keypoint in enumerate(keypoints): contours = [] for pt in (keypoint): x, y, prob = pt if prob > 0.05: contours.append([x, y]) contours = np.array(contours) box = cv2.boundingRect(contours) x, y, w, h, = box h = h * 1.2 w = w * 1.1 c, s = self._xywh2cs(x, y, w, h) r = 0 trans = self.get_affine_transform(c, s, r, self.image_size) input = cv2.warpAffine( image, trans, (int(self.image_size[0]), int(self.image_size[1])), flags=cv2.INTER_LINEAR) # cv2.imshow('keypoint_input', input) input_rgb = cv2.cvtColor(input, cv2.COLOR_BGR2RGB) input_rgb = np.divide(np.subtract(input_rgb.astype(np.float32) / 255, [0.406, 0.456, 0.485]), [0.225, 0.224, 0.229]) img = np.expand_dims(input_rgb.astype(np.float32).transpose((2, 0, 1)), axis=0) img = torch.from_numpy(img) if self.gpu_id != -1: img = img.cuda(self.gpu_id) output = self.model(img) output_numpy = output.detach().cpu().numpy() c = np.array([c]) s = np.array([s]) preds, maxvals = get_final_preds(self.cfg, output_numpy, c, s) all_preds.append(np.concatenate((preds, maxvals), axis=2)) if len(all_preds) > 0: res = np.concatenate(all_preds, axis=0) else: res = [] return res def draw_connect_keypoints(self, keypoints, w, h, draw_line=True): COCO_PERSON_KEYPOINT_NAMES = ( "nose", "left_eye", "right_eye", "left_ear", "right_ear", "left_shoulder", "right_shoulder", "left_elbow", "right_elbow", "left_wrist", "right_wrist", "left_hip", "right_hip", "left_knee", "right_knee", "left_ankle", "right_ankle", ) COCO_PERSON_KEYPOINT_COLORS = ( (102, 204, 255), (51, 153, 255), (102, 0, 204), (51, 102, 255), (153, 255, 204), (128, 229, 255), (153, 255, 153), (102, 255, 224), (255, 102, 0), (255, 255, 77), (153, 255, 204), (191, 255, 128), (255, 195, 77), (77, 204, 22), (22, 139, 77), (0, 56, 138), (138, 76, 23), ) KEYPOINT_CONNECTION_RULES = [ # face ("left_ear", "left_eye", (102, 204, 255)), ("right_ear", "right_eye", (51, 153, 255)), ("left_eye", "nose", (102, 0, 204)), ("nose", "right_eye", (51, 102, 255)), # upper-body ("left_shoulder", "right_shoulder", (255, 128, 0)), ("left_shoulder", "left_elbow", (153, 255, 204)), ("right_shoulder", "right_elbow", (128, 229, 255)), ("left_elbow", "left_wrist", (153, 255, 153)), ("right_elbow", "right_wrist", (102, 255, 224)), # lower-body ("left_hip", "right_hip", (255, 102, 0)), ("left_hip", "left_knee", (255, 255, 77)), ("right_hip", "right_knee", (153, 255, 204)), ("left_knee", "left_ankle", (191, 255, 128)), ("right_knee", "right_ankle", (255, 195, 77)), ] image = np.zeros((h, w, 3), dtype=np.uint8) for ix, keypoint in enumerate(keypoints): visible = {} for idx, pt in enumerate(keypoint): x, y, prob = pt keypoint_name = COCO_PERSON_KEYPOINT_NAMES[idx] if x < 0 or x >= w: continue if y < 0 or y >= h: continue if prob < 0.2: continue visible[keypoint_name] = (int(x), int(y), prob) for ix, (k, v) in enumerate(visible.items()): if not draw_line and ( k == 'nose' or k == 'left_eye' or k == 'right_eye' or k == 'left_ear' or k == 'right_ear'): continue cv2.circle(image, (v[0], v[1]), 10, color=COCO_PERSON_KEYPOINT_COLORS[ix], thickness=cv2.FILLED) # cv2.putText(image, '{}%'.format(int(v[2] * 100)), (int(v[0] + 5), int(v[1])), cv2.FONT_HERSHEY_COMPLEX, 0.5, # (0, 255, 0), 1) if not draw_line: continue for kp0, kp1, color in KEYPOINT_CONNECTION_RULES: if kp0 in visible and kp1 in visible: x0, y0, _ = visible[kp0] x1, y1, _ = visible[kp1] color = (color[2], color[1], color[0]) cv2.line(image, (x0, y0), (x1, y1), color=color, thickness=10, lineType=cv2.LINE_AA) try: ls_x, ls_y, _ = visible["left_shoulder"] rs_x, rs_y, _ = visible["right_shoulder"] mid_shoulder_x, mid_shoulder_y = int((ls_x + rs_x) / 2), int((ls_y + rs_y) / 2) except KeyError: pass else: # draw line from nose to mid-shoulder nose_x, nose_y, _ = visible.get("nose", (None, None, None)) if nose_x is not None: cv2.line(image, (nose_x, nose_y), (mid_shoulder_x, mid_shoulder_y), color=(0, 0, 255), thickness=10, lineType=cv2.LINE_AA) try: # draw line from mid-shoulder to mid-hip lh_x, lh_y, _ = visible["left_hip"] rh_x, rh_y, _ = visible["right_hip"] except KeyError: pass else: mid_hip_x, mid_hip_y = int((lh_x + rh_x) / 2), int((lh_y + rh_y) / 2) cv2.line(image, (mid_hip_x, mid_hip_y), (mid_shoulder_x, mid_shoulder_y), color=(0, 0, 255), thickness=10, lineType=cv2.LINE_AA) return image class Generator_Hair(object): def __init__(self, gpu, device_id): 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.model_dir = modelRoot self.dst_height = 768 self.dst_width = 768 generator_hair_path = os.path.join(self.model_dir, "master_hair_v8_onlyhair_nowarp_all_0709.pt") # default master_hair_v8_onlyhair_blur_aug_0413.pt master_hair_v8_onlyhair_min_0223 # "master_hair_v8_onlyhair_nowarp_all_0709" # master_hair_v8_onlyhair_blur_aug_all_0703 self.net = torch.jit.load(generator_hair_path, torch.device('cpu')).to(self.device) # print("generate hair net: ", self.net) self.net.eval() def Generator_Hair_inference(self, ref_rgb_8uc3_512, ref_matting_8uc3_512, ref_baldseg_8uc3_512, ref_landmark_f1k2_512, user_baldseg_8uc3_512, user_bald_8uc3_512, user_landmark_f1k2_512): """ 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 user_baldseg_8uc3_512：用户图 光头分割 mask， uint8 （0-255） user_bald_8uc3_512： 用户图 光头， uint8 （0-255） user_landmark_f1k2_512： 用户图 关键点 1k*2 float32 output： hair_gene_8uc3_512: 生成 发型图， uint8 （0-255） """ nohair_mask = user_baldseg_8uc3_512.copy() user_bald_mask = (user_baldseg_8uc3_512 == [0, 255, 0]).all(axis=2) user_pts137 = landmark_processor.pts_1k_to_137(user_landmark_f1k2_512).astype(np.int32) # 2 ********************** nohair_mask ********************** # Label mouth cv2.fillPoly(nohair_mask, np.concatenate((user_pts137[47:35:-1], user_pts137[56:64], [user_pts137[48], user_pts137[22]]))[ np.newaxis, :, :], (255, 255, 0)) cv2.fillPoly(nohair_mask, np.concatenate((user_pts137[22:37], user_pts137[56:47:-1]))[np.newaxis, :, :], (255, 0, 128)) cv2.fillPoly(nohair_mask, user_pts137[88:104][np.newaxis, :, :], (255, 0, 255)) # eye cv2.fillPoly(nohair_mask, user_pts137[105:121][np.newaxis, :, :], (255, 0, 255)) # eye # Label nose cv2.fillPoly(nohair_mask, user_pts137[64:79][np.newaxis, :, :], (255, 255, 255)) # Label eyebrow cv2.fillPoly(nohair_mask, user_pts137[121:129][np.newaxis, :, :], (255, 128, 0)) cv2.fillPoly(nohair_mask, user_pts137[129:137][np.newaxis, :, :], (255, 128, 0)) # 5 ********************** bald_img ********************** bald_img = (user_bald_mask[:, :, np.newaxis] * user_bald_8uc3_512).astype(np.uint8) # 8 ********************** another pose hair_image ********************** another_pose_image = ref_rgb_8uc3_512.copy() another_nohair_pose_mask = ref_baldseg_8uc3_512.copy() 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_512).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_512.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_nohair_mask = torch.from_numpy( (nohair_mask.astype(np.float32) / 255).transpose(2, 0, 1)[np.newaxis, :, :, :]).to(self.device) # input_nohair_image = torch.from_numpy( (nohair_image.astype(np.float32) / 255).transpose(2, 0, 1)[np.newaxis, :, :, :]).to(self.device) input_bald_img = torch.from_numpy( (bald_img.astype(np.float32) / 255).transpose(2, 0, 1)[np.newaxis, :, :, :]).to(self.device) input_another_pose_hair_image = torch.from_numpy( (another_pose_hair_image.astype(np.float32) / 255).transpose(2, 0, 1)[np.newaxis, :, :, :]).to(self.device) zero_tensor = torch.from_numpy(np.zeros((1, 3, self.dst_height, self.dst_width), dtype=np.float32)).to( self.device) with torch.no_grad(): fake_image = self.net(input_nohair_mask, input_bald_img, input_another_pose_hair_image, zero_tensor) fake_image_numpy = (fake_image[0].cpu().detach().numpy().transpose((1, 2, 0)) * 255).astype(np.uint8).copy() return fake_image_numpy def Generator_Hair_inference_use_pref(self, another_pose_hair_image, user_baldseg_8uc3_768, user_bald_8uc3_768, user_landmark_f1k2_768, gender="boy"): """ input： 图像尺寸基于： 人脸 512， 图像均为3通道 another_pose_hair_image: 参考图 处理好 numpy，float32 user_baldseg_8uc3_512：用户图 光头分割 mask， uint8 （0-255） user_bald_8uc3_512： 用户图 光头， uint8 （0-255） user_landmark_f1k2_512： 用户图 关键点 1k*2 float32 output： hair_gene_8uc3_512: 生成 发型图， uint8 （0-255） """ nohair_mask = user_baldseg_8uc3_768.copy() # user_bald_mask = (user_baldseg_8uc3_768 == [0, 255, 0]).all(axis=2) user_pts137 = landmark_processor.pts_1k_to_137(user_landmark_f1k2_768).astype(np.int32) # 2 ********************** nohair_mask ********************** # Label mouth cv2.fillPoly(nohair_mask, np.concatenate((user_pts137[47:35:-1], user_pts137[56:64], [user_pts137[48], user_pts137[22]]))[ np.newaxis, :, :], (255, 255, 0)) cv2.fillPoly(nohair_mask, np.concatenate((user_pts137[22:37], user_pts137[56:47:-1]))[np.newaxis, :, :], (255, 0, 128)) cv2.fillPoly(nohair_mask, user_pts137[88:104][np.newaxis, :, :], (255, 0, 255)) # eye cv2.fillPoly(nohair_mask, user_pts137[105:121][np.newaxis, :, :], (255, 0, 255)) # eye # Label nose cv2.fillPoly(nohair_mask, user_pts137[64:79][np.newaxis, :, :], (255, 255, 255)) # Label eyebrow cv2.fillPoly(nohair_mask, user_pts137[121:129][np.newaxis, :, :], (255, 128, 0)) cv2.fillPoly(nohair_mask, user_pts137[129:137][np.newaxis, :, :], (255, 128, 0)) # # 4 ********************** nohair_image ********************** nohair_image_orig_blur = cv2.blur(user_bald_8uc3_768, (self.dst_width // 5, self.dst_width // 5)) # bald_img = nohair_image_orig_blur.copy() clear_mask = np.zeros((self.dst_height, self.dst_width, 3), dtype=np.uint8) cv2.fillPoly(clear_mask, np.concatenate((user_pts137[96:88:-1], user_pts137[105:114], user_pts137[2::-1], user_pts137[21:19:-1]))[np.newaxis, :, :], (1, 1, 1)) bald_img = nohair_image_orig_blur * (1 - clear_mask) + user_bald_8uc3_768 * clear_mask # inter_res = np.concatenate((user_baldseg_8uc3_768, user_bald_8uc3_768, bald_img), axis=1) # cv2.imshow("gen hair inter_res: ", inter_res) # cv2.waitKey() # 改动 for blur version # nohair_image_orig_blur = cv2.blur(user_bald_8uc3_768, (self.dst_width // 5, self.dst_width // 5)) # bald_img = nohair_image_orig_blur.copy() input_nohair_mask = torch.from_numpy((nohair_mask.astype(np.float32) / 255).transpose(2, 0, 1)[np.newaxis, :, :, :]).to(self.device) input_bald_img = torch.from_numpy((bald_img.astype(np.float32) / 255).transpose(2, 0, 1)[np.newaxis, :, :, :]).to(self.device) input_another_pose_hair_image = torch.from_numpy(another_pose_hair_image.transpose(2, 0, 1)[np.newaxis, :, :, :]).to(self.device) zero_tensor = torch.from_numpy(np.zeros((1, 3, self.dst_height, self.dst_width), dtype=np.float32)).to( self.device) with torch.no_grad(): fake_image = self.net(input_nohair_mask, input_bald_img, input_another_pose_hair_image, zero_tensor) fake_image_numpy = (fake_image[0].cpu().detach().numpy().transpose((1, 2, 0)) * 255).astype(np.uint8).copy() # hair__align_concat_show = np.concatenate((nohair_mask, bald_img, (another_pose_hair_image*255).astype(np.uint8), fake_image_numpy), axis=1) # hair__align_concat_show = cv2.resize(hair__align_concat_show, (0, 0), fx=0.5, fy=0.5) # cv2.imshow("hair__align_concat_show", hair__align_concat_show) # cv2.waitKey(0) return fake_image_numpy class Generator_Fusion_Res(Process_Data): def __init__(self, gpu, device_id): 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.model_dir = modelRoot self.output_img_size = 768 fusion_model_path = os.path.join(self.model_dir, 'hair_fusion_0427.pt') # v1: 0412 v2:zxm_v2_use_paired_data_remap_add_encode_04_12_768 self.load_fusion_model(fusion_model_path) def load_fusion_model(self, hair_fusion_model_path): self.fusion_model = torch.jit.load(hair_fusion_model_path, map_location='cpu').to(self.device) def inference_girl(self, user_generator_hair_8uc3_orisize, user_generator_matte_8uc3_orisize, user_generator_landmark_8uc3_orisize, user_bald_mask_8uc3_orisize): """ input: user_generator_hair_8uc3_orisize: 换发型结果 原图尺寸 8uc3 user_generator_matte_8uc3_orisize: 换发型结果matting alpha图 原图尺寸 8uc3 output: fusion_res """ ratio = 0.5 h_offset = 0.45 user_pts1k = user_generator_landmark_8uc3_orisize.astype(np.int32) user_real = user_generator_hair_8uc3_orisize user_hair_mask = user_generator_matte_8uc3_orisize # user_bald_mask_8uc3_orisize = cv2.imread("/media/DATA_4T/test_hair/debug_fusion/5c_mask_xxq/00B4D4B8-8F26-7DF9-51D0-41492D30DDF720201230_boy_20_None_0.png") # cv2.imshow("user_bald_mask_8uc3_orisize", user_bald_mask_8uc3_orisize) # cv2.imshow("user_generator_hair_8uc3_orisize", user_generator_hair_8uc3_orisize) # cv2.imshow("user_generator_matte_8uc3_orisize", user_generator_matte_8uc3_orisize) # cv2.waitKey() user_bald_mask = user_bald_mask_8uc3_orisize user_hair_mask_orig = user_hair_mask[:, :, 0].copy() user_real_save = user_real.copy() orig_h, orig_w, _ = user_real_save.shape M_user = landmark_processor.get_transform_mat_hair_ratio_v1(user_pts1k, self.output_img_size, ratio, h_offset) inv_M_user = cv2.invertAffineTransform(M_user) user_real = cv2.warpAffine(user_real, M_user, (self.output_img_size, self.output_img_size)) user_hair_mask = cv2.warpAffine(user_hair_mask, M_user, (self.output_img_size, self.output_img_size), flags=cv2.INTER_CUBIC) user_bald_mask = cv2.warpAffine(user_bald_mask, M_user, (self.output_img_size, self.output_img_size), flags=cv2.INTER_NEAREST) random_v = 15 kernel_e = np.ones((random_v, random_v), np.uint8) kernel_d = np.ones((random_v + 10, random_v + 10), np.uint8) user_hair_mask_erode = cv2.erode(user_hair_mask[:, :, 0], kernel_e, iterations=1) user_hair_mask = cv2.dilate(user_hair_mask[:, :, 0], kernel_d, iterations=1) blend = np.clip(user_real * (user_hair_mask_erode[:, :, np.newaxis] / 255), 0, 255) lab_img_paf = cv2.cvtColor(user_real, cv2.COLOR_BGR2LAB) lab_img_paf[:, :, 1] = 0 lab_img_paf[:, :, 2] = 0 # img_paf = np.clip(lab_img_paf * (user_hair_mask[:, :, np.newaxis] / 255) + user_real * ( # 1 - user_hair_mask[:, :, np.newaxis] / 255), 0, 255) img_paf = user_real.copy() img_paf[user_hair_mask > 0] = lab_img_paf[user_hair_mask > 0] # cv2.imshow("blend", blend.astype(np.uint8)) # cv2.imshow("user_hair_mask", user_hair_mask.astype(np.uint8)) # cv2.imshow("img_paf", img_paf.astype(np.uint8)) # cv2.waitKey() # img_paf = cv2.imread("/media/DATA_4T/test_hair/debug_fusion/user_debug/fusion_test/img_paf.png") # blend = cv2.imread("/media/DATA_4T/test_hair/debug_fusion/user_debug/fusion_test/blend.png") condition = blend # np.concatenate((user_mask, user_hair_mask), axis=2) condition = (condition.astype(np.float32) / 255).transpose((2, 0, 1))[np.newaxis, :, :, :] input_paf = (img_paf.astype(np.float32) / 255).transpose((2, 0, 1))[np.newaxis, :, :, :] condition = torch.from_numpy(condition).to(self.device) input_paf = torch.from_numpy(input_paf).to(self.device) # ******************* forward ******************* test_fake = self.fusion_model(input_paf, condition) test_res = (test_fake.detach()[0].to('cpu').numpy().transpose((1, 2, 0)) * 255).astype(np.uint8).copy() input_out_show = np.concatenate((img_paf, user_bald_mask, (blend).astype(np.uint8), test_res), axis=1) # cv2.imshow("input_out_show", input_out_show) # cv2.waitKey() # cv2.imwrite("/media/DATA_4T/test_hair/debug_fusion/user_debug/test_new_res.png", test_res) test_res_orig = user_real_save.copy() user_hair_mask_e_orig = np.zeros(user_real_save[:, :, 0].shape, dtype=np.uint8) user_hair_mask_d_orig = np.zeros(user_real_save[:, :, 0].shape, dtype=np.uint8) cv2.warpAffine(test_res, inv_M_user, (orig_w, orig_h), dst=test_res_orig, borderMode=cv2.BORDER_TRANSPARENT) cv2.warpAffine(user_hair_mask_erode, inv_M_user, (orig_w, orig_h), dst=user_hair_mask_e_orig) cv2.warpAffine(user_hair_mask, inv_M_user, (orig_w, orig_h), dst=user_hair_mask_d_orig) # cv2.imshow("user_hair_mask", user_hair_mask) # cv2.imshow("user_hair_mask_d_orig", user_hair_mask_d_orig) # cv2.imshow("user_hair_mask_erode", user_hair_mask_erode) # cv2.imshow("user_hair_mask_e_orig", user_hair_mask_e_orig) # cv2.waitKey() user_bald_mask_b_orig = np.zeros(user_real_save[:, :, 0].shape, dtype=np.uint8) user_bald_mask_b = np.zeros(user_bald_mask[:, :, 0].shape, dtype=np.uint8) user_bald_mask_b[user_bald_mask[:, :, 0] > 0] = 1 user_bald_mask_b[user_bald_mask[:, :, 1] > 0] = 1 user_bald_mask_b[user_bald_mask[:, :, 2] > 0] = 1 user_bald_mask_b = cv2.dilate(user_bald_mask_b, np.ones((25, 25))) cv2.warpAffine(user_bald_mask_b, inv_M_user, (orig_w, orig_h), dst=user_bald_mask_b_orig) # user_hair_mask_d_orig_b = cv2.blur(user_hair_mask_d_orig, (30, 30)) user_hair_mask_d_orig_b = user_hair_mask_d_orig user_hair_mask_e_orig_b = cv2.blur(user_hair_mask_e_orig, (15, 15)) # user_hair_mask_d_orig[~(user_bald_mask_b_orig.astype(bool))] = user_hair_mask_e_orig[~(user_bald_mask_b_orig.astype(bool))] # # user_hair_mask_d_orig_b = cv2.blur(user_hair_mask_d_orig, (15, 15)) user_hair_mask_d_orig_b[~(user_bald_mask_b_orig.astype(bool))] = user_hair_mask_e_orig[ ~(user_bald_mask_b_orig.astype(bool))] user_hair_mask_d_orig_b = cv2.blur(user_hair_mask_d_orig_b, (15, 15)) test_res_clear_orig = user_real_save * ( 1 - user_hair_mask_d_orig_b[:, :, np.newaxis] / 255) + test_res_orig * ( user_hair_mask_d_orig_b[:, :, np.newaxis] / 255) test_res_clear_orig = (np.clip(test_res_clear_orig, 0, 255)).astype(np.uint8) test_res_orig_lab = cv2.cvtColor(test_res_clear_orig, cv2.COLOR_BGR2LAB) user_real_save_lab = cv2.cvtColor(user_real_save, cv2.COLOR_BGR2LAB) test_res_orig_lab[:, :, 0] = user_real_save_lab[:, :, 0] fusion_res = cv2.cvtColor(test_res_orig_lab, cv2.COLOR_LAB2BGR) user_hair_mask_e_orig_b[~(user_bald_mask_b_orig.astype(bool))] = user_hair_mask_d_orig_b[~(user_bald_mask_b_orig.astype(bool))] fusion_res = test_res_clear_orig * ( 1 - user_hair_mask_e_orig_b[:, :, np.newaxis] / 255) + fusion_res * ( user_hair_mask_e_orig_b[:, :, np.newaxis] / 255) fusion_res = (np.clip(fusion_res, 0, 255)).astype(np.uint8) # con = np.concatenate((user_hair_mask_orig, user_hair_mask_d_orig_b), axis=1) # con = cv2.resize(con, (0, 0), fx=1 / 4, fy=1 / 4, interpolation=cv2.INTER_CUBIC) # con_res = np.concatenate(( user_real_save, fusion_res), axis=1) # con_res = cv2.resize(con_res, (0, 0), fx=1 / 4, fy=1 / 4, interpolation=cv2.INTER_CUBIC) # # cv2.imshow("con", con) # cv2.imshow("con_res", con_res) # cv2.waitKey() # cv2.imwrite("/media/DATA_4T/test_hair/debug_fusion/user_debug/con_new1_fuison.png", con_res) return fusion_res def inference_boy(self, user_generator_hair_8uc3_orisize, user_generator_matte_8uc3_orisize, user_generator_landmark_8uc3_orisize, user_bald_mask_8uc3_orisize): """ input: user_generator_hair_8uc3_orisize: 换发型结果 原图尺寸 8uc3 user_generator_matte_8uc3_orisize: 换发型结果matting alpha图 原图尺寸 8uc3 output: fusion_res """ ratio = 0.6 h_offset = 0.65 user_pts1k = user_generator_landmark_8uc3_orisize.astype(np.int32) user_real = user_generator_hair_8uc3_orisize user_hair_mask = user_generator_matte_8uc3_orisize # user_bald_mask_8uc3_orisize = cv2.imread("/media/DATA_4T/test_hair/debug_fusion/5c_mask_xxq/00B4D4B8-8F26-7DF9-51D0-41492D30DDF720201230_boy_20_None_0.png") # cv2.imshow("user_bald_mask_8uc3_orisize", user_bald_mask_8uc3_orisize) # cv2.imshow("user_generator_hair_8uc3_orisize", user_generator_hair_8uc3_orisize) # cv2.imshow("user_generator_matte_8uc3_orisize", user_generator_matte_8uc3_orisize) # cv2.waitKey() user_bald_mask = user_bald_mask_8uc3_orisize user_hair_mask_orig = user_hair_mask[:, :, 0].copy() user_real_save = user_real.copy() orig_h, orig_w, _ = user_real_save.shape M_user = landmark_processor.get_transform_mat_hair_ratio_v1(user_pts1k, self.output_img_size, ratio, h_offset) inv_M_user = cv2.invertAffineTransform(M_user) user_real = cv2.warpAffine(user_real, M_user, (self.output_img_size, self.output_img_size)) user_hair_mask = cv2.warpAffine(user_hair_mask, M_user, (self.output_img_size, self.output_img_size), flags=cv2.INTER_CUBIC) user_bald_mask = cv2.warpAffine(user_bald_mask, M_user, (self.output_img_size, self.output_img_size), flags=cv2.INTER_NEAREST) random_v = 15 kernel_e = np.ones((random_v, random_v), np.uint8) kernel_d = np.ones((random_v + 10, random_v + 10), np.uint8) user_hair_mask_erode = cv2.erode(user_hair_mask[:, :, 0], kernel_e, iterations=1) user_hair_mask = cv2.dilate(user_hair_mask[:, :, 0], kernel_d, iterations=1) blend = np.clip(user_real * (user_hair_mask_erode[:, :, np.newaxis] / 255), 0, 255) lab_img_paf = cv2.cvtColor(user_real, cv2.COLOR_BGR2LAB) lab_img_paf[:, :, 1] = 0 lab_img_paf[:, :, 2] = 0 # img_paf = np.clip(lab_img_paf * (user_hair_mask[:, :, np.newaxis] / 255) + user_real * ( # 1 - user_hair_mask[:, :, np.newaxis] / 255), 0, 255) img_paf = user_real.copy() img_paf[user_hair_mask > 0] = lab_img_paf[user_hair_mask > 0] # cv2.imshow("blend", blend.astype(np.uint8)) # cv2.imshow("user_hair_mask", user_hair_mask.astype(np.uint8)) # cv2.imshow("img_paf", img_paf.astype(np.uint8)) # cv2.waitKey() # img_paf = cv2.imread("/media/DATA_4T/test_hair/debug_fusion/user_debug/fusion_test/img_paf.png") # blend = cv2.imread("/media/DATA_4T/test_hair/debug_fusion/user_debug/fusion_test/blend.png") condition = blend # np.concatenate((user_mask, user_hair_mask), axis=2) condition = (condition.astype(np.float32) / 255).transpose((2, 0, 1))[np.newaxis, :, :, :] input_paf = (img_paf.astype(np.float32) / 255).transpose((2, 0, 1))[np.newaxis, :, :, :] condition = torch.from_numpy(condition).to(self.device) input_paf = torch.from_numpy(input_paf).to(self.device) # ******************* forward ******************* test_fake = self.fusion_model(input_paf, condition) test_res = (test_fake.detach()[0].to('cpu').numpy().transpose((1, 2, 0)) * 255).astype(np.uint8).copy() input_out_show = np.concatenate((img_paf, user_bald_mask, (blend).astype(np.uint8), test_res), axis=1) # cv2.imshow("input_out_show", input_out_show) # cv2.waitKey() # cv2.imwrite("/media/DATA_4T/test_hair/debug_fusion/user_debug/test_new_res.png", test_res) test_res_orig = user_real_save.copy() user_hair_mask_e_orig = np.zeros(user_real_save[:, :, 0].shape, dtype=np.uint8) user_hair_mask_d_orig = np.zeros(user_real_save[:, :, 0].shape, dtype=np.uint8) cv2.warpAffine(test_res, inv_M_user, (orig_w, orig_h), dst=test_res_orig, borderMode=cv2.BORDER_TRANSPARENT) cv2.warpAffine(user_hair_mask_erode, inv_M_user, (orig_w, orig_h), dst=user_hair_mask_e_orig) cv2.warpAffine(user_hair_mask, inv_M_user, (orig_w, orig_h), dst=user_hair_mask_d_orig) # cv2.imshow("user_hair_mask", user_hair_mask) # cv2.imshow("user_hair_mask_d_orig", user_hair_mask_d_orig) # cv2.imshow("user_hair_mask_erode", user_hair_mask_erode) # cv2.imshow("user_hair_mask_e_orig", user_hair_mask_e_orig) # cv2.waitKey() user_bald_mask_b_orig = np.zeros(user_real_save[:, :, 0].shape, dtype=np.uint8) user_bald_mask_b = np.zeros(user_bald_mask[:, :, 0].shape, dtype=np.uint8) user_bald_mask_b[user_bald_mask[:, :, 0] > 0] = 1 user_bald_mask_b[user_bald_mask[:, :, 1] > 0] = 1 user_bald_mask_b[user_bald_mask[:, :, 2] > 0] = 1 user_bald_mask_b = cv2.dilate(user_bald_mask_b, np.ones((10, 10))) cv2.warpAffine(user_bald_mask_b, inv_M_user, (orig_w, orig_h), dst=user_bald_mask_b_orig) user_hair_mask_d_orig_b = user_hair_mask_d_orig user_hair_mask_d_orig_b[~(user_bald_mask_b_orig.astype(bool))] = user_hair_mask_e_orig[ ~(user_bald_mask_b_orig.astype(bool))] user_hair_mask_d_orig_b = cv2.blur(user_hair_mask_d_orig_b, (15, 15)) user_hair_mask_e_orig_b = cv2.blur(user_hair_mask_e_orig, (15, 15)) test_res_clear_orig = user_real_save * ( 1 - user_hair_mask_d_orig_b[:, :, np.newaxis] / 255) + test_res_orig * ( user_hair_mask_d_orig_b[:, :, np.newaxis] / 255) test_res_clear_orig = (np.clip(test_res_clear_orig, 0, 255)).astype(np.uint8) test_res_orig_lab = cv2.cvtColor(test_res_clear_orig, cv2.COLOR_BGR2LAB) user_real_save_lab = cv2.cvtColor(user_real_save, cv2.COLOR_BGR2LAB) test_res_orig_lab[:, :, 0] = user_real_save_lab[:, :, 0] fusion_res = cv2.cvtColor(test_res_orig_lab, cv2.COLOR_LAB2BGR) user_hair_mask_e_orig_b[~(user_bald_mask_b_orig.astype(bool))] = user_hair_mask_d_orig_b[ ~(user_bald_mask_b_orig.astype(bool))] fusion_res = test_res_clear_orig * ( 1 - user_hair_mask_e_orig_b[:, :, np.newaxis] / 255) + fusion_res * ( user_hair_mask_e_orig_b[:, :, np.newaxis] / 255) fusion_res = (np.clip(fusion_res, 0, 255)).astype(np.uint8) # con = np.concatenate((user_hair_mask_orig, user_hair_mask_d_orig_b), axis=1) # con = cv2.resize(con, (0, 0), fx=1 / 4, fy=1 / 4, interpolation=cv2.INTER_CUBIC) # # con_res = np.concatenate((user_real_save, fusion_res), axis=1) # con_res = cv2.resize(con_res, (0, 0), fx=1 / 4, fy=1 / 4, interpolation=cv2.INTER_CUBIC) # # cv2.imshow("con", con) # cv2.imshow("con_res", con_res) # cv2.waitKey() return fusion_res def inference(self, user_generator_hair_8uc3_orisize, user_generator_matte_8uc3_orisize, user_generator_landmark_8uc3_orisize, user_bald_mask_8uc3_orisize, ratio): """ input: user_generator_hair_8uc3_orisize: 换发型结果 原图尺寸 8uc3 user_generator_matte_8uc3_orisize: 换发型结果matting alpha图 原图尺寸 8uc3 output: fusion_res """ user_pts1k = user_generator_landmark_8uc3_orisize.astype(np.int32) user_real = user_generator_hair_8uc3_orisize user_hair_mask = user_generator_matte_8uc3_orisize # user_bald_mask_8uc3_orisize = cv2.imread("/media/DATA_4T/test_hair/debug_fusion/5c_mask_xxq/00B4D4B8-8F26-7DF9-51D0-41492D30DDF720201230_boy_20_None_0.png") # cv2.imshow("user_bald_mask_8uc3_orisize", user_bald_mask_8uc3_orisize) # cv2.imshow("user_generator_hair_8uc3_orisize", user_generator_hair_8uc3_orisize) # cv2.imshow("user_generator_matte_8uc3_orisize", user_generator_matte_8uc3_orisize) # cv2.waitKey() user_bald_mask = user_bald_mask_8uc3_orisize user_hair_mask_orig = user_hair_mask[:, :, 0].copy() user_real_save = user_real.copy() orig_h, orig_w, _ = user_real_save.shape if ratio == 0: M_user = self.get_hair_M_boy_v1(user_pts1k) elif ratio == 1: M_user = self.get_hair_M_girl_v1(user_pts1k) elif ratio == 2: M_user = self.get_hair_M_girl_v2(user_pts1k) else: M_user = self.get_hair_M_girl_v1(user_pts1k) # M_user = landmark_processor.get_transform_mat_hair_ratio_v1(user_pts1k, self.output_img_size, ratio, # h_offset) inv_M_user = cv2.invertAffineTransform(M_user) user_real = cv2.warpAffine(user_real, M_user, (self.output_img_size, self.output_img_size)) user_hair_mask = cv2.warpAffine(user_hair_mask, M_user, (self.output_img_size, self.output_img_size), flags=cv2.INTER_CUBIC) user_bald_mask = cv2.warpAffine(user_bald_mask, M_user, (self.output_img_size, self.output_img_size), flags=cv2.INTER_NEAREST) random_v = 15 kernel_e = np.ones((random_v, random_v), np.uint8) kernel_d = np.ones((random_v + 10, random_v + 10), np.uint8) user_hair_mask_erode = cv2.erode(user_hair_mask[:, :, 0], kernel_e, iterations=1) user_hair_mask = cv2.dilate(user_hair_mask[:, :, 0], kernel_d, iterations=1) blend = np.clip(user_real * (user_hair_mask_erode[:, :, np.newaxis] / 255), 0, 255) lab_img_paf = cv2.cvtColor(user_real, cv2.COLOR_BGR2LAB) lab_img_paf[:, :, 1] = 0 lab_img_paf[:, :, 2] = 0 # img_paf = np.clip(lab_img_paf * (user_hair_mask[:, :, np.newaxis] / 255) + user_real * ( # 1 - user_hair_mask[:, :, np.newaxis] / 255), 0, 255) img_paf = user_real.copy() img_paf[user_hair_mask > 0] = lab_img_paf[user_hair_mask > 0] # cv2.imshow("blend", blend.astype(np.uint8)) # cv2.imshow("user_hair_mask", user_hair_mask.astype(np.uint8)) # cv2.imshow("img_paf", img_paf.astype(np.uint8)) # cv2.waitKey() # img_paf = cv2.imread("/media/DATA_4T/test_hair/debug_fusion/user_debug/fusion_test/img_paf.png") # blend = cv2.imread("/media/DATA_4T/test_hair/debug_fusion/user_debug/fusion_test/blend.png") condition = blend # np.concatenate((user_mask, user_hair_mask), axis=2) condition = (condition.astype(np.float32) / 255).transpose((2, 0, 1))[np.newaxis, :, :, :] input_paf = (img_paf.astype(np.float32) / 255).transpose((2, 0, 1))[np.newaxis, :, :, :] condition = torch.from_numpy(condition).to(self.device) input_paf = torch.from_numpy(input_paf).to(self.device) # ******************* forward ******************* test_fake = self.fusion_model(input_paf, condition) test_res = (test_fake.detach()[0].to('cpu').numpy().transpose((1, 2, 0)) * 255).astype(np.uint8).copy() input_out_show = np.concatenate((img_paf, user_bald_mask, (blend).astype(np.uint8), test_res), axis=1) ratio = 1536. / max(input_out_show.shape[:2]) input_out_show = cv2.resize(input_out_show, (0, 0), fx=ratio, fy=ratio) # cv2.imshow("fusion_input_out_show", input_out_show) # cv2.waitKey() test_res_orig = user_real_save.copy() user_hair_mask_e_orig = np.zeros(user_real_save[:, :, 0].shape, dtype=np.uint8) user_hair_mask_d_orig = np.zeros(user_real_save[:, :, 0].shape, dtype=np.uint8) cv2.warpAffine(test_res, inv_M_user, (orig_w, orig_h), dst=test_res_orig, borderMode=cv2.BORDER_TRANSPARENT) cv2.warpAffine(user_hair_mask_erode, inv_M_user, (orig_w, orig_h), dst=user_hair_mask_e_orig) cv2.warpAffine(user_hair_mask, inv_M_user, (orig_w, orig_h), dst=user_hair_mask_d_orig) # cv2.imshow("user_hair_mask", user_hair_mask) # cv2.imshow("user_hair_mask_d_orig", user_hair_mask_d_orig) # cv2.imshow("user_hair_mask_erode", user_hair_mask_erode) # cv2.imshow("user_hair_mask_e_orig", user_hair_mask_e_orig) # cv2.waitKey() user_bald_mask_b_orig = np.zeros(user_real_save[:, :, 0].shape, dtype=np.uint8) user_bald_mask_b = np.zeros(user_bald_mask[:, :, 0].shape, dtype=np.uint8) user_bald_mask_b[user_bald_mask[:, :, 0] > 0] = 1 user_bald_mask_b[user_bald_mask[:, :, 1] > 0] = 1 user_bald_mask_b[user_bald_mask[:, :, 2] > 0] = 1 if ratio != 0: user_bald_mask_b = cv2.dilate(user_bald_mask_b, np.ones((25, 25))) # default 25 cv2.warpAffine(user_bald_mask_b, inv_M_user, (orig_w, orig_h), dst=user_bald_mask_b_orig) user_hair_mask_d_orig_b = user_hair_mask_d_orig user_hair_mask_e_orig_b = cv2.blur(user_hair_mask_e_orig, (15, 15)) user_hair_mask_d_orig_b[~(user_bald_mask_b_orig.astype(bool))] = user_hair_mask_e_orig[ ~(user_bald_mask_b_orig.astype(bool))] user_hair_mask_d_orig_b = cv2.blur(user_hair_mask_d_orig_b, (15, 15)) else: user_bald_mask_b = cv2.dilate(user_bald_mask_b, np.ones((10, 10))) cv2.warpAffine(user_bald_mask_b, inv_M_user, (orig_w, orig_h), dst=user_bald_mask_b_orig) user_hair_mask_d_orig_b = user_hair_mask_d_orig user_hair_mask_d_orig_b[~(user_bald_mask_b_orig.astype(bool))] = user_hair_mask_e_orig[ ~(user_bald_mask_b_orig.astype(bool))] user_hair_mask_d_orig_b = cv2.blur(user_hair_mask_d_orig_b, (15, 15)) user_hair_mask_e_orig_b = cv2.blur(user_hair_mask_e_orig, (15, 15)) test_res_clear_orig = user_real_save * ( 1 - user_hair_mask_d_orig_b[:, :, np.newaxis] / 255) + test_res_orig * ( user_hair_mask_d_orig_b[:, :, np.newaxis] / 255) test_res_clear_orig = (np.clip(test_res_clear_orig, 0, 255)).astype(np.uint8) test_res_orig_lab = cv2.cvtColor(test_res_clear_orig, cv2.COLOR_BGR2LAB) user_real_save_lab = cv2.cvtColor(user_real_save, cv2.COLOR_BGR2LAB) test_res_orig_lab[:, :, 0] = user_real_save_lab[:, :, 0] fusion_res = cv2.cvtColor(test_res_orig_lab, cv2.COLOR_LAB2BGR) user_hair_mask_e_orig_b[~(user_bald_mask_b_orig.astype(bool))] = user_hair_mask_d_orig_b[~(user_bald_mask_b_orig.astype(bool))] fusion_res = test_res_clear_orig * ( 1 - user_hair_mask_e_orig_b[:, :, np.newaxis] / 255) + fusion_res * ( user_hair_mask_e_orig_b[:, :, np.newaxis] / 255) fusion_res = (np.clip(fusion_res, 0, 255)).astype(np.uint8) return fusion_res class Change_Hair_Color(object): def __init__(self, gpu, device_id): 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.model_dir = modelRoot self.dst_height = 768 self.dst_width = 768 generator_hair_path = os.path.join(self.model_dir, "zxm_v1_encode_refer_add_body_0107_use_gendata_from_hisd_0721.pt") #last: master_gen_hair_zxm_change_hair_color_v1_encode_refer_add_body self.net = torch.jit.load(generator_hair_path, torch.device('cpu')).to(self.device) print("hair color net: ", self.net) self.net.eval() def Change_Hair_inference(self, user_rgb_8uc3_change_color_768, user_matting_8uc3_change_color_768, ref_rgb_8uc3_change_color_768, ref_matting_8uc3_change_color_768): """ input： 图像尺寸基于： 人脸 768， 图像均为3通道 user_rgb_8uc3_change_color_768: 用户图， size 768， uint8 （0-255） user_matting_8uc3_change_color_768: 用户图 matting结果 size 768， uint8 （0-255） user_landmark_f1k2_change_color_768: 用户图 关键点 1k*2 size 768， float32 ref_rgb_8uc3_change_color_768: 参考图， size 768， uint8 （0-255） ref_matting_8uc3_change_color_768: 参考图， size 768， uint8 （0-255） ref_landmark_f1k2_change_color_768: 参考图， 关键点 1k*2 size 768， float32 output： hair_gene_color_8uc3_768: 头发换颜色生成图， size 768， uint8 （0-255） """ # cv2.imshow("user_rgb_8uc3_change_color_768", user_rgb_8uc3_change_color_768) # cv2.imshow("user_matting_8uc3_change_color_768", user_matting_8uc3_change_color_768) # cv2.imshow("ref_rgb_8uc3_change_color_768", ref_rgb_8uc3_change_color_768) # cv2.imshow("ref_matting_8uc3_change_color_768", ref_matting_8uc3_change_color_768) user_lab_user = cv2.cvtColor(user_rgb_8uc3_change_color_768, cv2.COLOR_BGR2LAB) user_lab_user[:, :, 1] = 0 user_lab_user[:, :, 2] = 0 img_paf = user_rgb_8uc3_change_color_768.copy() img_paf[user_matting_8uc3_change_color_768[:, :, 0].astype(bool)] = user_lab_user[user_matting_8uc3_change_color_768[:, :, 0].astype(bool)] img_paf = (np.clip(img_paf, 0, 255)).astype(np.uint8) tmp = ref_matting_8uc3_change_color_768[:, :, 0] tmp[tmp < 125] = 0 # tmp[tmp < 1] = 0 # mean_color = cv2.mean(refer_real, tmp) condit_hair = (ref_rgb_8uc3_change_color_768 * (tmp[:, :, np.newaxis] / 255)).astype(np.uint8) # cv2.imshow("condit_hair", condit_hair) # cv2.imshow("img_paf", img_paf) # cv2.waitKey() condition = condit_hair # np.concatenate((user_mask, user_hair_mask), axis=2) condition = (condition.astype(np.float32) / 255).transpose((2, 0, 1))[np.newaxis, :, :, :] input_paf = (img_paf.astype(np.float32) / 255).transpose((2, 0, 1))[np.newaxis, :, :, :] condition = torch.from_numpy(condition).to(self.device) input_paf = torch.from_numpy(input_paf).to(self.device) # ******************* forward ******************* with torch.no_grad(): test_fake = self.net(input_paf, condition) hair_gene_color_8uc3_768 = (test_fake.detach()[0].to('cpu').numpy().transpose((1, 2, 0)) * 255).astype(np.uint8).copy() # hair_colr_mid_show = np.concatenate((condit_hair, img_paf, hair_gene_color_8uc3_768), axis=1) # hair_colr_mid_show = cv2.resize(hair_colr_mid_show, (0, 0), fx=0.8, fy=0.8) # cv2.imshow("hair_colr_mid_show", hair_colr_mid_show) # cv2.waitKey(0) user_matting_mask_d = cv2.dilate(user_matting_8uc3_change_color_768, np.ones((30, 30), np.uint8)) user_matting_mask_d_b = cv2.blur(user_matting_mask_d, (30, 30)) # hair_gene_color_8uc3_768 = user_rgb_8uc3_change_color_768 * (1 - user_matting_8uc3_change_color_768 / 255) + test_res * (user_matting_8uc3_change_color_768 / 255) # hair_gene_color_8uc3_768 = (np.clip(hair_gene_color_8uc3_768, 0, 255)).astype(np.uint8) return hair_gene_color_8uc3_768, user_matting_mask_d_b class BodySeg(): def __init__(self, gpu_id=0): self.gpu_id = gpu_id self.device = torch.device(f"cuda:{gpu_id}") self.model = DeepLab() model_io.load_model_by_path("./weights/human_seg_deeplabv3_288_384_sigmoid_msc.pth", self.model, gpu_id=gpu_id) self.model.to(self.device) self.model.eval() self.output_img_size = [288, 384] self.resize_ratio = 2 def getM(self, center, angle, sx, sy): angle = math.radians(angle) alpha = math.cos(angle) beta = math.sin(angle) M = [[sx * alpha, sx * beta, (1 - sx * alpha) * center[0] - sx * beta * center[1]], [-sy * beta, sy * alpha, sy * beta * center[0] + (1 - sy * alpha) * center[1]]] return np.array(M) def forward(self, frame): # frame = cv2.imread(img_path).astype(np.float32) / 255 height, width = frame.shape[:2] x0, y0, x1, y1 = 0, 0, frame.shape[1], frame.shape[0] center_x, center_y = (x0 + x1) / 2, (y0 + y1) / 2 random_scalex = min(self.output_img_size[0] * self.resize_ratio * 1.0 / (x1 - x0), self.output_img_size[1] * self.resize_ratio * 1.0 / (y1 - y0)) random_scaley = random_scalex M = self.getM((center_x, center_y), 0, random_scalex, random_scaley) M[:, 2] += [self.output_img_size[0] * self.resize_ratio / 2 - center_x, self.output_img_size[1] * self.resize_ratio / 2 - center_y] crop_img = cv2.warpAffine(frame, M, (self.output_img_size[0] * self.resize_ratio, self.output_img_size[1] * self.resize_ratio)) crop_img_resize = cv2.resize(crop_img, (0, 0), fx=1 / self.resize_ratio, fy=1 / self.resize_ratio) input_tensor = torch.from_numpy(crop_img_resize.transpose([2, 0, 1])[np.newaxis]).to(self.device) input_tensor = input_tensor.float() output_tensor, _ = self.model(input_tensor) output_tensor = F.interpolate(output_tensor, size=[self.output_img_size[1] * self.resize_ratio, self.output_img_size[0] * self.resize_ratio], mode='bilinear', align_corners=True) # out_mask = torch.max(output_tensor[:1], 1)[1].detach().cpu().numpy().squeeze().astype(np.float32) output_tensor = torch.sigmoid(output_tensor) out_mask = output_tensor[0, 0].detach().cpu().numpy().squeeze().astype(np.float32) input_output_show = np.concatenate((crop_img, np.repeat(out_mask[:, :, np.newaxis], 3, axis=2)), axis=1) # cv2.imshow("input_output_show", input_output_show) M_inv = cv2.invertAffineTransform(M) # mask_rawsize = cv2.warpAffine(out_mask, M_inv, (frame.shape[1], frame.shape[0])) mask_rawsize = cv2.warpAffine(out_mask, M_inv, (frame.shape[1], frame.shape[0]), flags=cv2.INTER_NEAREST) mask_rawsize = np.repeat(mask_rawsize[:, :, np.newaxis], 3, axis=2) return mask_rawsize def forward_kpts(self, frame, landmarks_1k): M = landmark_processor.get_transform_mat_bodyseg(landmarks_1k, output_size=self.output_img_size[0]*self.resize_ratio, ratio=0.7, offset=(0, 45)) crop_img = cv2.warpAffine(frame, M, (self.output_img_size[0] * self.resize_ratio, self.output_img_size[1] * self.resize_ratio)).astype(np.float32) / 255 crop_img_resize = cv2.resize(crop_img, (0, 0), fx=1 / self.resize_ratio, fy=1 / self.resize_ratio) input_tensor = torch.from_numpy(crop_img_resize.transpose([2, 0, 1])[np.newaxis]).to(self.device) input_tensor = input_tensor.float() output_tensor, _ = self.model(input_tensor) output_tensor = F.interpolate(output_tensor, size=[self.output_img_size[1] * self.resize_ratio, self.output_img_size[0] * self.resize_ratio], mode='bilinear', align_corners=True) # out_mask = torch.max(output_tensor[:1], 1)[1].detach().cpu().numpy().squeeze().astype(np.float32) output_tensor = torch.sigmoid(output_tensor) out_mask = output_tensor[0, 0].detach().cpu().numpy().squeeze().astype(np.float32) # input_output_show = np.concatenate((crop_img, np.repeat(out_mask[:, :, np.newaxis], 3, axis=2)), axis=1) # cv2.imshow("input_output_show", input_output_show) # cv2.waitKey() M_inv = cv2.invertAffineTransform(M) # mask_rawsize = cv2.warpAffine(out_mask, M_inv, (frame.shape[1], frame.shape[0])) mask_rawsize = cv2.warpAffine(out_mask, M_inv, (frame.shape[1], frame.shape[0]), flags=cv2.INTER_NEAREST) mask_rawsize = np.repeat(mask_rawsize[:, :, np.newaxis], 3, axis=2) return mask_rawsize