xsl
443cfa298f
包含： - hair_service_sd: 主服务（换发型/换发色/生发，端口8801） - photo_service: LoRA调度+训练（端口32678） - hair_grow_service: 调试测试页（端口8888，含4个测试页） - 批量训练脚本（batch_train_hairstyles.py） - 发际线mask自动识别（hairline_mask.py，4种方案） - 手绘mask换发型（hair_swap_manual.py） - 文档：README.md + LARGE_FILES.md + docs/ 大文件（模型权重200G、训练数据123G）已排除，见 LARGE_FILES.md OSS/COS密钥已脱敏为环境变量，原文件备份在本地
3781 lines
183 KiB
Python
3781 lines
183 KiB
Python
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
|
||
|