| import torch | |
| import torch.nn as nn | |
| class BodyModelParams(nn.Module): | |
| def __init__(self, num_frames, model_type='smpl'): | |
| super(BodyModelParams, self).__init__() | |
| self.num_frames = num_frames | |
| self.model_type = model_type | |
| self.params_dim = { | |
| 'betas': 10, | |
| 'global_orient': 3, | |
| 'transl': 3, | |
| } | |
| if model_type == 'smpl': | |
| self.params_dim.update({ | |
| 'body_pose': 69, | |
| }) | |
| else: | |
| assert ValueError(f'Unknown model type {model_type}, exiting!') | |
| self.param_names = self.params_dim.keys() | |
| for param_name in self.param_names: | |
| if param_name == 'betas': | |
| param = nn.Embedding(1, self.params_dim[param_name]) | |
| param.weight.data.fill_(0) | |
| param.weight.requires_grad = False | |
| setattr(self, param_name, param) | |
| else: | |
| param = nn.Embedding(num_frames, self.params_dim[param_name]) | |
| param.weight.data.fill_(0) | |
| param.weight.requires_grad = False | |
| setattr(self, param_name, param) | |
| def init_parameters(self, param_name, data, requires_grad=False): | |
| getattr(self, param_name).weight.data = data[..., :self.params_dim[param_name]] | |
| getattr(self, param_name).weight.requires_grad = requires_grad | |
| def set_requires_grad(self, param_name, requires_grad=True): | |
| getattr(self, param_name).weight.requires_grad = requires_grad | |
| def forward(self, frame_ids): | |
| params = {} | |
| for param_name in self.param_names: | |
| if param_name == 'betas': | |
| params[param_name] = getattr(self, param_name)(torch.zeros_like(frame_ids)) | |
| else: | |
| params[param_name] = getattr(self, param_name)(frame_ids) | |
| return params |