ExPose Decoder Decode latent representation to rotation. Args: num_angles (int): Joint num. dtype: dtype. mean (torch.tensor): Mean value for params.
| 147 | |
| 148 | |
| 149 | class ContinuousRotReprDecoder: |
| 150 | """ExPose Decoder Decode latent representation to rotation. |
| 151 | |
| 152 | Args: |
| 153 | num_angles (int): Joint num. |
| 154 | dtype: dtype. |
| 155 | mean (torch.tensor): Mean value for params. |
| 156 | """ |
| 157 | def __init__(self, num_angles, dtype=torch.float32, mean=None): |
| 158 | self.num_angles = num_angles |
| 159 | self.dtype = dtype |
| 160 | |
| 161 | if isinstance(mean, dict): |
| 162 | mean = mean.get('cont_rot_repr', None) |
| 163 | if mean is None: |
| 164 | mean = torch.tensor([1.0, 0.0, 0.0, 1.0, 0.0, 0.0], |
| 165 | dtype=self.dtype).unsqueeze(dim=0).expand( |
| 166 | self.num_angles, -1).contiguous().view(-1) |
| 167 | if not torch.is_tensor(mean): |
| 168 | mean = torch.tensor(mean) |
| 169 | mean = mean.reshape(-1, 6) |
| 170 | |
| 171 | if mean.shape[0] < self.num_angles: |
| 172 | mean = mean.repeat(self.num_angles // mean.shape[0] + 1, |
| 173 | 1).contiguous() |
| 174 | mean = mean[:self.num_angles] |
| 175 | elif mean.shape[0] > self.num_angles: |
| 176 | mean = mean[:self.num_angles] |
| 177 | |
| 178 | mean = mean.reshape(-1) |
| 179 | self.mean = mean |
| 180 | |
| 181 | def get_mean(self): |
| 182 | return self.mean.clone() |
| 183 | |
| 184 | def get_dim_size(self): |
| 185 | return self.num_angles * 6 |
| 186 | |
| 187 | def __call__(self, module_input): |
| 188 | batch_size = module_input.shape[0] |
| 189 | reshaped_input = module_input.view(-1, 6) |
| 190 | rot_mats = rot6d_to_rotmat(reshaped_input) |
| 191 | # aa = rot6d_to_aa(reshaped_input) |
| 192 | # return aa.view(batch_size,-1,3) |
| 193 | return rot_mats.view(batch_size, -1, 3, 3) |
| 194 | |
| 195 | |
| 196 | class ExPoseHead(BaseModule): |