| 107 | |
| 108 | |
| 109 | class Decoder(nn.Module): |
| 110 | def __init__(self, channels): |
| 111 | super(Decoder, self).__init__() |
| 112 | self.config = Config() |
| 113 | DecoderBlock = eval('BasicDecBlk') |
| 114 | LateralBlock = eval('BasicLatBlk') |
| 115 | |
| 116 | self.decoder_block4 = DecoderBlock(channels[0], channels[1]) |
| 117 | self.decoder_block3 = DecoderBlock(channels[1], channels[2]) |
| 118 | self.decoder_block2 = DecoderBlock(channels[2], channels[3]) |
| 119 | self.decoder_block1 = DecoderBlock(channels[3], channels[3]//2) |
| 120 | |
| 121 | self.lateral_block4 = LateralBlock(channels[1], channels[1]) |
| 122 | self.lateral_block3 = LateralBlock(channels[2], channels[2]) |
| 123 | self.lateral_block2 = LateralBlock(channels[3], channels[3]) |
| 124 | |
| 125 | if self.config.ms_supervision: |
| 126 | self.conv_ms_spvn_4 = nn.Conv2d(channels[1], 1, 1, 1, 0) |
| 127 | self.conv_ms_spvn_3 = nn.Conv2d(channels[2], 1, 1, 1, 0) |
| 128 | self.conv_ms_spvn_2 = nn.Conv2d(channels[3], 1, 1, 1, 0) |
| 129 | self.conv_out1 = nn.Sequential(nn.Conv2d(channels[3]//2, 1, 1, 1, 0)) |
| 130 | |
| 131 | def forward(self, features): |
| 132 | x, x1, x2, x3, x4 = features |
| 133 | outs = [] |
| 134 | p4 = self.decoder_block4(x4) |
| 135 | _p4 = F.interpolate(p4, size=x3.shape[2:], mode='bilinear', align_corners=True) |
| 136 | _p3 = _p4 + self.lateral_block4(x3) |
| 137 | |
| 138 | p3 = self.decoder_block3(_p3) |
| 139 | _p3 = F.interpolate(p3, size=x2.shape[2:], mode='bilinear', align_corners=True) |
| 140 | _p2 = _p3 + self.lateral_block3(x2) |
| 141 | |
| 142 | p2 = self.decoder_block2(_p2) |
| 143 | _p2 = F.interpolate(p2, size=x1.shape[2:], mode='bilinear', align_corners=True) |
| 144 | _p1 = _p2 + self.lateral_block2(x1) |
| 145 | |
| 146 | _p1 = self.decoder_block1(_p1) |
| 147 | _p1 = F.interpolate(_p1, size=x.shape[2:], mode='bilinear', align_corners=True) |
| 148 | p1_out = self.conv_out1(_p1) |
| 149 | |
| 150 | if self.config.ms_supervision: |
| 151 | outs.append(self.conv_ms_spvn_4(p4)) |
| 152 | outs.append(self.conv_ms_spvn_3(p3)) |
| 153 | outs.append(self.conv_ms_spvn_2(p2)) |
| 154 | outs.append(p1_out) |
| 155 | return outs |
| 156 | |
| 157 | |
| 158 | class RefUNet(nn.Module): |