def __init__(self, nf_in, nf_per_level, nf_out, use_skip_sparse=True, use_skip_dense=True):
nn.Module.__init__(self)
assert (type(nf_per_level) is list)
data_dim = 3
self.use_skip_sparse = True
self.use_skip_dense = True
self.use_aspp = True
#self.use_bias = True
self.use_bias = False
modules = []
# 32 - 16
for level in range(len(nf_per_level)):
nf_in = nf_in if level == 0 else nf_per_level[level-1]
modules.append(PreEncoderLayer(nf_in, nf_per_level[level]))
self.process_sparse = nn.Sequential(*modules)
# print("TSDF SPARSE: ", self.process_sparse)
nf = nf_per_level[-1]
# 16 -> 8
nf0 = nf * 2
self.encode_dense0 = nn.Sequential(
nn.Conv3d(nf, nf0, kernel_size=4, stride=2, padding=1, bias=self.use_bias),
nn.BatchNorm3d(nf0),
nn.ReLU(True)
)
# if self.use_aspp:
self.bottleneck = ASPP(nf0, nf0)
# else:
# self.bottleneck = nn.Sequential(
# nn.Conv3d(nf0, nf0, kernel_size=1, bias=self.use_bias),
# nn.BatchNorm3d(nf0),
# nn.ReLU(True)
# )
# 8 -> 16
nf1 = nf0
if self.use_skip_dense:
nf1 += nf0
nf2 = nf1*2 // 3
self.decode_dense0 = nn.Sequential(
nn.ConvTranspose3d(nf1, nf2, kernel_size=4, stride=2, padding=1, bias=self.use_bias),
nn.BatchNorm3d(nf2),
nn.ReLU(True)
)
self.final = nn.Sequential(
nn.Conv3d(nf2, nf_out, kernel_size=1, bias=self.use_bias),
nn.BatchNorm3d(nf_out),
nn.ReLU(True)
)
# occ prediction
self.occpred = nn.Sequential(
nn.Conv3d(nf_out, 1, kernel_size=1, bias=self.use_bias)
)
self.sdfpred = nn.Sequential(
nn.Conv3d(nf_out, 1, kernel_size=1, bias=self.use_bias)
)
def __init__(self,
nf_in,
nf,
truncation=3,
pass_occ=True,
pass_feats=True):
nn.Module.__init__(self)
data_dim = 3
self.pass_occ = pass_occ
self.pass_feats = pass_feats
self.nf_in = nf_in
self.nf = nf
self.truncation = truncation
self.use_aspp = True
self.p1 = nn.Sequential(
nn.ConvTranspose3d(nf_in,
nf,
kernel_size=FSIZE0,
stride=1,
padding=1,
bias=False), nn.BatchNorm3d(nf), nn.ReLU(True))
if self.use_aspp:
self.p2 = ASPP(nf, nf)
self.p3 = nn.Sequential(
nn.Conv3d(nf,
nf,
kernel_size=3,
stride=1,
padding=1,
bias=False), nn.BatchNorm3d(nf),
nn.ReLU(inplace=True))
else:
self.p2 = UNet(nPlanes=[nf, nf, nf], reps=1)
self.p3 = nn.Sequential(
nn.Conv3d(nf * 3,
nf,
kernel_size=3,
stride=1,
padding=1,
bias=False), nn.BatchNorm3d(nf),
nn.ReLU(inplace=True))
# self.p3 = nn.Sequential(
# nn.Conv3d(nf, nf, kernel_size=3, stride=1, padding=1, bias=False),
# nn.BatchNorm3d(nf),
# nn.ReLU(inplace=True)
# )
self.occ = nn.Conv3d(nf,
1,
kernel_size=1,
stride=1,
padding=0,
bias=False)
def __init__(self, nf_in, nf, nf_out):
nn.Module.__init__(self)
data_dim = 3
self.use_aspp = False
if self.use_aspp:
self.aspp = ASPP(nf_in, nf_in)
self.sdf = nn.Conv3d(nf_in, nf_out, kernel_size=1, stride=1, padding=0, bias=False)
else:
self.p1 = nn.Sequential(
nn.ConvTranspose3d(nf_in, nf, kernel_size=FSIZE0, stride=1, padding=1, bias=False),
nn.BatchNorm3d(nf),
nn.ReLU(True)
)
self.p2 = UNet(nPlanes=[nf, nf, nf], reps=1)
self.sdf = nn.Conv3d(nf*3, nf_out, kernel_size=1, stride=1, padding=0, bias=False)
def __init__(self,
nf_in,
nf_per_level,
nf_out,
input_volume_size,
use_skip_sparse=True,
use_skip_dense=True):
nn.Module.__init__(self)
assert (type(nf_per_level) is list)
data_dim = 3
self.use_skip_sparse = use_skip_sparse
self.use_skip_dense = use_skip_dense
self.use_aspp = False
#self.use_bias = True
self.use_bias = False
modules = []
volume_sizes = [(np.array(input_volume_size) // (k + 1)).tolist()
for k in range(len(nf_per_level))]
for level in range(len(nf_per_level)):
nf_in = nf_in if level == 0 else nf_per_level[level - 1]
modules.append(PreEncoderLayer(nf_in, nf_per_level[level]))
self.process_sparse = nn.Sequential(*modules)
# print("TSDF SPARSE: ", self.process_sparse)
nf = nf_per_level[-1]
# 16 -> 8
nf0 = nf * 3 // 2
self.encode_dense0 = nn.Sequential(
nn.Conv3d(nf,
nf0,
kernel_size=4,
stride=2,
padding=1,
bias=self.use_bias), nn.BatchNorm3d(nf0), nn.ReLU(True))
nf1 = nf * 2
if self.use_aspp:
self.aspp = ASPP(nf0, nf1)
# 8 -> 4
self.encode_dense1 = nn.Sequential(
nn.Conv3d(nf1,
nf1,
kernel_size=4,
stride=2,
padding=1,
bias=self.use_bias), nn.BatchNorm3d(nf1),
nn.ReLU(True))
else:
# 8 -> 4
self.encode_dense1 = nn.Sequential(
nn.Conv3d(nf0,
nf1,
kernel_size=4,
stride=2,
padding=1,
bias=self.use_bias), nn.BatchNorm3d(nf1),
nn.ReLU(True))
# 4 -> 4
nf2 = nf1
self.bottleneck_dense2 = nn.Sequential(
nn.Conv3d(nf1, nf2, kernel_size=1, bias=self.use_bias),
nn.BatchNorm3d(nf2), nn.ReLU(True))
# 4 -> 8
nf3 = nf2 if not self.use_skip_dense else nf1 + nf2
nf4 = nf3 // 2
self.decode_dense3 = nn.Sequential(
nn.ConvTranspose3d(nf3,
nf4,
kernel_size=4,
stride=2,
padding=1,
bias=self.use_bias), nn.BatchNorm3d(nf4),
nn.ReLU(True))
# 8 -> 16
if self.use_skip_dense:
nf4 += nf0
nf5 = nf4 // 2
self.decode_dense4 = nn.Sequential(
nn.ConvTranspose3d(nf4,
nf5,
kernel_size=4,
stride=2,
padding=1,
bias=self.use_bias), nn.BatchNorm3d(nf5),
nn.ReLU(True))
self.final = nn.Sequential(
nn.Conv3d(nf5, nf_out, kernel_size=1, bias=self.use_bias),
nn.BatchNorm3d(nf_out), nn.ReLU(True))
# occ prediction
self.occpred = nn.Sequential(
nn.Conv3d(nf_out, 1, kernel_size=1, bias=self.use_bias))
self.sdfpred = nn.Sequential(
nn.Conv3d(nf_out, 1, kernel_size=1, bias=self.use_bias))