pengsida/NeuralBody

1import torch2import torch.nn.functional as F3import numpy as np4 5 6def ppts_to_pts(ppts, bw, A):7    """transform points from the pose space to the zero space"""8    sh = ppts.shape9    bw = bw.permute(0, 2, 1)10    A = torch.bmm(bw, A.view(sh[0], 24, -1))11    A = A.view(sh[0], -1, 4, 4)12    pts = ppts - A[..., :3, 3]13    R_inv = torch.inverse(A[..., :3, :3])14    pts = torch.sum(R_inv * pts[:, :, None], dim=3)15    return pts16 17 18def grid_sample_blend_weights(grid_coords, bw):19    # the blend weight is indexed by xyz20    grid_coords = grid_coords[:, None, None]21    bw = F.grid_sample(bw,22                       grid_coords,23                       padding_mode='border',24                       align_corners=True)25    bw = bw[:, :, 0, 0]26    return bw27 28 29def bounds_grid_sample_blend_weights(pts, bw, bounds):30    """grid sample blend weights"""31    pts = pts.clone()32 33    # interpolate blend weights34    min_xyz = bounds[:, 0]35    max_xyz = bounds[:, 1]36    bounds = max_xyz[:, None] - min_xyz[:, None]37    grid_coords = (pts - min_xyz[:, None]) / bounds38    grid_coords = grid_coords * 2 - 139    # convert xyz to zyx, since the blend weight is indexed by xyz40    grid_coords = grid_coords[..., [2, 1, 0]]41 42    # the blend weight is indexed by xyz43    bw = bw.permute(0, 4, 1, 2, 3)44    grid_coords = grid_coords[:, None, None]45    bw = F.grid_sample(bw,46                       grid_coords,47                       padding_mode='border',48                       align_corners=True)49    bw = bw[:, :, 0, 0]50 51    return bw52 53 54def grid_sample_A_blend_weights(nf_grid_coords, bw):55    """56    nf_grid_coords: batch_size x N_samples x 24 x 357    bw: batch_size x 24 x 64 x 64 x 6458    """59    bws = []60    for i in range(24):61        nf_grid_coords_ = nf_grid_coords[:, :, i]62        nf_grid_coords_ = nf_grid_coords_[:, None, None]63        bw_ = F.grid_sample(bw[:, i:i + 1],64                            nf_grid_coords_,65                            padding_mode='border',66                            align_corners=True)67        bw_ = bw_[:, :, 0, 0]68        bws.append(bw_)69    bw = torch.cat(bws, dim=1)70    return bw71 72 73def ppts_to_pts(pts, bw, A):74    """transform points from the pose space to the t pose"""75    sh = pts.shape76    bw = bw.permute(0, 2, 1)77    A = torch.bmm(bw, A.view(sh[0], 24, -1))78    A = A.view(sh[0], -1, 4, 4)79    pts = pts - A[..., :3, 3]80    R_inv = torch.inverse(A[..., :3, :3])81    pts = torch.sum(R_inv * pts[:, :, None], dim=3)82    return pts83