(self, geo_feat_channels, tex_feat_channels, feat_channel_up, mlp_hidden_channels, mlp_hidden_layers, use_tex=True, tex_channels=3, posenc=0)
| 123 | |
| 124 | class AutoEncoderGroupSkip(nn.Module): |
| 125 | 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: |
| 126 | super().__init__() |
| 127 | self.use_tex = use_tex |
| 128 | |
| 129 | self.geo_encoder = nn.Conv3d(1, geo_feat_channels, kernel_size=4, stride=2, padding=1, bias=True) |
| 130 | if use_tex: |
| 131 | self.tex_encoder = nn.Conv3d(tex_channels + 1, tex_feat_channels, kernel_size=4, stride=2, padding=1, bias=True) |
| 132 | out_channels = geo_feat_channels + tex_feat_channels if use_tex else geo_feat_channels |
| 133 | self.norm = nn.InstanceNorm2d(out_channels) |
| 134 | |
| 135 | self.geo_feat_dim = geo_feat_channels |
| 136 | self.tex_feat_dim = tex_feat_channels |
| 137 | |
| 138 | self.geo_convs = TriplaneGroupResnetBlock( |
| 139 | geo_feat_channels, feat_channel_up, ks=5, input_norm=False, input_act=False |
| 140 | ) |
| 141 | self.geo_decoder = DecoderMLPSkipConcat(feat_channel_up, 1, mlp_hidden_channels, mlp_hidden_layers) |
| 142 | |
| 143 | if use_tex: |
| 144 | self.tex_convs = TriplaneGroupResnetBlock( |
| 145 | tex_feat_channels, feat_channel_up, ks=5, input_norm=False, input_act=False |
| 146 | ) |
| 147 | self.tex_decoder = DecoderMLPSkipConcat(feat_channel_up, tex_channels, mlp_hidden_channels, mlp_hidden_layers, posenc=posenc) |
| 148 | |
| 149 | self.register_buffer("aabb", torch.tensor([-1, -1, -1, 1, 1, 1], dtype=torch.float32)) |
| 150 | |
| 151 | def geo_parameters(self): |
| 152 | return list(self.geo_encoder.parameters()) + list(self.geo_convs.parameters()) + list(self.geo_decoder.parameters()) |
nothing calls this directly
no test coverage detected