(self, x, feat_maps, aabb=None)
| 190 | return feat |
| 191 | |
| 192 | def decode(self, x, feat_maps, aabb=None): |
| 193 | # x [N, 3] |
| 194 | if aabb is None: |
| 195 | aabb = self.aabb |
| 196 | x = 2 * (x - aabb[:3]) / (aabb[3:] - aabb[:3]) - 1 # [-1, 1] |
| 197 | |
| 198 | h_geo = 0 |
| 199 | h_tex = 0 |
| 200 | |
| 201 | coords_list = [[0, 1], [0, 2], [1, 2]] |
| 202 | |
| 203 | geo_feat_maps = [fm[:, :self.geo_feat_dim] for fm in feat_maps] |
| 204 | geo_feat_maps = self.geo_convs(geo_feat_maps) |
| 205 | for i in range(3): |
| 206 | h_geo += self.sample_feature_plane2D(geo_feat_maps[i], x[..., coords_list[i]]) # (N, C) |
| 207 | |
| 208 | if self.use_tex: |
| 209 | tex_feat_maps = [fm[:, self.geo_feat_dim:] for fm in feat_maps] |
| 210 | tex_feat_maps = self.tex_convs(tex_feat_maps) |
| 211 | for i in range(3): |
| 212 | h_tex += self.sample_feature_plane2D(tex_feat_maps[i], x[..., coords_list[i]]) |
| 213 | |
| 214 | h_geo = self.geo_decoder(h_geo) # (N, 1) |
| 215 | if self.use_tex: |
| 216 | h_tex = self.tex_decoder(h_tex).sigmoid() # (N, 1) |
| 217 | h = torch.cat([h_geo, h_tex], dim=1) |
| 218 | else: |
| 219 | h = h_geo |
| 220 | return h |
| 221 | |
| 222 | def forward(self, vol, x, aabb=None): |
| 223 | feat_map = self.encode(vol) |
no test coverage detected