| 225 | |
| 226 | |
| 227 | class AutoEncoderGroupPBR(nn.Module): |
| 228 | def __init__(self, geo_feat_channels, tex_feat_channels, feat_channel_up, mlp_hidden_channels, mlp_hidden_layers, use_tex=True, tex_channels=3, posenc=0) -> None: |
| 229 | super().__init__() |
| 230 | self.use_tex = use_tex |
| 231 | |
| 232 | self.geo_encoder = nn.Conv3d(1, geo_feat_channels, kernel_size=4, stride=2, padding=1, bias=True) |
| 233 | if use_tex: |
| 234 | self.tex_encoder = nn.Conv3d(tex_channels + 1, tex_feat_channels, kernel_size=4, stride=2, padding=1, bias=True) |
| 235 | out_channels = geo_feat_channels + tex_feat_channels if use_tex else geo_feat_channels |
| 236 | self.norm = nn.InstanceNorm2d(out_channels) |
| 237 | |
| 238 | self.geo_feat_dim = geo_feat_channels |
| 239 | self.tex_feat_dim = tex_feat_channels |
| 240 | |
| 241 | self.geo_convs = TriplaneGroupResnetBlock( |
| 242 | geo_feat_channels, feat_channel_up, ks=5, input_norm=False, input_act=False |
| 243 | ) |
| 244 | self.geo_decoder = DecoderMLPSkipConcat(feat_channel_up, 1, mlp_hidden_channels, mlp_hidden_layers) |
| 245 | |
| 246 | if use_tex: |
| 247 | self.tex_convs = nn.Sequential( |
| 248 | TriplaneGroupResnetBlock(tex_feat_channels, feat_channel_up, ks=3, input_norm=False, input_act=False), |
| 249 | TriplaneGroupResnetBlock(feat_channel_up, feat_channel_up, ks=3, input_norm=True, input_act=True), |
| 250 | ) |
| 251 | self.rgb_decoder = DecoderMLPSkipConcat(feat_channel_up, 3, mlp_hidden_channels, mlp_hidden_layers, posenc=posenc) |
| 252 | self.mr_decoder = DecoderMLPSkipConcat(feat_channel_up, 2, mlp_hidden_channels, mlp_hidden_layers, posenc=posenc) |
| 253 | self.normal_decoder = DecoderMLPSkipConcat(feat_channel_up, 3, mlp_hidden_channels, mlp_hidden_layers, posenc=posenc) |
| 254 | |
| 255 | self.register_buffer("aabb", torch.tensor([-1, -1, -1, 1, 1, 1], dtype=torch.float32)) |
| 256 | |
| 257 | def geo_parameters(self): |
| 258 | return list(self.geo_encoder.parameters()) + list(self.geo_convs.parameters()) + list(self.geo_decoder.parameters()) |
| 259 | |
| 260 | def tex_parameters(self): |
| 261 | return list(self.tex_encoder.parameters()) + list(self.tex_convs.parameters()) + \ |
| 262 | list(self.rgb_decoder.parameters()) + list(self.mr_decoder.parameters()) + list(self.normal_decoder.parameters()) |
| 263 | |
| 264 | def reset_aabb(self, aabb): |
| 265 | print("set net aabb:", aabb) |
| 266 | if not isinstance(aabb, torch.Tensor): |
| 267 | aabb = torch.tensor(aabb, dtype=torch.float32) |
| 268 | # self.register_buffer("aabb", aabb.to(self.encoder.weight.device)) |
| 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() |