class Evaluate:
def __init__(self, arguments, cuda=False):
"""
Evaluate
:param arguments: argparse options
:param cuda: gpu speed up
"""
self.args = arguments
location = self.args.lightcnn
self.lightcnn_inst = utils.load_lightcnn(location, cuda)
self.cuda = cuda
self.parsing = self.args.parsing_checkpoint
self.max_itr = arguments.total_eval_steps
self.learning_rate = arguments.eval_learning_rate
self.losses = []
self.prev_path = "./output/eval"
self.model_path = "../unity/models"
self.clean()
self.imitator = Imitator("neural imitator", arguments, clean=False)
self.l2_c = (torch.ones((512, 512)), torch.ones((512, 512)))
if cuda:
self.imitator.cuda()
self.imitator.eval()
self.imitator.load_checkpoint(args.imitator_model, False, cuda=cuda)
def _init_l1_l2(self, y):
"""
init reference photo l1 & l2
:param y: input photo, numpy array [H, W, C]
"""
y_ = cv2.resize(y, dsize=(128, 128), interpolation=cv2.INTER_LINEAR)
y_ = np.swapaxes(y_, 0, 2).astype(np.float32)
y_ = np.mean(y_, axis=0)[np.newaxis, np.newaxis, :, :]
y_ = torch.from_numpy(y_)
if self.cuda:
y_ = y_.cuda()
self.l1_y = y_
y = y[np.newaxis, :, :, ]
y = np.swapaxes(y, 1, 2)
y = np.swapaxes(y, 1, 3)
y = torch.from_numpy(y)
if self.cuda:
y = y.cuda()
self.l2_y = y / 255.
def discrim_l1(self, y_):
"""
content loss evaluated by lightcnn
:param y_: generated image, torch tensor [B, C, W, H]
:return: l1 loss
"""
y_ = F.max_pool2d(y_, kernel_size=(4, 4), stride=4) # 512->128
y_ = torch.mean(y_, dim=1).view(1, 1, 128, 128) # gray
return utils.discriminative_loss(self.l1_y, y_, self.lightcnn_inst)
def discrim_l2(self, y_):
"""
facial semantic feature loss
evaluate loss use l1 at pixel space
:param y_: generated image, tensor [B, C, W, H]
:return: l1 loss in pixel space
"""
# [eyebrow，eye，nose，teeth，up lip，lower lip]
w_r = [1.1, 1.1, 1., 0.7, 1., 1.]
w_g = [1.1, 1.1, 1., 0.7, 1., 1.]
part1, _ = faceparsing_tensor(self.l2_y,
self.parsing,
w_r,
cuda=self.cuda)
y_ = y_.transpose(2, 3)
part2, _ = faceparsing_tensor(y_, self.parsing, w_g, cuda=self.cuda)
self.l2_c = (part1 * 10, part2 * 10)
return F.l1_loss(part1, part2)
def evaluate_ls(self, y_):
"""
评估损失Ls
由于l1表示的是余弦距离的损失， 其范围又在0-1之间 所以这里使用1-l1
(余弦距离越大 表示越接近)
:param y_: generated image, tensor [b,c,w,h]
:return: ls, description
"""
l1 = self.discrim_l1(y_)
l2 = self.discrim_l2(y_)
alpha = self.args.eval_alpha
ls = alpha * (1 - l1) + l2
info = "l1:{0:.3f} l2:{1:.3f} ls:{2:.3f}".format(l1, l2, ls)
self.losses.append((l1.item(), l2.item() / 3, ls.item()))
return ls, info
def itr_train(self, y):
"""
iterator train
:param y: numpy array, image [H, W, C]
"""
param_cnt = self.args.params_cnt
t_params = 0.5 * torch.ones((1, param_cnt), dtype=torch.float32)
if self.cuda:
t_params = t_params.cuda()
t_params.requires_grad = True
self.losses.clear()
lr = self.learning_rate
self._init_l1_l2(y)
m_progress = tqdm(range(1, self.max_itr + 1))
for i in m_progress:
y_ = self.imitator(t_params)
loss, info = self.evaluate_ls(y_)
loss.backward()
if i == 1:
self.output(t_params, y, 0)
t_params.data = t_params.data - lr * t_params.grad.data
t_params.data = t_params.data.clamp(0., 1.)
t_params.grad.zero_()
m_progress.set_description(info)
if i % self.args.eval_prev_freq == 0:
x = i / float(self.max_itr)
lr = self.learning_rate * (1 - x) + 1e-2
self.output(t_params, y, i)
self.plot()
self.plot()
log.info("steps:{0} params:{1}".format(self.max_itr, t_params.data))
return t_params
def output(self, x, refer, step):
"""
capture for result
:param x: generated image with grad, torch tensor [b,params]
:param refer: reference picture
:param step: train step
"""
self.write(x)
y_ = self.imitator(x)
y_ = y_.cpu().detach().numpy()
y_ = np.squeeze(y_, axis=0)
y_ = np.swapaxes(y_, 0, 2) * 255
y_ = y_.astype(np.uint8)
im1 = self.l2_c[0]
im2 = self.l2_c[1]
np_im1 = im1.cpu().detach().numpy()
np_im2 = im2.cpu().detach().numpy()
f_im1 = ops.fill_gray(np_im1)
f_im2 = ops.fill_gray(np_im2)
image_ = ops.merge_4image(refer, y_, f_im1, f_im2, transpose=False)
path = os.path.join(self.prev_path, "eval_{0}.jpg".format(step))
cv2.imwrite(path, image_)
def write(self, params):
"""
生成二进制文件 能够在unity里还原出来
:param params: 捏脸参数 tensor [batch, params_cnt]
"""
np_param = params.cpu().detach().numpy()
np_param = np_param[0]
list_param = np_param.tolist()
dataset = self.args.path_to_dataset
shape = utils.curr_roleshape(dataset)
path = os.path.join(self.model_path, "eval.bytes")
f = open(path, 'wb')
write_layer(f, shape, list_param)
f.close()
def clean(self):
"""
clean for new iter
"""
ops.clear_files(self.prev_path)
ops.clear_files(self.model_path)
def plot(self):
"""
plot loss
"""
count = len(self.losses)
if count > 0:
plt.style.use('seaborn-whitegrid')
x = range(count)
y1 = []
y2 = []
for it in self.losses:
y1.append(it[0])
y2.append(it[1])
plt.plot(x, y1, color='r', label='l1')
plt.plot(x, y2, color='g', label='l2')
plt.ylabel("loss")
plt.xlabel('step')
plt.legend()
path = os.path.join(self.prev_path, "loss.png")
plt.savefig(path)
plt.close('all')
imitator = Imitator("neural imitator", args)
if cuda:
imitator.cuda()
imitator.batch_train(cuda)
elif args.phase == "train_extractor":
log.info('feature extractor train mode')
extractor = Extractor("neural extractor", args)
if cuda:
extractor.cuda()
extractor.batch_train(cuda)
elif args.phase == "inference_imitator":
log.info("inference imitator")
imitator = Imitator("neural imitator", args, clean=False)
if cuda:
imitator.cuda()
imitator.load_checkpoint(args.imitator_model, True, cuda=cuda)
elif args.phase == "prev_imitator":
log.info("preview imitator")
imitator = Imitator("neural imitator", args, clean=False)
imitator.load_checkpoint(args.imitator_model, False, cuda=False)
dataset = FaceDataset(args)
name, param, img = dataset.get_picture()
param = np.array(param, dtype=np.float32)
b_param = param[np.newaxis, :]
log.info(b_param.shape)
t_param = torch.from_numpy(b_param)
output = imitator(t_param)
output = output.cpu().detach().numpy()
output = np.squeeze(output, axis=0)
output = output.swapaxes(0, 2) * 255
cv2.imwrite('./output/{0}.jpg'.format(name), output)
class Evaluate:
def __init__(self, arguments, cuda=False):
"""
Evaluate
:param arguments: argparse options
:param cuda: gpu speed up
"""
self.args = arguments
location = self.args.lightcnn
self.lightcnn_inst = utils.load_lightcnn(location, cuda)
self.cuda = cuda
self.parsing = self.args.parsing_checkpoint
self.max_itr = arguments.total_eval_steps
self.learning_rate = arguments.eval_learning_rate
self.losses = []
self.prev_path = "./output/eval"
self.clean()
self.imitator = Imitator("neural imitator", arguments, clean=False)
if cuda:
self.imitator.cuda()
self.imitator.eval()
self.imitator.load_checkpoint(args.imitator_model, False, cuda=cuda)
def discrim_l1(self, y, y_):
"""
content loss evaluated by lightcnn
:param y: input photo, numpy array [H, W, C]
:param y_: generated image, torch tensor [B, C, W, H]
:return: l1 loss
"""
y = cv2.resize(y, dsize=(128, 128), interpolation=cv2.INTER_LINEAR)
y = np.swapaxes(y, 0, 2).astype(np.float32)
y = np.mean(y, axis=0)[np.newaxis, np.newaxis, :, :]
y = torch.from_numpy(y)
if self.cuda:
y = y.cuda()
y_ = F.max_pool2d(y_, kernel_size=(4, 4), stride=4) # 512->128
y_ = torch.mean(y_, dim=1).view(1, 1, 128, 128) # gray
return utils.discriminative_loss(y, y_, self.lightcnn_inst)
def discrim_l2(self, y, y_, step):
"""
facial semantic feature loss
evaluate loss use l1 at pixel space
:param y: input photo, numpy array [H, W, C]
:param y_: generated image, tensor [B, C, W, H]
:param step: train step
:return: l1 loss in pixel space
"""
img1 = parse_evaluate(y.astype(np.uint8),
cp=self.parsing,
cuda=self.cuda)
y_ = y_.cpu().detach().numpy()
y_ = np.squeeze(y_, axis=0)
y_ = np.swapaxes(y_, 0, 2) * 255
img2 = parse_evaluate(y_.astype(np.uint8),
cp=self.parsing,
cuda=self.cuda)
edge_img1 = utils.img_edge(img1).astype(np.float32)
edge_img2 = utils.img_edge(img2).astype(np.float32)
w_g = 1.0
w_r = 1.0
if step % self.args.eval_prev_freq == 0:
path = os.path.join(self.prev_path, "l2_{0}.jpg".format(step))
edge1_v3 = 255. - ops.fill_grey(edge_img1)
edge2_v3 = 255. - ops.fill_grey(edge_img2)
merge = ops.merge_4image(y,
y_,
edge1_v3,
edge2_v3,
transpose=False)
cv2.imwrite(path, merge)
return np.mean(np.abs(w_r * edge_img1 - w_g * edge_img2))
def evaluate_ls(self, y, y_, step):
"""
评估损失Ls
:param y: input photo, numpy array
:param y_: generated image, tensor [b,c,w,h]
:param step: train step
:return: ls, description
"""
l1 = self.discrim_l1(y, y_)
l2 = self.discrim_l2(y, y_, step)
alpha = self.args.eval_alpha
ls = alpha * l1 + l2
info = "l1:{0:.3f} l2:{1:.3f} ls:{2:.3f}".format(alpha * l1, l2, ls)
self.losses.append((l1.item() * alpha, l2.item(), ls.item()))
return ls, info
def itr_train(self, y):
"""
iterator train
:param y: numpy array, image [H, W, C]
"""
param_cnt = self.args.params_cnt
t_params = 0.5 * torch.ones((1, param_cnt), dtype=torch.float32)
if self.cuda:
t_params = t_params.cuda()
t_params.requires_grad = True
self.losses.clear()
lr = self.learning_rate
progress = tqdm(range(self.max_itr), initial=0, total=self.max_itr)
for i in progress:
y_ = self.imitator.forward(t_params)
loss, info = self.evaluate_ls(y, y_, i)
loss.backward()
t_params.data = t_params.data - lr * t_params.grad.data
t_params.data = t_params.data.clamp(0., 1.)
t_params.grad.zero_()
progress.set_description(info)
if self.max_itr % 100 == 0:
x = i / float(self.max_itr)
lr = self.learning_rate * (x**2 - 2 * x + 1) + 1e-4
self.plot()
log.info("steps:{0} params:{1}".format(self.max_itr, t_params.data))
return t_params
def output(self, x, refer):
"""
capture for result
:param x: generated image with grad, torch tensor [b,params]
:param refer: reference picture
"""
self.write(x)
y_ = self.imitator.forward(x)
y_ = y_.cpu().detach().numpy()
y_ = np.squeeze(y_, axis=0)
y_ = np.swapaxes(y_, 0, 2) * 255.
y_ = y_.astype(np.uint8)
image_ = ops.merge_image(refer, y_, transpose=False)
path = os.path.join(self.prev_path, "eval.jpg")
cv2.imwrite(path, image_)
def write(self, params):
"""
生成二进制文件 能够在unity里还原出来
:param params: 捏脸参数 tensor [batch, 95]
"""
np_param = params.cpu().detach().numpy()
np_param = np_param[0]
list_param = np_param.tolist()
dataset = self.args.path_to_dataset
shape = utils.curr_roleshape(dataset)
f = open("../unity/models/eval.bytes", 'wb')
write_layer(f, shape, list_param)
f.close()
def clean(self):
"""
clean for new iter
"""
ops.clear_files(self.prev_path)
def plot(self):
plt.style.use('seaborn-whitegrid')
x = range(self.max_itr)
y1 = []
y2 = []
for it in self.losses:
y1.append(it[0])
y2.append(it[1])
plt.plot(x, y1, color='r', label='l1')
plt.plot(x, y2, color='g', label='l2')
plt.ylabel("loss")
plt.xlabel('step')
plt.legend()
path = os.path.join(self.prev_path, "eval.png")
plt.savefig(path)