Spaces:
Sleeping
Sleeping
| import os | |
| import sys | |
| import cv2 | |
| import argparse | |
| from pathlib import Path | |
| import torch | |
| import numpy as np | |
| from data_loader import load_dir | |
| from facemodel import Face_3DMM | |
| from util import * | |
| from render_3dmm import Render_3DMM | |
| # torch.autograd.set_detect_anomaly(True) | |
| 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, img_paths = 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 | |
| ) | |
| # only use one image per 40 to do fit the focal length | |
| sel_ids = np.arange(0, num_frames, 40) | |
| sel_num = sel_ids.shape[0] | |
| arg_focal = 1600 | |
| arg_landis = 1e5 | |
| print(f'[INFO] fitting focal length...') | |
| # fit the focal length | |
| for focal in range(600, 1500, 100): | |
| 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] -= 7 | |
| focal_length = lms.new_zeros(1, requires_grad=False) | |
| focal_length.data += focal | |
| set_requires_grad([id_para, exp_para, euler_angle, trans]) | |
| optimizer_idexp = torch.optim.Adam([id_para, exp_para], lr=0.1) | |
| optimizer_frame = torch.optim.Adam([euler_angle, trans], lr=0.1) | |
| for iter in range(2000): | |
| id_para_batch = id_para.expand(sel_num, -1) | |
| geometry = model_3dmm.get_3dlandmarks( | |
| id_para_batch, exp_para, euler_angle, trans, focal_length, cxy | |
| ) | |
| proj_geo = forward_transform(geometry, euler_angle, trans, focal_length, cxy) | |
| loss_lan = cal_lan_loss(proj_geo[:, :, :2], lms[sel_ids].detach()) | |
| loss = loss_lan | |
| optimizer_frame.zero_grad() | |
| loss.backward() | |
| optimizer_frame.step() | |
| # if iter % 100 == 0: | |
| # print(focal, 'pose', iter, loss.item()) | |
| for iter in range(2500): | |
| id_para_batch = id_para.expand(sel_num, -1) | |
| geometry = model_3dmm.get_3dlandmarks( | |
| id_para_batch, exp_para, euler_angle, trans, focal_length, cxy | |
| ) | |
| proj_geo = forward_transform(geometry, euler_angle, trans, focal_length, cxy) | |
| loss_lan = cal_lan_loss(proj_geo[:, :, :2], lms[sel_ids].detach()) | |
| loss_regid = torch.mean(id_para * id_para) | |
| loss_regexp = torch.mean(exp_para * exp_para) | |
| loss = loss_lan + loss_regid * 0.5 + loss_regexp * 0.4 | |
| optimizer_idexp.zero_grad() | |
| optimizer_frame.zero_grad() | |
| loss.backward() | |
| optimizer_idexp.step() | |
| optimizer_frame.step() | |
| # if iter % 100 == 0: | |
| # print(focal, 'poseidexp', iter, loss_lan.item(), loss_regid.item(), loss_regexp.item()) | |
| if iter % 1500 == 0 and iter >= 1500: | |
| for param_group in optimizer_idexp.param_groups: | |
| param_group["lr"] *= 0.2 | |
| for param_group in optimizer_frame.param_groups: | |
| param_group["lr"] *= 0.2 | |
| print(focal, loss_lan.item(), torch.mean(trans[:, 2]).item()) | |
| if loss_lan.item() < arg_landis: | |
| arg_landis = loss_lan.item() | |
| arg_focal = focal | |
| print("[INFO] find best focal:", arg_focal) | |
| print(f'[INFO] coarse fitting...') | |
| # for all frames, do a coarse fitting ??? | |
| 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 | |
| ) # not optimized in this block ??? | |
| 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] -= 7 # ??? | |
| focal_length = lms.new_zeros(1, requires_grad=True) | |
| focal_length.data += arg_focal | |
| set_requires_grad([id_para, exp_para, tex_para, euler_angle, trans, light_para]) | |
| optimizer_idexp = torch.optim.Adam([id_para, exp_para], lr=0.1) | |
| optimizer_frame = torch.optim.Adam([euler_angle, trans], lr=1) | |
| for iter in range(1500): | |
| id_para_batch = id_para.expand(num_frames, -1) | |
| geometry = model_3dmm.get_3dlandmarks( | |
| id_para_batch, exp_para, euler_angle, trans, focal_length, cxy | |
| ) | |
| proj_geo = forward_transform(geometry, euler_angle, trans, focal_length, cxy) | |
| loss_lan = cal_lan_loss(proj_geo[:, :, :2], lms.detach()) | |
| loss = loss_lan | |
| optimizer_frame.zero_grad() | |
| loss.backward() | |
| optimizer_frame.step() | |
| if iter == 1000: | |
| for param_group in optimizer_frame.param_groups: | |
| param_group["lr"] = 0.1 | |
| # if iter % 100 == 0: | |
| # print('pose', iter, loss.item()) | |
| for param_group in optimizer_frame.param_groups: | |
| param_group["lr"] = 0.1 | |
| for iter in range(2000): | |
| id_para_batch = id_para.expand(num_frames, -1) | |
| geometry = model_3dmm.get_3dlandmarks( | |
| id_para_batch, exp_para, euler_angle, trans, focal_length, cxy | |
| ) | |
| proj_geo = forward_transform(geometry, euler_angle, trans, focal_length, cxy) | |
| loss_lan = cal_lan_loss(proj_geo[:, :, :2], lms.detach()) | |
| loss_regid = torch.mean(id_para * id_para) | |
| loss_regexp = torch.mean(exp_para * exp_para) | |
| loss = loss_lan + loss_regid * 0.5 + loss_regexp * 0.4 | |
| optimizer_idexp.zero_grad() | |
| optimizer_frame.zero_grad() | |
| loss.backward() | |
| optimizer_idexp.step() | |
| optimizer_frame.step() | |
| # if iter % 100 == 0: | |
| # print('poseidexp', iter, loss_lan.item(), loss_regid.item(), loss_regexp.item()) | |
| if iter % 1000 == 0 and iter >= 1000: | |
| for param_group in optimizer_idexp.param_groups: | |
| param_group["lr"] *= 0.2 | |
| for param_group in optimizer_frame.param_groups: | |
| param_group["lr"] *= 0.2 | |
| print(loss_lan.item(), torch.mean(trans[:, 2]).item()) | |
| print(f'[INFO] fitting light...') | |
| batch_size = 32 | |
| device_default = torch.device("cuda:0") | |
| device_render = torch.device("cuda:0") | |
| renderer = Render_3DMM(arg_focal, h, w, batch_size, device_render) | |
| sel_ids = np.arange(0, num_frames, int(num_frames / batch_size))[:batch_size] | |
| imgs = [] | |
| for sel_id in sel_ids: | |
| imgs.append(cv2.imread(img_paths[sel_id])[:, :, ::-1]) | |
| imgs = np.stack(imgs) | |
| sel_imgs = torch.as_tensor(imgs).cuda() | |
| sel_lms = lms[sel_ids] | |
| sel_light = light_para.new_zeros((batch_size, 27), requires_grad=True) | |
| set_requires_grad([sel_light]) | |
| optimizer_tl = torch.optim.Adam([tex_para, sel_light], lr=0.1) | |
| optimizer_id_frame = torch.optim.Adam([euler_angle, trans, exp_para, id_para], lr=0.01) | |
| for iter in range(71): | |
| sel_exp_para, sel_euler, sel_trans = ( | |
| exp_para[sel_ids], | |
| euler_angle[sel_ids], | |
| trans[sel_ids], | |
| ) | |
| sel_id_para = id_para.expand(batch_size, -1) | |
| geometry = model_3dmm.get_3dlandmarks( | |
| sel_id_para, sel_exp_para, sel_euler, sel_trans, focal_length, cxy | |
| ) | |
| proj_geo = forward_transform(geometry, sel_euler, sel_trans, focal_length, cxy) | |
| loss_lan = cal_lan_loss(proj_geo[:, :, :2], sel_lms.detach()) | |
| loss_regid = torch.mean(id_para * id_para) | |
| loss_regexp = torch.mean(sel_exp_para * sel_exp_para) | |
| sel_tex_para = tex_para.expand(batch_size, -1) | |
| sel_texture = model_3dmm.forward_tex(sel_tex_para) | |
| geometry = model_3dmm.forward_geo(sel_id_para, sel_exp_para) | |
| rott_geo = forward_rott(geometry, sel_euler, sel_trans) | |
| render_imgs = renderer( | |
| rott_geo.to(device_render), | |
| sel_texture.to(device_render), | |
| sel_light.to(device_render), | |
| ) | |
| render_imgs = render_imgs.to(device_default) | |
| mask = (render_imgs[:, :, :, 3]).detach() > 0.0 | |
| render_proj = sel_imgs.clone() | |
| render_proj[mask] = render_imgs[mask][..., :3].byte() | |
| loss_col = cal_col_loss(render_imgs[:, :, :, :3], sel_imgs.float(), mask) | |
| if iter > 50: | |
| loss = loss_col + loss_lan * 0.05 + loss_regid * 1.0 + loss_regexp * 0.8 | |
| else: | |
| loss = loss_col + loss_lan * 3 + loss_regid * 2.0 + loss_regexp * 1.0 | |
| optimizer_tl.zero_grad() | |
| optimizer_id_frame.zero_grad() | |
| loss.backward() | |
| optimizer_tl.step() | |
| optimizer_id_frame.step() | |
| if iter % 50 == 0 and iter > 0: | |
| for param_group in optimizer_id_frame.param_groups: | |
| param_group["lr"] *= 0.2 | |
| for param_group in optimizer_tl.param_groups: | |
| param_group["lr"] *= 0.2 | |
| # print(iter, loss_col.item(), loss_lan.item(), loss_regid.item(), loss_regexp.item()) | |
| light_mean = torch.mean(sel_light, 0).unsqueeze(0).repeat(num_frames, 1) | |
| light_para.data = light_mean | |
| exp_para = exp_para.detach() | |
| euler_angle = euler_angle.detach() | |
| trans = trans.detach() | |
| light_para = light_para.detach() | |
| print(f'[INFO] fine frame-wise fitting...') | |
| for i in range(int((num_frames - 1) / batch_size + 1)): | |
| if (i + 1) * batch_size > num_frames: | |
| start_n = num_frames - batch_size | |
| sel_ids = np.arange(num_frames - batch_size, num_frames) | |
| else: | |
| start_n = i * batch_size | |
| sel_ids = np.arange(i * batch_size, i * batch_size + batch_size) | |
| imgs = [] | |
| for sel_id in sel_ids: | |
| imgs.append(cv2.imread(img_paths[sel_id])[:, :, ::-1]) | |
| imgs = np.stack(imgs) | |
| sel_imgs = torch.as_tensor(imgs).cuda() | |
| sel_lms = lms[sel_ids] | |
| sel_exp_para = exp_para.new_zeros((batch_size, exp_dim), requires_grad=True) | |
| sel_exp_para.data = exp_para[sel_ids].clone() | |
| sel_euler = euler_angle.new_zeros((batch_size, 3), requires_grad=True) | |
| sel_euler.data = euler_angle[sel_ids].clone() | |
| sel_trans = trans.new_zeros((batch_size, 3), requires_grad=True) | |
| sel_trans.data = trans[sel_ids].clone() | |
| sel_light = light_para.new_zeros((batch_size, 27), requires_grad=True) | |
| sel_light.data = light_para[sel_ids].clone() | |
| set_requires_grad([sel_exp_para, sel_euler, sel_trans, sel_light]) | |
| optimizer_cur_batch = torch.optim.Adam( | |
| [sel_exp_para, sel_euler, sel_trans, sel_light], lr=0.005 | |
| ) | |
| sel_id_para = id_para.expand(batch_size, -1).detach() | |
| sel_tex_para = tex_para.expand(batch_size, -1).detach() | |
| pre_num = 5 | |
| if i > 0: | |
| pre_ids = np.arange(start_n - pre_num, start_n) | |
| for iter in range(50): | |
| geometry = model_3dmm.get_3dlandmarks( | |
| sel_id_para, sel_exp_para, sel_euler, sel_trans, focal_length, cxy | |
| ) | |
| proj_geo = forward_transform(geometry, sel_euler, sel_trans, focal_length, cxy) | |
| loss_lan = cal_lan_loss(proj_geo[:, :, :2], sel_lms.detach()) | |
| loss_regexp = torch.mean(sel_exp_para * sel_exp_para) | |
| sel_geometry = model_3dmm.forward_geo(sel_id_para, sel_exp_para) | |
| sel_texture = model_3dmm.forward_tex(sel_tex_para) | |
| geometry = model_3dmm.forward_geo(sel_id_para, sel_exp_para) | |
| rott_geo = forward_rott(geometry, sel_euler, sel_trans) | |
| render_imgs = renderer( | |
| rott_geo.to(device_render), | |
| sel_texture.to(device_render), | |
| sel_light.to(device_render), | |
| ) | |
| render_imgs = render_imgs.to(device_default) | |
| mask = (render_imgs[:, :, :, 3]).detach() > 0.0 | |
| loss_col = cal_col_loss(render_imgs[:, :, :, :3], sel_imgs.float(), mask) | |
| if i > 0: | |
| geometry_lap = model_3dmm.forward_geo_sub( | |
| id_para.expand(batch_size + pre_num, -1).detach(), | |
| torch.cat((exp_para[pre_ids].detach(), sel_exp_para)), | |
| model_3dmm.rigid_ids, | |
| ) | |
| rott_geo_lap = forward_rott( | |
| geometry_lap, | |
| torch.cat((euler_angle[pre_ids].detach(), sel_euler)), | |
| torch.cat((trans[pre_ids].detach(), sel_trans)), | |
| ) | |
| loss_lap = cal_lap_loss( | |
| [rott_geo_lap.reshape(rott_geo_lap.shape[0], -1).permute(1, 0)], [1.0] | |
| ) | |
| else: | |
| geometry_lap = model_3dmm.forward_geo_sub( | |
| id_para.expand(batch_size, -1).detach(), | |
| sel_exp_para, | |
| model_3dmm.rigid_ids, | |
| ) | |
| rott_geo_lap = forward_rott(geometry_lap, sel_euler, sel_trans) | |
| loss_lap = cal_lap_loss( | |
| [rott_geo_lap.reshape(rott_geo_lap.shape[0], -1).permute(1, 0)], [1.0] | |
| ) | |
| if iter > 30: | |
| loss = loss_col * 0.5 + loss_lan * 1.5 + loss_lap * 100000 + loss_regexp * 1.0 | |
| else: | |
| loss = loss_col * 0.5 + loss_lan * 8 + loss_lap * 100000 + loss_regexp * 1.0 | |
| optimizer_cur_batch.zero_grad() | |
| loss.backward() | |
| optimizer_cur_batch.step() | |
| # if iter % 10 == 0: | |
| # print( | |
| # i, | |
| # iter, | |
| # loss_col.item(), | |
| # loss_lan.item(), | |
| # loss_lap.item(), | |
| # loss_regexp.item(), | |
| # ) | |
| print(str(i) + " of " + str(int((num_frames - 1) / batch_size + 1)) + " done") | |
| render_proj = sel_imgs.clone() | |
| render_proj[mask] = render_imgs[mask][..., :3].byte() | |
| exp_para[sel_ids] = sel_exp_para.clone() | |
| euler_angle[sel_ids] = sel_euler.clone() | |
| trans[sel_ids] = sel_trans.clone() | |
| light_para[sel_ids] = sel_light.clone() | |
| 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("params saved") | |