| 269 | self.aabb = aabb.to(self.geo_encoder.weight.device) |
| 270 | |
| 271 | def encode(self, vol): |
| 272 | geo_feat = self.geo_encoder(vol[:, :1]) |
| 273 | if self.use_tex: |
| 274 | tex_feat = self.tex_encoder(vol) |
| 275 | vol_feat = torch.cat([geo_feat, tex_feat], dim=1) |
| 276 | else: |
| 277 | vol_feat = geo_feat |
| 278 | |
| 279 | xy_feat = vol_feat.mean(dim=4) |
| 280 | xz_feat = vol_feat.mean(dim=3) |
| 281 | yz_feat = vol_feat.mean(dim=2) |
| 282 | |
| 283 | xy_feat = (self.norm(xy_feat) * 0.5).tanh() |
| 284 | xz_feat = (self.norm(xz_feat) * 0.5).tanh() |
| 285 | yz_feat = (self.norm(yz_feat) * 0.5).tanh() |
| 286 | |
| 287 | return [xy_feat, xz_feat, yz_feat] |
| 288 | |
| 289 | def sample_feature_plane2D(self, feat_map, x): |
| 290 | """Sample feature map at given coordinates""" |