| 365 | |
| 366 | |
| 367 | class SMPLLayer(SMPL): |
| 368 | def __init__(self, *args, **kwargs) -> None: |
| 369 | # Just create a SMPL module without any member variables |
| 370 | super(SMPLLayer, self).__init__( |
| 371 | create_body_pose=False, |
| 372 | create_betas=False, |
| 373 | create_global_orient=False, |
| 374 | create_transl=False, |
| 375 | *args, |
| 376 | **kwargs, |
| 377 | ) |
| 378 | |
| 379 | def forward(self, |
| 380 | betas: Optional[Tensor] = None, |
| 381 | body_pose: Optional[Tensor] = None, |
| 382 | global_orient: Optional[Tensor] = None, |
| 383 | transl: Optional[Tensor] = None, |
| 384 | return_verts=True, |
| 385 | return_full_pose: bool = False, |
| 386 | pose2rot: bool = True, |
| 387 | **kwargs) -> SMPLOutput: |
| 388 | """Forward pass for the SMPL model. |
| 389 | |
| 390 | Parameters |
| 391 | ---------- |
| 392 | global_orient: torch.tensor, optional, shape Bx3 |
| 393 | If given, ignore the member variable and use it as the global |
| 394 | rotation of the body. Useful if someone wishes to predicts this |
| 395 | with an external model. (default=None) |
| 396 | betas: torch.tensor, optional, shape Bx10 |
| 397 | If given, ignore the member variable `betas` and use it |
| 398 | instead. For example, it can used if shape parameters |
| 399 | `betas` are predicted from some external model. |
| 400 | (default=None) |
| 401 | body_pose: torch.tensor, optional, shape Bx(J*3) |
| 402 | If given, ignore the member variable `body_pose` and use it |
| 403 | instead. For example, it can used if someone predicts the |
| 404 | pose of the body joints are predicted from some external model. |
| 405 | It should be a tensor that contains joint rotations in |
| 406 | axis-angle format. (default=None) |
| 407 | transl: torch.tensor, optional, shape Bx3 |
| 408 | If given, ignore the member variable `transl` and use it |
| 409 | instead. For example, it can used if the translation |
| 410 | `transl` is predicted from some external model. |
| 411 | (default=None) |
| 412 | return_verts: bool, optional |
| 413 | Return the vertices. (default=True) |
| 414 | return_full_pose: bool, optional |
| 415 | Returns the full axis-angle pose vector (default=False) |
| 416 | |
| 417 | Returns |
| 418 | ------- |
| 419 | """ |
| 420 | device, dtype = self.shapedirs.device, self.shapedirs.dtype |
| 421 | if global_orient is None: |
| 422 | batch_size = 1 |
| 423 | global_orient = torch.zeros(3, device=device, dtype=dtype).view( |
| 424 | 1, 1, 3).expand(batch_size, 1, 1).contiguous() |