|
|
|
from numpy.core.numeric import require |
|
from numpy.lib.function_base import quantile |
|
import torch |
|
import numpy as np |
|
from facemodel import Face_3DMM |
|
from data_loader import load_dir |
|
from util import * |
|
import os |
|
import sys |
|
import cv2 |
|
import imageio |
|
import argparse |
|
|
|
dir_path = os.path.dirname(os.path.realpath(__file__)) |
|
|
|
|
|
def set_requires_grad(tensor_list): |
|
for tensor in tensor_list: |
|
tensor.requires_grad = True |
|
|
|
|
|
parser = argparse.ArgumentParser() |
|
parser.add_argument( |
|
"--path", type=str, default="obama/ori_imgs", help="idname of target person") |
|
parser.add_argument('--img_h', type=int, default=512, help='image height') |
|
parser.add_argument('--img_w', type=int, default=512, help='image width') |
|
parser.add_argument('--frame_num', type=int, |
|
default=11000, help='image number') |
|
args = parser.parse_args() |
|
start_id = 0 |
|
end_id = args.frame_num |
|
|
|
lms = load_dir(args.path, start_id, end_id) |
|
num_frames = lms.shape[0] |
|
h, w = args.img_h, args.img_w |
|
cxy = torch.tensor((w/2.0, h/2.0), dtype=torch.float).cuda() |
|
id_dim, exp_dim, tex_dim, point_num = 100, 79, 100, 34650 |
|
model_3dmm = Face_3DMM(os.path.join(dir_path, '3DMM'), |
|
id_dim, exp_dim, tex_dim, point_num) |
|
lands_info = np.loadtxt(os.path.join( |
|
dir_path, '3DMM', 'lands_info.txt'), dtype=np.int32) |
|
lands_info = torch.as_tensor(lands_info).cuda() |
|
|
|
focal = 1150 |
|
|
|
id_para = lms.new_zeros((1, id_dim), requires_grad=True) |
|
exp_para = lms.new_zeros((num_frames, exp_dim), requires_grad=True) |
|
tex_para = lms.new_zeros((1, tex_dim), requires_grad=True) |
|
euler_angle = lms.new_zeros((num_frames, 3), requires_grad=True) |
|
trans = lms.new_zeros((num_frames, 3), requires_grad=True) |
|
light_para = lms.new_zeros((num_frames, 27), requires_grad=True) |
|
trans.data[:, 2] -= 600 |
|
focal_length = lms.new_zeros(1, requires_grad=True) |
|
focal_length.data += focal |
|
|
|
set_requires_grad([id_para, exp_para, tex_para, |
|
euler_angle, trans, light_para]) |
|
|
|
sel_ids = np.arange(0, num_frames, 10) |
|
sel_num = sel_ids.shape[0] |
|
arg_focal = 0.0 |
|
arg_landis = 1e5 |
|
for focal in range(500, 1500, 50): |
|
id_para = lms.new_zeros((1, id_dim), requires_grad=True) |
|
exp_para = lms.new_zeros((sel_num, exp_dim), requires_grad=True) |
|
euler_angle = lms.new_zeros((sel_num, 3), requires_grad=True) |
|
trans = lms.new_zeros((sel_num, 3), requires_grad=True) |
|
trans.data[:, 2] -= 600 |
|
focal_length = lms.new_zeros(1, requires_grad=False) |
|
focal_length.data += focal |
|
set_requires_grad([id_para, exp_para, euler_angle, trans]) |
|
|
|
optimizer_id = torch.optim.Adam([id_para], lr=.3) |
|
optimizer_exp = torch.optim.Adam([exp_para], lr=.3) |
|
optimizer_frame = torch.optim.Adam( |
|
[euler_angle, trans], lr=.3) |
|
iter_num = 2000 |
|
|
|
for iter in range(iter_num): |
|
id_para_batch = id_para.expand(sel_num, -1) |
|
geometry = model_3dmm.forward_geo_sub( |
|
id_para_batch, exp_para, lands_info[-51:].long()) |
|
proj_geo = forward_transform( |
|
geometry, euler_angle, trans, focal_length, cxy) |
|
loss_lan = cal_lan_loss( |
|
proj_geo[:, :, :2], lms[sel_ids, -51:, :].detach()) |
|
loss_regid = torch.mean(id_para*id_para)*8 |
|
loss_regexp = torch.mean(exp_para*exp_para)*0.5 |
|
loss = loss_lan + loss_regid + loss_regexp |
|
optimizer_id.zero_grad() |
|
optimizer_exp.zero_grad() |
|
optimizer_frame.zero_grad() |
|
loss.backward() |
|
if iter > 1000: |
|
optimizer_id.step() |
|
optimizer_exp.step() |
|
optimizer_frame.step() |
|
print(focal, loss_lan.item(), torch.mean(trans[:, 2]).item()) |
|
if loss_lan.item() < arg_landis: |
|
arg_landis = loss_lan.item() |
|
arg_focal = focal |
|
|
|
sel_ids = np.arange(0, num_frames) |
|
sel_num = sel_ids.shape[0] |
|
id_para = lms.new_zeros((1, id_dim), requires_grad=True) |
|
exp_para = lms.new_zeros((sel_num, exp_dim), requires_grad=True) |
|
euler_angle = lms.new_zeros((sel_num, 3), requires_grad=True) |
|
trans = lms.new_zeros((sel_num, 3), requires_grad=True) |
|
trans.data[:, 2] -= 600 |
|
focal_length = lms.new_zeros(1, requires_grad=False) |
|
focal_length.data += arg_focal |
|
set_requires_grad([id_para, exp_para, euler_angle, trans]) |
|
|
|
optimizer_id = torch.optim.Adam([id_para], lr=.3) |
|
optimizer_exp = torch.optim.Adam([exp_para], lr=.3) |
|
optimizer_frame = torch.optim.Adam( |
|
[euler_angle, trans], lr=.3) |
|
iter_num = 2000 |
|
|
|
for iter in range(iter_num): |
|
id_para_batch = id_para.expand(sel_num, -1) |
|
geometry = model_3dmm.forward_geo_sub( |
|
id_para_batch, exp_para, lands_info[-51:].long()) |
|
proj_geo = forward_transform( |
|
geometry, euler_angle, trans, focal_length, cxy) |
|
loss_lan = cal_lan_loss( |
|
proj_geo[:, :, :2], lms[sel_ids, -51:, :].detach()) |
|
loss_regid = torch.mean(id_para*id_para)*8 |
|
loss_regexp = torch.mean(exp_para*exp_para)*0.5 |
|
loss = loss_lan + loss_regid + loss_regexp |
|
optimizer_id.zero_grad() |
|
optimizer_exp.zero_grad() |
|
optimizer_frame.zero_grad() |
|
loss.backward() |
|
if iter > 1000: |
|
optimizer_id.step() |
|
optimizer_exp.step() |
|
optimizer_frame.step() |
|
print(arg_focal, loss_lan.item(), torch.mean(trans[:, 2]).item()) |
|
|
|
|
|
torch.save({'id': id_para.detach().cpu(), 'exp': exp_para.detach().cpu(), |
|
'euler': euler_angle.detach().cpu(), 'trans': trans.detach().cpu(), |
|
'focal': focal_length.detach().cpu()}, os.path.join(os.path.dirname(args.path), 'track_params.pt')) |
|
print('face tracking params saved') |
|
|
