cnywt/SyncTalk

1import torch2import torch.nn as nn3import render_util4import geo_transform5import numpy as np6 7 8def compute_tri_normal(geometry, tris):9    geometry = geometry.permute(0, 2, 1)10    tri_1 = tris[:, 0]11    tri_2 = tris[:, 1]12    tri_3 = tris[:, 2]13 14    vert_1 = torch.index_select(geometry, 2, tri_1)15    vert_2 = torch.index_select(geometry, 2, tri_2)16    vert_3 = torch.index_select(geometry, 2, tri_3)17 18    nnorm = torch.cross(vert_2-vert_1, vert_3-vert_1, 1)19    normal = nn.functional.normalize(nnorm).permute(0, 2, 1)20    return normal21 22 23class Compute_normal_base(torch.autograd.Function):24    @staticmethod25    def forward(ctx, normal):26        normal_b, = render_util.normal_base_forward(normal)27        ctx.save_for_backward(normal)28        return normal_b29 30    @staticmethod31    def backward(ctx, grad_normal_b):32        normal, = ctx.saved_tensors33        grad_normal, = render_util.normal_base_backward(grad_normal_b, normal)34        return grad_normal35 36 37class Normal_Base(torch.nn.Module):38    def __init__(self):39        super(Normal_Base, self).__init__()40 41    def forward(self, normal):42        return Compute_normal_base.apply(normal)43 44 45def preprocess_render(geometry, euler, trans, cam, tris, vert_tris, ori_img):46    point_num = geometry.shape[1]47    rott_geo = geo_transform.euler_trans_geo(geometry, euler, trans)48    proj_geo = geo_transform.proj_geo(rott_geo, cam)49    rot_tri_normal = compute_tri_normal(rott_geo, tris)50    rot_vert_normal = torch.index_select(rot_tri_normal, 1, vert_tris)51    is_visible = -torch.bmm(rot_vert_normal.reshape(-1, 1, 3),52                            nn.functional.normalize(rott_geo.reshape(-1, 3, 1))).reshape(-1, point_num)53    is_visible[is_visible < 0.01] = -154    pixel_valid = torch.zeros((ori_img.shape[0], ori_img.shape[1]*ori_img.shape[2]),55                              dtype=torch.float32, device=ori_img.device)56    return rott_geo, proj_geo, rot_tri_normal, is_visible, pixel_valid57 58 59class Render_Face(torch.autograd.Function):60    @staticmethod61    def forward(ctx, proj_geo, texture, nbl, ori_img, is_visible, tri_inds,62                pixel_valid):63        batch_size, h, w, _ = ori_img.shape64        ori_img = ori_img.view(batch_size, -1, 3)65        ori_size = torch.cat((torch.ones((batch_size, 1), dtype=torch.int32, device=ori_img.device)*h,66                              torch.ones((batch_size, 1), dtype=torch.int32, device=ori_img.device)*w),67                             dim=1).view(-1)68        tri_index, tri_coord, render, real = render_util.render_face_forward(69            proj_geo, ori_img, ori_size, texture, nbl, is_visible, tri_inds, pixel_valid)70        ctx.save_for_backward(ori_img, ori_size, proj_geo, texture, nbl,71                              tri_inds, tri_index, tri_coord)72        return render, real73 74    @staticmethod75    def backward(ctx, grad_render, grad_real):76        ori_img, ori_size, proj_geo, texture, nbl, tri_inds, tri_index, tri_coord = \77            ctx.saved_tensors78        grad_proj_geo, grad_texture, grad_nbl = render_util.render_face_backward(79            grad_render, grad_real, ori_img, ori_size, proj_geo, texture, nbl, tri_inds,80            tri_index, tri_coord)81        return grad_proj_geo, grad_texture, grad_nbl, None, None, None, None82 83 84class Render_RGB(nn.Module):85    def __init__(self):86        super(Render_RGB, self).__init__()87 88    def forward(self, proj_geo, texture, nbl, ori_img, is_visible, tri_inds, pixel_valid):89        return Render_Face.apply(proj_geo, texture, nbl, ori_img, is_visible,90                                 tri_inds, pixel_valid)91 92 93def cal_land(proj_geo, is_visible, lands_info, land_num):94    land_index, = render_util.update_contour(95        lands_info, is_visible, land_num)96    proj_land = torch.index_select(97        proj_geo.reshape(-1, 3), 0, land_index)[:, :2].reshape(-1, land_num, 2)98    return proj_land99 100 101class Render_Land(nn.Module):102    def __init__(self):103        super(Render_Land, self).__init__()104        lands_info = np.loadtxt('../data/3DMM/lands_info.txt', dtype=np.int32)105        self.lands_info = torch.as_tensor(lands_info).cuda()106        tris = np.loadtxt('../data/3DMM/tris.txt', dtype=np.int64)107        self.tris = torch.as_tensor(tris).cuda() - 1108        vert_tris = np.loadtxt('../data/3DMM/vert_tris.txt', dtype=np.int64)109        self.vert_tris = torch.as_tensor(vert_tris).cuda()110        self.normal_baser = Normal_Base().cuda()111        self.renderer = Render_RGB().cuda()112 113    def render_mesh(self, geometry, euler, trans, cam, ori_img, light):114        batch_size, h, w, _ = ori_img.shape115        ori_img = ori_img.view(batch_size, -1, 3)116        ori_size = torch.cat((torch.ones((batch_size, 1), dtype=torch.int32, device=ori_img.device)*h,117                              torch.ones((batch_size, 1), dtype=torch.int32, device=ori_img.device)*w),118                             dim=1).view(-1)119        rott_geo, proj_geo, rot_tri_normal, _, _ = preprocess_render(120            geometry, euler, trans, cam, self.tris, self.vert_tris, ori_img)121        tri_nb = self.normal_baser(rot_tri_normal.contiguous())122        nbl = torch.bmm(tri_nb, (light.reshape(-1, 9, 3))123                        [:, :, 0].unsqueeze(-1).repeat(1, 1, 3))124        texture = torch.ones_like(geometry) * 200125        render, = render_util.render_mesh(126            proj_geo, ori_img, ori_size, texture, nbl, self.tris)127        return render.view(batch_size, h, w, 3).byte()128 129    def cal_loss_rgb(self, geometry, euler, trans, cam, ori_img, light, texture, lands):130        rott_geo, proj_geo, rot_tri_normal, is_visible, pixel_valid = \131            preprocess_render(geometry, euler, trans, cam,132                              self.tris, self.vert_tris, ori_img)133        tri_nb = self.normal_baser(rot_tri_normal.contiguous())134        nbl = torch.bmm(tri_nb, light.reshape(-1, 9, 3))135        render, real = self.renderer(136            proj_geo, texture, nbl, ori_img, is_visible, self.tris, pixel_valid)137        proj_land = cal_land(proj_geo, is_visible,138                             self.lands_info, lands.shape[1])139        col_minus = torch.norm((render-real).reshape(-1, 3),140                               dim=1).reshape(ori_img.shape[0], -1)141        col_dis = torch.mean(col_minus*pixel_valid) / \142            (torch.mean(pixel_valid)+0.00001)143        land_dists = torch.norm(144            (proj_land-lands).reshape(-1, 2), dim=1).reshape(ori_img.shape[0], -1)145        lan_dis = torch.mean(land_dists)146        return col_dis, lan_dis147