cnywt/SyncTalk

1import torch2import torch.nn as nn3import torch.nn.functional as F4 5 6def compute_tri_normal(geometry, tris):7    tri_1 = tris[:, 0]8    tri_2 = tris[:, 1]9    tri_3 = tris[:, 2]10    vert_1 = torch.index_select(geometry, 1, tri_1)11    vert_2 = torch.index_select(geometry, 1, tri_2)12    vert_3 = torch.index_select(geometry, 1, tri_3)13    nnorm = torch.cross(vert_2-vert_1, vert_3-vert_1, 2)14    normal = nn.functional.normalize(nnorm)15    return normal16 17 18def euler2rot(euler_angle):19    batch_size = euler_angle.shape[0]20    theta = euler_angle[:, 0].reshape(-1, 1, 1)21    phi = euler_angle[:, 1].reshape(-1, 1, 1)22    psi = euler_angle[:, 2].reshape(-1, 1, 1)23    one = torch.ones(batch_size, 1, 1).to(euler_angle.device)24    zero = torch.zeros(batch_size, 1, 1).to(euler_angle.device)25    rot_x = torch.cat((26        torch.cat((one, zero, zero), 1),27        torch.cat((zero, theta.cos(), theta.sin()), 1),28        torch.cat((zero, -theta.sin(), theta.cos()), 1),29    ), 2)30    rot_y = torch.cat((31        torch.cat((phi.cos(), zero, -phi.sin()), 1),32        torch.cat((zero, one, zero), 1),33        torch.cat((phi.sin(), zero, phi.cos()), 1),34    ), 2)35    rot_z = torch.cat((36        torch.cat((psi.cos(), -psi.sin(), zero), 1),37        torch.cat((psi.sin(), psi.cos(), zero), 1),38        torch.cat((zero, zero, one), 1)39    ), 2)40    return torch.bmm(rot_x, torch.bmm(rot_y, rot_z))41 42 43def rot_trans_pts(geometry, rot, trans):44    rott_geo = torch.bmm(rot, geometry.permute(0, 2, 1)) + trans[:, :, None]45    return rott_geo.permute(0, 2, 1)46 47 48def cal_lap_loss(tensor_list, weight_list):49    lap_kernel = torch.Tensor(50        (-0.5, 1.0, -0.5)).unsqueeze(0).unsqueeze(0).float().to(tensor_list[0].device)51    loss_lap = 052    for i in range(len(tensor_list)):53        in_tensor = tensor_list[i]54        in_tensor = in_tensor.view(-1, 1, in_tensor.shape[-1])55        out_tensor = F.conv1d(in_tensor, lap_kernel)56        loss_lap += torch.mean(out_tensor**2)*weight_list[i]57    return loss_lap58 59 60def proj_pts(rott_geo, focal_length, cxy):61    cx, cy = cxy[0], cxy[1]62    X = rott_geo[:, :, 0]63    Y = rott_geo[:, :, 1]64    Z = rott_geo[:, :, 2]65    fxX = focal_length*X66    fyY = focal_length*Y67    proj_x = -fxX/Z + cx68    proj_y = fyY/Z + cy69    return torch.cat((proj_x[:, :, None], proj_y[:, :, None], Z[:, :, None]), 2)70 71 72def forward_transform(geometry, euler_angle, trans, focal_length, cxy):73    rot = euler2rot(euler_angle)74    rott_geo = rot_trans_pts(geometry, rot, trans)75    proj_geo = proj_pts(rott_geo, focal_length, cxy)76    return proj_geo77 78 79def cal_lan_loss(proj_lan, gt_lan):80    return torch.mean((proj_lan-gt_lan)**2)