| 59 | self.aabb = aabb.to(self.geo_encoder.weight.device) |
| 60 | |
| 61 | def encode(self, vol): |
| 62 | geo_feat = self.geo_encoder(vol[:, :1]) |
| 63 | if self.use_tex: |
| 64 | tex_feat = self.tex_encoder(vol) |
| 65 | vol_feat = torch.cat([geo_feat, tex_feat], dim=1) |
| 66 | else: |
| 67 | vol_feat = geo_feat |
| 68 | |
| 69 | xy_feat = vol_feat.mean(dim=4) |
| 70 | xz_feat = vol_feat.mean(dim=3) |
| 71 | yz_feat = vol_feat.mean(dim=2) |
| 72 | |
| 73 | xy_feat = (self.norm(xy_feat) * 0.5).tanh() |
| 74 | xz_feat = (self.norm(xz_feat) * 0.5).tanh() |
| 75 | yz_feat = (self.norm(yz_feat) * 0.5).tanh() |
| 76 | |
| 77 | return [xy_feat, xz_feat, yz_feat] |
| 78 | |
| 79 | def sample_feature_plane2D(self, feat_map, x): |
| 80 | """Sample feature map at given coordinates""" |