138 lines
5.1 KiB
Python
138 lines
5.1 KiB
Python
"""Aggregation modules - D-DFFNet (with depth side outputs)
|
|
|
|
Original source: D-DFFNet/model/ours/aggregation.py
|
|
"""
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
class BasicConv2d(nn.Module):
|
|
def __init__(self, in_planes, out_planes, kernel_size, stride=1, padding=0, dilation=1):
|
|
super(BasicConv2d, self).__init__()
|
|
self.conv = nn.Conv2d(in_planes, out_planes,
|
|
kernel_size=kernel_size, stride=stride,
|
|
padding=padding, dilation=dilation, bias=False)
|
|
self.bn = nn.BatchNorm2d(out_planes)
|
|
self.relu = nn.ReLU(inplace=True)
|
|
|
|
def forward(self, x):
|
|
x = self.conv(x)
|
|
x = self.bn(x)
|
|
return x
|
|
|
|
|
|
class aggregation_4feat_side_loss_sa_d(nn.Module):
|
|
"""Aggregation module for D-DFFNet (PDnet_ResNest_4F_SL_LA_D)
|
|
|
|
Features 4-level feature fusion with spatial attention,
|
|
side outputs, and depth side outputs.
|
|
"""
|
|
def __init__(self, channel):
|
|
super(aggregation_4feat_side_loss_sa_d, self).__init__()
|
|
self.relu = nn.ReLU(True)
|
|
|
|
self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
|
|
self.conv_upsample1 = BasicConv2d(channel, channel, 3, padding=1)
|
|
self.conv_upsample2 = BasicConv2d(channel, channel, 3, padding=1)
|
|
self.conv_upsample3 = BasicConv2d(channel, channel, 3, padding=1)
|
|
self.conv_upsample4 = BasicConv2d(channel, channel, 3, padding=1)
|
|
self.conv_upsample5 = BasicConv2d(2 * channel, 2 * channel, 3, padding=1)
|
|
self.conv_upsample6 = BasicConv2d(channel, channel, 3, padding=1)
|
|
self.conv_upsample7 = BasicConv2d(3 * channel, 3 * channel, 3, padding=1)
|
|
|
|
self.conv_concat2 = BasicConv2d(2 * channel, 2 * channel, 3, padding=1)
|
|
self.conv_concat3 = BasicConv2d(3 * channel, 3 * channel, 3, padding=1)
|
|
self.conv_concat4 = BasicConv2d(4 * channel, 4 * channel, 3, padding=1)
|
|
|
|
self.conv4 = BasicConv2d(4 * channel, 4 * channel, 3, padding=1)
|
|
self.conv5 = nn.Conv2d(4 * channel, 1, 1)
|
|
|
|
self.side1 = nn.Conv2d(32, 1, 1)
|
|
self.side2 = nn.Conv2d(32, 1, 1)
|
|
self.side3 = nn.Conv2d(32, 1, 1)
|
|
self.side4 = nn.Conv2d(32, 1, 1)
|
|
|
|
self.side5 = nn.Conv2d(32, 1, 1)
|
|
self.side6 = nn.Conv2d(32, 1, 1)
|
|
self.side7 = nn.Conv2d(32, 1, 1)
|
|
self.side8 = nn.Conv2d(32, 1, 1)
|
|
|
|
self.spatial_attention1 = nn.Sequential(
|
|
nn.Conv2d(1, 32, 3, padding=1, bias=False),
|
|
nn.Sigmoid()
|
|
)
|
|
self.spatial_attention2 = nn.Sequential(
|
|
nn.Conv2d(1, 32, 3, padding=1, bias=False),
|
|
nn.Sigmoid()
|
|
)
|
|
self.spatial_attention3 = nn.Sequential(
|
|
nn.Conv2d(1, 32, 3, padding=1, bias=False),
|
|
nn.Sigmoid()
|
|
)
|
|
self.spatial_attention4 = nn.Sequential(
|
|
nn.Conv2d(1, 32, 3, padding=1, bias=False),
|
|
nn.Sigmoid()
|
|
)
|
|
|
|
def forward(self, x1, x2, x3, x4):
|
|
"""
|
|
Args:
|
|
x1: deepest feature (2048 -> 32 channels)
|
|
x2: mid-deep feature (1024 -> 32 channels)
|
|
x3: mid feature (512 -> 32 channels)
|
|
x4: shallow feature (256 -> 32 channels)
|
|
Returns:
|
|
out: final prediction (1 channel)
|
|
side: list of 4 side outputs for training
|
|
depth: list of 4 depth side outputs for training
|
|
"""
|
|
x1_1 = x1
|
|
x2_1 = self.conv_upsample1(self.upsample(x1)) * x2
|
|
x3_1 = self.conv_upsample2(self.upsample(self.upsample(x1))) \
|
|
* self.conv_upsample3(self.upsample(x2)) * x3
|
|
x4_1 = self.conv_upsample2(self.upsample(self.upsample(self.upsample(x1)))) \
|
|
* self.conv_upsample3(self.upsample(self.upsample(x2))) \
|
|
* self.conv_upsample6(self.upsample(x3)) * x4
|
|
|
|
x1_1 = self.conv_upsample4(self.upsample(x1_1))
|
|
|
|
side = []
|
|
depth = []
|
|
|
|
side1 = self.side1(x1_1)
|
|
side2 = self.side2(x2_1)
|
|
side3 = self.side3(x3_1)
|
|
side4 = self.side4(x4_1)
|
|
|
|
side5 = self.side5(x1_1)
|
|
side6 = self.side6(x2_1)
|
|
side7 = self.side7(x3_1)
|
|
side8 = self.side8(x4_1)
|
|
|
|
side.append(side1)
|
|
side.append(side2)
|
|
side.append(side3)
|
|
side.append(side4)
|
|
|
|
depth.append(side5)
|
|
depth.append(side6)
|
|
depth.append(side7)
|
|
depth.append(side8)
|
|
|
|
x1_1 = torch.mul(self.spatial_attention1(side1), x1_1)
|
|
x2_1 = torch.mul(self.spatial_attention2(side2), x2_1)
|
|
x3_1 = torch.mul(self.spatial_attention3(side3), x3_1)
|
|
x4_1 = torch.mul(self.spatial_attention4(side4), x4_1)
|
|
|
|
x2_2 = torch.cat((x2_1, x1_1), 1)
|
|
x2_2 = self.conv_concat2(x2_2)
|
|
x3_2 = torch.cat((x3_1, self.conv_upsample5(self.upsample(x2_2))), 1)
|
|
x3_2 = self.conv_concat3(x3_2)
|
|
x4_2 = torch.cat((x4_1, self.conv_upsample7(self.upsample(x3_2))), 1)
|
|
x4_2 = self.conv_concat4(x4_2)
|
|
|
|
x = self.conv4(x4_2)
|
|
x = self.conv5(x)
|
|
|
|
return x, side, depth
|