153 lines
6.1 KiB
Python
153 lines
6.1 KiB
Python
import os
|
|
import torch
|
|
import cv2
|
|
import os.path as osp
|
|
import numpy as np
|
|
from PIL import Image
|
|
from .CSD_MT.options import Options
|
|
from .CSD_MT.model import CSD_MT
|
|
from .faceutils.face_parsing.model import BiSeNet
|
|
import torchvision.transforms as transforms
|
|
from . import faceutils as futils
|
|
|
|
|
|
# load face_parsing model
|
|
# n_classes = 19
|
|
# face_paseing_model = BiSeNet(n_classes=n_classes)
|
|
# save_pth = osp.join('faceutils/face_parsing/res/cp', '79999_iter.pth')
|
|
# face_paseing_model.load_state_dict(torch.load(save_pth,map_location='cpu'))
|
|
# face_paseing_model.eval()
|
|
|
|
# # load makeup transfer model
|
|
# parser = Options()
|
|
# opts = parser.parse()
|
|
# makeup_model = CSD_MT(opts)
|
|
# ep0, total_it = makeup_model.resume('CSD_MT/weights/CSD_MT.pth')
|
|
# makeup_model.eval()
|
|
|
|
def crop_image(image):
|
|
up_ratio = 0.2 / 0.85 # delta_size / face_size
|
|
down_ratio = 0.15 / 0.85 # delta_size / face_size
|
|
width_ratio = 0.2 / 0.85 # delta_size / face_size
|
|
|
|
#image = Image.fromarray(image)
|
|
face = futils.dlib.detect(image)
|
|
|
|
if not face:
|
|
raise ValueError("No face !")
|
|
|
|
face_on_image = face[0]
|
|
|
|
image, face, crop_face = futils.dlib.crop(image, face_on_image, up_ratio, down_ratio, width_ratio)
|
|
np_image = np.array(image)
|
|
return np_image
|
|
|
|
def get_face_parsing(x,face_paseing_model):
|
|
to_tensor = transforms.Compose([
|
|
transforms.ToTensor(),
|
|
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
|
|
])
|
|
with torch.no_grad():
|
|
img = Image.fromarray(x)
|
|
image = img.resize((512, 512), Image.BILINEAR)
|
|
img = to_tensor(image)
|
|
img = torch.unsqueeze(img, 0)
|
|
out = face_paseing_model(img)[0]
|
|
parsing = out.squeeze(0).cpu().numpy().argmax(0)
|
|
return parsing
|
|
|
|
|
|
def split_parse(opts,parse):
|
|
h, w = parse.shape
|
|
result = np.zeros([h, w, opts.semantic_dim])
|
|
result[:, :, 0][np.where(parse == 0)] = 1
|
|
result[:, :, 0][np.where(parse == 16)] = 1
|
|
result[:, :, 0][np.where(parse == 17)] = 1
|
|
result[:, :, 0][np.where(parse == 18)] = 1
|
|
result[:, :, 0][np.where(parse == 9)] = 1
|
|
result[:, :, 1][np.where(parse == 1)] = 1
|
|
result[:, :, 2][np.where(parse == 2)] = 1
|
|
result[:, :, 2][np.where(parse == 3)] = 1
|
|
result[:, :, 3][np.where(parse == 4)] = 1
|
|
result[:, :, 3][np.where(parse == 5)] = 1
|
|
result[:, :, 1][np.where(parse == 6)] = 1
|
|
result[:, :, 4][np.where(parse == 7)] = 1
|
|
result[:, :, 4][np.where(parse == 8)] = 1
|
|
result[:, :, 5][np.where(parse == 10)] = 1
|
|
result[:, :, 6][np.where(parse == 11)] = 1
|
|
result[:, :, 7][np.where(parse == 12)] = 1
|
|
result[:, :, 8][np.where(parse == 13)] = 1
|
|
result[:, :, 9][np.where(parse == 14)] = 1
|
|
result[:, :, 9][np.where(parse == 15)] = 1
|
|
result = np.array(result)
|
|
return result
|
|
|
|
|
|
def local_masks(opts,split_parse):
|
|
h, w, c = split_parse.shape
|
|
all_mask = np.zeros([h, w])
|
|
all_mask[np.where(split_parse[:, :, 0] == 0)] = 1
|
|
all_mask[np.where(split_parse[:, :, 3] == 1)] = 0
|
|
all_mask[np.where(split_parse[:, :, 6] == 1)] = 0
|
|
all_mask = np.expand_dims(all_mask, axis=2) # Expansion of the dimension
|
|
all_mask = np.concatenate((all_mask, all_mask, all_mask), axis=2)
|
|
return all_mask
|
|
|
|
|
|
|
|
def load_data_from_image(non_makeup_img, makeup_img,opts,face_paseing_model):
|
|
non_makeup_img=crop_image(non_makeup_img)
|
|
makeup_img = crop_image(makeup_img)
|
|
non_makeup_img=cv2.resize(non_makeup_img,(opts.resize_size,opts.resize_size))
|
|
makeup_img = cv2.resize(makeup_img, (opts.resize_size, opts.resize_size))
|
|
non_makeup_parse = get_face_parsing(non_makeup_img,face_paseing_model)
|
|
non_makeup_parse = cv2.resize(non_makeup_parse, (opts.resize_size, opts.resize_size),interpolation=cv2.INTER_NEAREST)
|
|
makeup_parse = get_face_parsing(makeup_img,face_paseing_model)
|
|
makeup_parse = cv2.resize(makeup_parse, (opts.resize_size, opts.resize_size),interpolation=cv2.INTER_NEAREST)
|
|
|
|
non_makeup_split_parse = split_parse(opts,non_makeup_parse)
|
|
makeup_split_parse = split_parse(opts,makeup_parse)
|
|
|
|
non_makeup_all_mask = local_masks(opts,
|
|
non_makeup_split_parse)
|
|
makeup_all_mask = local_masks(opts,
|
|
makeup_split_parse)
|
|
|
|
non_makeup_img = non_makeup_img / 127.5 - 1
|
|
non_makeup_img = np.transpose(non_makeup_img, (2, 0, 1))
|
|
non_makeup_split_parse = np.transpose(non_makeup_split_parse, (2, 0, 1))
|
|
|
|
makeup_img = makeup_img / 127.5 - 1
|
|
makeup_img = np.transpose(makeup_img, (2, 0, 1))
|
|
makeup_split_parse = np.transpose(makeup_split_parse, (2, 0, 1))
|
|
|
|
non_makeup_img=torch.from_numpy(non_makeup_img).type(torch.FloatTensor)
|
|
non_makeup_img = torch.unsqueeze(non_makeup_img, 0)
|
|
non_makeup_split_parse = torch.from_numpy(non_makeup_split_parse).type(torch.FloatTensor)
|
|
non_makeup_split_parse = torch.unsqueeze(non_makeup_split_parse, 0)
|
|
non_makeup_all_mask = np.transpose(non_makeup_all_mask, (2, 0, 1))
|
|
|
|
makeup_img = torch.from_numpy(makeup_img).type(torch.FloatTensor)
|
|
makeup_img = torch.unsqueeze(makeup_img, 0)
|
|
makeup_split_parse = torch.from_numpy(makeup_split_parse).type(torch.FloatTensor)
|
|
makeup_split_parse = torch.unsqueeze(makeup_split_parse, 0)
|
|
makeup_all_mask = np.transpose(makeup_all_mask, (2, 0, 1))
|
|
|
|
data = {'non_makeup_color_img': non_makeup_img,
|
|
'non_makeup_split_parse':non_makeup_split_parse,
|
|
'non_makeup_all_mask': torch.unsqueeze(torch.from_numpy(non_makeup_all_mask).type(torch.FloatTensor), 0),
|
|
|
|
'makeup_color_img': makeup_img,
|
|
'makeup_split_parse': makeup_split_parse,
|
|
'makeup_all_mask': torch.unsqueeze(torch.from_numpy(makeup_all_mask).type(torch.FloatTensor), 0)
|
|
}
|
|
return data
|
|
|
|
def makeup_transfer256(non_makeup_image, makeup_image,makeup_model,opts,face_paseing_model):
|
|
data=load_data_from_image(non_makeup_image, makeup_image, opts=opts,face_paseing_model=face_paseing_model)
|
|
with torch.no_grad():
|
|
transfer_tensor=makeup_model.test_pair(data)
|
|
transfer_img=transfer_tensor[0].cpu().float().numpy()
|
|
transfer_img = np.transpose((transfer_img/ 2 + 0.5)*255., (1, 2, 0))
|
|
transfer_img=np.clip(transfer_img+0.5,0,255).astype(np.uint8)
|
|
return transfer_img |