| |
| |
| |
| |
|
|
| import os |
| import pdb |
|
|
| import numpy as np |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
|
|
| class CanoBlendWeightVolume(nn.Module): |
| def __init__(self, data_path): |
| super().__init__() |
|
|
| data = np.load(data_path) |
|
|
| diff_weight_volume = data["diff_weight_volume"] |
| diff_weight_volume = diff_weight_volume.transpose((3, 0, 1, 2))[None] |
| |
|
|
| self.register_buffer( |
| "diff_weight_volume", torch.from_numpy(diff_weight_volume).to(torch.float32) |
| ) |
| self.res_x, self.res_y, self.res_z = self.diff_weight_volume.shape[2:] |
| self.joint_num = self.diff_weight_volume.shape[1] |
|
|
| ori_weight_volume = torch.from_numpy( |
| data["ori_weight_volume"].transpose((3, 0, 1, 2))[None] |
| ).to(torch.float32) |
| self.register_buffer("ori_weight_volume", ori_weight_volume) |
|
|
| if "sdf_volume" in data: |
| smpl_sdf_volume = data["sdf_volume"] |
| if len(smpl_sdf_volume.shape) == 3: |
| smpl_sdf_volume = smpl_sdf_volume[..., None] |
| smpl_sdf_volume = smpl_sdf_volume.transpose((3, 0, 1, 2))[None] |
| self.register_buffer( |
| "smpl_sdf_volume", torch.from_numpy(smpl_sdf_volume).to(torch.float32) |
| ) |
|
|
| volume_bounds = torch.from_numpy(data["volume_bounds"]).to(torch.float32) |
| center = torch.from_numpy(data["center"]).to(torch.float32) |
| smpl_bounds = torch.from_numpy(data["smpl_bounds"]).to(torch.float32) |
|
|
| self.register_buffer("volume_bounds", volume_bounds) |
| self.register_buffer("center", center) |
| self.register_buffer("smpl_bounds", smpl_bounds) |
|
|
| volume_len = self.volume_bounds[1] - self.volume_bounds[0] |
| self.voxel_size = volume_len / torch.tensor( |
| [self.res_x - 1, self.res_y - 1, self.res_z - 1] |
| ).to(volume_len) |
|
|
| def forward(self, pts, requires_scale=True, volume_type="diff"): |
| """ |
| :param pts: (B, N, 3), x y z |
| :param requires_scale: bool, scale pts to [0, 1] |
| :return: (B, N, 55) |
| """ |
| if requires_scale: |
| pts = (pts - self.volume_bounds[None, None, 0]) / ( |
| self.volume_bounds[1] - self.volume_bounds[0] |
| )[None, None] |
|
|
| B, N, _ = pts.shape |
| grid = 2 * pts - 1 |
| grid = grid[..., [2, 1, 0]] |
| grid = grid[:, :, None, None] |
|
|
| weight_volume = ( |
| self.diff_weight_volume if volume_type == "diff" else self.ori_weight_volume |
| ) |
|
|
| base_w = F.grid_sample( |
| weight_volume.expand(B, -1, -1, -1, -1), |
| grid, |
| mode="bilinear", |
| padding_mode="border", |
| align_corners=True, |
| ) |
| base_w = base_w[:, :, :, 0, 0].permute(0, 2, 1) |
| return base_w |
|
|
| def forward_weight_grad(self, pts, requires_scale=True): |
| """ |
| :param pts: (B, N, 3) |
| :param requires_scale: bool, scale pts to [0, 1] |
| :return: (B, N, 24) |
| """ |
| if requires_scale: |
| pts = (pts - self.volume_bounds[None, None, 0]) / ( |
| self.volume_bounds[1] - self.volume_bounds[0] |
| )[None, None] |
| B, N, _ = pts.shape |
| grid = 2 * pts - 1 |
| grid = grid.reshape(-1, 3)[:, [2, 1, 0]] |
| grid = grid[None, :, None, None] |
|
|
| base_g = F.grid_sample( |
| self.base_gradient_volume.view( |
| self.joint_num * 3, self.res_x, self.res_y, self.res_z |
| )[None].expand(B, -1, -1, -1, -1), |
| grid, |
| mode="nearest", |
| padding_mode="border", |
| align_corners=True, |
| ) |
| base_g = base_g[:, :, :, 0, 0].permute(0, 2, 1).reshape(B, N, -1, 3) |
| return base_g |
|
|
| def forward_sdf(self, pts, requires_scale=True): |
| if requires_scale: |
| pts = (pts - self.volume_bounds[None, None, 0]) / ( |
| self.volume_bounds[1] - self.volume_bounds[0] |
| )[None, None] |
| B, N, _ = pts.shape |
| grid = 2 * pts - 1 |
| grid = grid.reshape(-1, 3)[:, [2, 1, 0]] |
| grid = grid[None, :, None, None] |
|
|
| sdf = F.grid_sample( |
| self.smpl_sdf_volume.expand(B, -1, -1, -1, -1), |
| grid, |
| padding_mode="border", |
| align_corners=True, |
| ) |
| sdf = sdf[:, :, :, 0, 0].permute(0, 2, 1) |
|
|
| return sdf |
|
|
