| 58 | |
| 59 | |
| 60 | class VoxelDecoderBlock(nn.Module): |
| 61 | def __init__(self, in_channels, hidden_channels, hidden_channels_high_res, num_classes, down_xyz, pos_num_freq): |
| 62 | super().__init__() |
| 63 | self.in_channels = in_channels |
| 64 | self.hidden_channels = hidden_channels |
| 65 | self.hidden_channels_high_res = hidden_channels_high_res |
| 66 | self.num_classes = num_classes |
| 67 | self.down_x, self.down_y, self.down_z = down_xyz |
| 68 | self.pos_num_freq = pos_num_freq |
| 69 | |
| 70 | self.conv_in = nn.Sequential( |
| 71 | nn.Conv3d(in_channels, hidden_channels, kernel_size=1), |
| 72 | nn.LeakyReLU(1e-1, inplace=True), |
| 73 | ) |
| 74 | self.conv_layers = nn.ModuleList() |
| 75 | self.conv_layers.append( |
| 76 | nn.Sequential( |
| 77 | nn.Conv3d(self.hidden_channels, self.hidden_channels, kernel_size=3, padding=1), |
| 78 | nn.InstanceNorm3d(self.hidden_channels), |
| 79 | nn.LeakyReLU(1e-1, inplace=True), |
| 80 | nn.Conv3d(self.hidden_channels, self.hidden_channels, kernel_size=3, padding=1), |
| 81 | nn.InstanceNorm3d(self.hidden_channels), |
| 82 | nn.LeakyReLU(1e-1, inplace=True), |
| 83 | ) |
| 84 | ) |
| 85 | self.conv_layers.append( |
| 86 | nn.Sequential( |
| 87 | nn.Conv3d(self.hidden_channels, self.hidden_channels, kernel_size=1), |
| 88 | nn.InstanceNorm3d(self.hidden_channels), |
| 89 | nn.LeakyReLU(1e-1, inplace=True), |
| 90 | nn.Conv3d(self.hidden_channels, self.hidden_channels, kernel_size=1), |
| 91 | nn.InstanceNorm3d(self.hidden_channels), |
| 92 | nn.LeakyReLU(1e-1, inplace=True), |
| 93 | ) |
| 94 | ) |
| 95 | self.transpose_conv = nn.ConvTranspose3d( |
| 96 | self.hidden_channels, self.hidden_channels_high_res, |
| 97 | kernel_size=(2 ** self.down_x, 2 ** self.down_y, 2 ** self.down_z), |
| 98 | stride=(2 ** self.down_x, 2 ** self.down_y, 2 ** self.down_z) |
| 99 | ) |
| 100 | self.final_conv = nn.Sequential( |
| 101 | nn.Conv3d( |
| 102 | self.hidden_channels_high_res + 4 * 2 * self.pos_num_freq, |
| 103 | self.num_classes, kernel_size=1 |
| 104 | ), |
| 105 | nn.InstanceNorm3d(self.num_classes), |
| 106 | nn.LeakyReLU(1e-1, inplace=True), |
| 107 | nn.Conv3d(self.num_classes, self.num_classes, kernel_size=1), |
| 108 | ) |
| 109 | |
| 110 | def forward(self, x): |
| 111 | b, t = x.shape[:2] |
| 112 | x = rearrange(x, 'b t x y z c -> (b t) c x y z') |
| 113 | x = self.conv_in(x) |
| 114 | for conv_layer in self.conv_layers: |
| 115 | x = x + conv_layer(x) |
| 116 | x = self.transpose_conv(x) |
| 117 | x = rearrange(x, '(b t) c x y z -> b t x y z c', b=b, t=t) |