(self, x, feat_maps, aabb=None)
| 297 | return feat |
| 298 | |
| 299 | def decode(self, x, feat_maps, aabb=None): |
| 300 | # x [N, 3] |
| 301 | if aabb is None: |
| 302 | aabb = self.aabb |
| 303 | x = 2 * (x - aabb[:3]) / (aabb[3:] - aabb[:3]) - 1 # [-1, 1] |
| 304 | |
| 305 | h_geo = 0 |
| 306 | h_tex = 0 |
| 307 | |
| 308 | coords_list = [[0, 1], [0, 2], [1, 2]] |
| 309 | |
| 310 | geo_feat_maps = [fm[:, :self.geo_feat_dim] for fm in feat_maps] |
| 311 | geo_feat_maps = self.geo_convs(geo_feat_maps) |
| 312 | for i in range(3): |
| 313 | h_geo += self.sample_feature_plane2D(geo_feat_maps[i], x[..., coords_list[i]]) # (N, C) |
| 314 | |
| 315 | if self.use_tex: |
| 316 | tex_feat_maps = [fm[:, self.geo_feat_dim:] for fm in feat_maps] |
| 317 | tex_feat_maps = self.tex_convs(tex_feat_maps) |
| 318 | for i in range(3): |
| 319 | h_tex += self.sample_feature_plane2D(tex_feat_maps[i], x[..., coords_list[i]]) |
| 320 | |
| 321 | h_geo = self.geo_decoder(h_geo) # (N, 1) |
| 322 | if self.use_tex: |
| 323 | h_rgb = self.rgb_decoder(h_tex) # (N, 3) |
| 324 | h_mr = self.mr_decoder(h_tex) # (N, 2) |
| 325 | h_normal = self.normal_decoder(h_tex) # (N, 3) |
| 326 | h = torch.cat([h_geo, h_rgb, h_mr, h_normal], dim=1) |
| 327 | else: |
| 328 | h = h_geo |
| 329 | return h |
| 330 | |
| 331 | def forward(self, vol, x, aabb=None): |
| 332 | feat_map = self.encode(vol) |
no test coverage detected