(
self,
resolution: Tuple[int] = 1024,
model_type: Optional[str] = 'smpl',
uv_param_path: Optional[str] = None,
obj_path: Optional[str] = None,
device: Union[torch.device, str] = 'cpu',
threshold_size: int = 512,
# TODO: Solved the sample bug when the resolution is too small.
# set threshold_size is just a temporary solution.
# TODO: add smplx_uv.npz and eval the warping & sampling of smplx
# model.
)
| 26 | class UVRenderer(nn.Module): |
| 27 | """Renderer for SMPL(x) UV map.""" |
| 28 | def __init__( |
| 29 | self, |
| 30 | resolution: Tuple[int] = 1024, |
| 31 | model_type: Optional[str] = 'smpl', |
| 32 | uv_param_path: Optional[str] = None, |
| 33 | obj_path: Optional[str] = None, |
| 34 | device: Union[torch.device, str] = 'cpu', |
| 35 | threshold_size: int = 512, |
| 36 | # TODO: Solved the sample bug when the resolution is too small. |
| 37 | # set threshold_size is just a temporary solution. |
| 38 | |
| 39 | # TODO: add smplx_uv.npz and eval the warping & sampling of smplx |
| 40 | # model. |
| 41 | ): |
| 42 | super().__init__() |
| 43 | self.threshold_size = threshold_size |
| 44 | num_verts = {'smpl': 6890, 'smplx': 10475} |
| 45 | self.NUM_VERTS = num_verts[model_type] |
| 46 | self.device = device |
| 47 | self.resolution = (resolution, resolution) if isinstance( |
| 48 | resolution, int) else resolution |
| 49 | self.uv_param_path = uv_param_path |
| 50 | self.obj_path = obj_path |
| 51 | if uv_param_path is not None: |
| 52 | check_path_suffix(uv_param_path, allowed_suffix=['npz']) |
| 53 | param_dict = dict(np.load(uv_param_path)) |
| 54 | |
| 55 | verts_uv = torch.Tensor(param_dict['verts_uv']) |
| 56 | verts_u, verts_v = torch.unbind(verts_uv, -1) |
| 57 | verts_v_ = 1 - verts_u.unsqueeze(-1) |
| 58 | verts_u_ = verts_v.unsqueeze(-1) |
| 59 | self.verts_uv = torch.cat([verts_u_, verts_v_], -1).to(self.device) |
| 60 | self.faces_uv = torch.LongTensor(param_dict['faces_uv']).to( |
| 61 | self.device) |
| 62 | |
| 63 | self.NUM_VT = self.verts_uv.shape[0] |
| 64 | |
| 65 | self.faces_tensor = torch.LongTensor(param_dict['faces'].astype( |
| 66 | np.int64)).to(self.device) |
| 67 | self.num_faces = self.faces_uv.shape[0] |
| 68 | elif obj_path is not None: |
| 69 | check_path_suffix(obj_path, allowed_suffix=['obj']) |
| 70 | mesh_template = load_objs_as_meshes([obj_path]) |
| 71 | self.faces_uv = mesh_template.textures.faces_uvs_padded()[0].to( |
| 72 | self.device) |
| 73 | self.verts_uv = mesh_template.textures.verts_uvs_padded()[0].to( |
| 74 | self.device) |
| 75 | self.NUM_VT = self.verts_uv.shape[0] |
| 76 | self.faces_tensor = mesh_template.faces_padded()[0].to(self.device) |
| 77 | self.num_faces = self.faces_uv.shape[0] |
| 78 | self.update_fragments() |
| 79 | self.update_face_uv_pixel() |
| 80 | |
| 81 | self = self.to(self.device) |
| 82 | |
| 83 | def to(self, device): |
| 84 | if isinstance(device, str): |
nothing calls this directly
no test coverage detected