def __init__(self, input_channels, task_channels, task):
super(ScalePredictionModule, self).__init__()
# Per task feature refinement + decoding
if input_channels == task_channels:
channels = input_channels
self.refinement = nn.Sequential(
TaskBasicBlock(channels, channels, task=task),
TaskBasicBlock(channels, channels, task=task))
else:
downsample = nn.Sequential(
# add feature and task estimation
TaskConv2d(input_channels,
task_channels,
1,
bias=False,
task=task),
TaskBatchNorm2d(task_channels, task=task))
self.refinement = nn.Sequential(
TaskBasicBlock(input_channels,
task_channels,
downsample=downsample,
task=task),
TaskBasicBlock(task_channels, task_channels, task=task))
def task_conv1x1(in_planes, out_planes, stride=1, task='shared'):
"""1x1 convolution"""
return TaskConv2d(in_planes,
out_planes,
kernel_size=1,
stride=stride,
bias=False,
task=task)
def task_conv3x3(in_planes,
out_planes,
stride=1,
groups=1,
dilation=1,
task='shared'):
"""3x3 convolution with padding"""
return TaskConv2d(in_planes,
out_planes,
kernel_size=3,
stride=stride,
padding=dilation,
groups=groups,
bias=False,
dilation=dilation,
task=task)
def __init__(self, input_channels):
super(NormalLayer, self).__init__()
self.pred = TaskConv2d(input_channels, 3, 1, task="normals")
self.normalize = Normalize()
def __init__(self, input_channels):
super(DepthLayer, self).__init__()
self.pred = TaskConv2d(input_channels, 1, 1, task="depth")
self.sigmoid = nn.Sigmoid()
def __init__(self, p, backbone, backbone_channels, max_depth=10.):
super(GeoDepthNet, self).__init__()
# General
self.tasks = p.TASKS.NAMES
self.auxilary_tasks = p.AUXILARY_TASKS.NAMES
self.num_scales = len(backbone_channels)
if 'resnet' in p['backbone']:
self.task_channels = [x // 4 for x in backbone_channels]
else:
self.task_channels = backbone_channels
print("task_channels: ", self.task_channels)
self.channels = backbone_channels
self.max_depth = p['max_depth'] if 'max_depth' in p.keys(
) else max_depth
self.use_gt_depth = p['use_gt_depth'] if 'use_gt_depth' in p.keys(
) else False
self.use_guidance = p['use_guidance'] if 'use_guidance' in p.keys(
) else False
self.guidance_reduce = p[
'guidance_reduce'] if 'guidance_reduce' in p.keys() else False
self.normal_loss = p['normal_loss'] if 'normal_loss' in p.keys(
) else False
# Backbone
self.backbone = backbone
# Initial task predictions at multiple scales
################################# Depth branch #############################################
self.scale_0_fea_depth = ScalePredictionModule(
self.channels[0] + self.task_channels[1] + 1 * 1,
self.task_channels[0],
task='depth')
self.scale_0_depth = DepthLayer(self.task_channels[0])
self.scale_1_fea_depth = ScalePredictionModule(
self.channels[1] + self.task_channels[2] + 1 * 1,
self.task_channels[1],
task='depth')
self.scale_1_depth = DepthLayer(self.task_channels[1])
self.scale_2_fea_depth = ScalePredictionModule(
self.channels[2] + self.task_channels[3] + 1 * 1,
self.task_channels[2],
task='depth')
self.scale_2_depth = DepthLayer(self.task_channels[2])
self.scale_3_fea_depth = ScalePredictionModule(self.channels[3],
self.task_channels[3],
task='depth')
self.scale_3_depth = DepthLayer(self.task_channels[3])
################################# Guidance branch #############################################
if self.use_guidance:
self.scale_0_guidance = ScalePredictionModule(
self.channels[0] + self.task_channels[1],
self.task_channels[0],
task='guidance')
if self.guidance_reduce:
self.scale_0_guidance_reduce_dims = nn.Sequential(
TaskConv2d(self.task_channels[0],
3,
1,
bias=True,
task='guidance'),
torch.nn.Sigmoid()) # easy to visualize
self.scale_1_guidance = ScalePredictionModule(
self.channels[1] + self.task_channels[2],
self.task_channels[1],
task='guidance')
self.scale_2_guidance = ScalePredictionModule(
self.channels[2] + self.task_channels[3],
self.task_channels[2],
task='guidance')
self.scale_3_guidance = ScalePredictionModule(
self.channels[3], self.task_channels[3], task='guidance')
# Depth Normal conversion modules
k_size = p['k_size'] if 'k_size' in p.keys() else 5
sample_num = p['sample_num'] if 'sample_num' in p.keys() else 40
self.scale_0_conversion = DepthNormalConversion(
k_size=k_size, dilation=1,
sample_num=sample_num) # Depth2NormalLight