| 162 | self.aabb = aabb.to(self.geo_encoder.weight.device) |
| 163 | |
| 164 | def encode(self, vol): |
| 165 | geo_feat = self.geo_encoder(vol[:, :1]) |
| 166 | if self.use_tex: |
| 167 | tex_feat = self.tex_encoder(vol) |
| 168 | vol_feat = torch.cat([geo_feat, tex_feat], dim=1) |
| 169 | else: |
| 170 | vol_feat = geo_feat |
| 171 | |
| 172 | xy_feat = vol_feat.mean(dim=4) |
| 173 | xz_feat = vol_feat.mean(dim=3) |
| 174 | yz_feat = vol_feat.mean(dim=2) |
| 175 | |
| 176 | xy_feat = (self.norm(xy_feat) * 0.5).tanh() |
| 177 | xz_feat = (self.norm(xz_feat) * 0.5).tanh() |
| 178 | yz_feat = (self.norm(yz_feat) * 0.5).tanh() |
| 179 | |
| 180 | return [xy_feat, xz_feat, yz_feat] |
| 181 | |
| 182 | def sample_feature_plane2D(self, feat_map, x): |
| 183 | """Sample feature map at given coordinates""" |