| 11 | |
| 12 | |
| 13 | class PoseDecoder(nn.Module): |
| 14 | def __init__( |
| 15 | self, |
| 16 | hidden_size=768, |
| 17 | mlp_ratio=4, |
| 18 | pose_encoding_type="absT_quaR", |
| 19 | ): |
| 20 | super().__init__() |
| 21 | |
| 22 | self.pose_encoding_type = pose_encoding_type |
| 23 | if self.pose_encoding_type == "absT_quaR": |
| 24 | self.target_dim = 7 |
| 25 | |
| 26 | self.mlp = Mlp( |
| 27 | in_features=hidden_size, |
| 28 | hidden_features=int(hidden_size * mlp_ratio), |
| 29 | out_features=self.target_dim, |
| 30 | drop=0, |
| 31 | ) |
| 32 | |
| 33 | def forward( |
| 34 | self, |
| 35 | pose_feat, |
| 36 | ): |
| 37 | """ |
| 38 | pose_feat: BxC |
| 39 | preliminary_cameras: cameras in opencv coordinate. |
| 40 | """ |
| 41 | |
| 42 | pred_cameras = self.mlp(pose_feat) # Bx7, 3 for absT, 4 for quaR |
| 43 | return pred_cameras |
| 44 | |
| 45 | |
| 46 | class PoseEncoder(nn.Module): |