def __init__(self, version, out_channels, ni, n1, n2, n3, n4, n5,
unpack_kernel, iconv_kernel, num_3d_feat):
super().__init__()
# Decoder
self.version = version
n1o, n1i = n1, n1 + ni + out_channels
n2o, n2i = n2, n2 + n1 + out_channels
n3o, n3i = n3, n3 + n2 + out_channels
n4o, n4i = n4, n4 + n3
n5o, n5i = n5, n5 + n4
self.unpack5 = UnpackLayerConv3d(n5,
n5o,
unpack_kernel[0],
d=num_3d_feat)
self.unpack4 = UnpackLayerConv3d(n5,
n4o,
unpack_kernel[1],
d=num_3d_feat)
self.unpack3 = UnpackLayerConv3d(n4,
n3o,
unpack_kernel[2],
d=num_3d_feat)
self.unpack2 = UnpackLayerConv3d(n3,
n2o,
unpack_kernel[3],
d=num_3d_feat)
self.unpack1 = UnpackLayerConv3d(n2,
n1o,
unpack_kernel[4],
d=num_3d_feat)
self.iconv5 = Conv2D(n5i, n5, iconv_kernel[0], 1)
self.iconv4 = Conv2D(n4i, n4, iconv_kernel[1], 1)
self.iconv3 = Conv2D(n3i, n3, iconv_kernel[2], 1)
self.iconv2 = Conv2D(n2i, n2, iconv_kernel[3], 1)
self.iconv1 = Conv2D(n1i, n1, iconv_kernel[4], 1)
# Depth Layers
self.unpack_disps = nn.PixelShuffle(2)
self.unpack_disp4 = nn.Upsample(scale_factor=2,
mode='nearest',
align_corners=None)
self.unpack_disp3 = nn.Upsample(scale_factor=2,
mode='nearest',
align_corners=None)
self.unpack_disp2 = nn.Upsample(scale_factor=2,
mode='nearest',
align_corners=None)
self.disp4_layer = InvDepth(n4, out_channels=out_channels)
self.disp3_layer = InvDepth(n3, out_channels=out_channels)
self.disp2_layer = InvDepth(n2, out_channels=out_channels)
self.disp1_layer = InvDepth(n1, out_channels=out_channels)
def __init__(self, version, in_channels, ni, n1, n2, n3, n4, n5,
pack_kernel, num_blocks, num_3d_feat, dropout):
super().__init__()
# Encoder
self.version = version
self.pre_calc = Conv2D(in_channels, ni, 5, 1)
self.pack1 = PackLayerConv3d(n1, pack_kernel[0], d=num_3d_feat)
self.pack2 = PackLayerConv3d(n2, pack_kernel[1], d=num_3d_feat)
self.pack3 = PackLayerConv3d(n3, pack_kernel[2], d=num_3d_feat)
self.pack4 = PackLayerConv3d(n4, pack_kernel[3], d=num_3d_feat)
self.pack5 = PackLayerConv3d(n5, pack_kernel[4], d=num_3d_feat)
self.conv1 = Conv2D(ni, n1, 7, 1)
self.conv2 = ResidualBlock(n1, n2, num_blocks[0], 1, dropout=dropout)
self.conv3 = ResidualBlock(n2, n3, num_blocks[1], 1, dropout=dropout)
self.conv4 = ResidualBlock(n3, n4, num_blocks[2], 1, dropout=dropout)
self.conv5 = ResidualBlock(n4, n5, num_blocks[3], 1, dropout=dropout)
def __init__(self, dropout=None, version=None, **kwargs):
super().__init__()
self.version = version[1:]
# Input/output channels
in_channels = 3
out_channels = 1
# Hyper-parameters
ni, no = 32, out_channels
n1, n2, n3, n4, n5 = 32, 32, 64, 128, 256
num_blocks = [2, 2, 3, 3]
pack_kernel = [5, 3, 3, 3, 3]
unpack_kernel = [3, 3, 3, 3, 3]
iconv_kernel = [3, 3, 3, 3, 3]
# Initial convolutional layer
self.pre_calc = Conv2D(in_channels, ni, 5, 1)
# Support for different versions
if self.version == 'A': # Channel concatenation
n1o, n1i = n1, n1 + ni + no
n2o, n2i = n2, n2 + n1 + no
n3o, n3i = n3, n3 + n2 + no
n4o, n4i = n4, n4 + n3
n5o, n5i = n5, n5 + n4
elif self.version == 'B': # Channel addition
n1o, n1i = n1, n1 + no
n2o, n2i = n2, n2 + no
n3o, n3i = n3 // 2, n3 // 2 + no
n4o, n4i = n4 // 2, n4 // 2
n5o, n5i = n5 // 2, n5 // 2
else:
raise ValueError('Unknown PackNet version {}'.format(version))
# Encoder
self.pack1 = PackLayerConv3d(n1, pack_kernel[0], 1)
self.pack2 = PackLayerConv3d(n2, pack_kernel[1], 1)
self.pack3 = PackLayerConv3d(n3, pack_kernel[2], 1)
self.pack4 = PackLayerConv3d(n4, pack_kernel[3], 1)
self.pack5 = PackLayerConv3d(n5, pack_kernel[4], 1)
self.conv1 = Conv2D(ni, n1, 7, 1)
self.conv2 = ResidualBlock(n1, n2, num_blocks[0], 1, dropout=dropout)
self.conv3 = ResidualBlock(n2, n3, num_blocks[1], 1, dropout=dropout)
self.conv4 = ResidualBlock(n3, n4, num_blocks[2], 1, dropout=dropout)
self.conv5 = ResidualBlock(n4, n5, num_blocks[3], 1, dropout=dropout)
# Decoder
self.unpack5 = UnpackLayerConv2d(n5, n5o, unpack_kernel[0])
self.unpack4 = UnpackLayerConv2d(n5, n4o, unpack_kernel[1])
self.unpack3 = UnpackLayerConv2d(n4, n3o, unpack_kernel[2])
self.unpack2 = UnpackLayerConv2d(n3, n2o, unpack_kernel[3])
self.unpack1 = UnpackLayerConv2d(n2, n1o, unpack_kernel[4])
self.iconv5 = Conv2D(n5i, n5, iconv_kernel[0], 1)
self.iconv4 = Conv2D(n4i, n4, iconv_kernel[1], 1)
self.iconv3 = Conv2D(n3i, n3, iconv_kernel[2], 1)
self.iconv2 = Conv2D(n2i, n2, iconv_kernel[3], 1)
self.iconv1 = Conv2D(n1i, n1, iconv_kernel[4], 1)
# Depth Layers
self.unpack_disps = nn.PixelShuffle(2)
self.unpack_disp4 = nn.Upsample(scale_factor=2,
mode='nearest',
align_corners=None)
self.unpack_disp3 = nn.Upsample(scale_factor=2,
mode='nearest',
align_corners=None)
self.unpack_disp2 = nn.Upsample(scale_factor=2,
mode='nearest',
align_corners=None)
self.disp4_layer = InvDepth(n4, out_channels=out_channels)
self.disp3_layer = InvDepth(n3, out_channels=out_channels)
self.disp2_layer = InvDepth(n2, out_channels=out_channels)
self.disp1_layer = InvDepth(n1, out_channels=out_channels)
self.init_weights()