| 148 | |
| 149 | class Decoder(nn.Module): |
| 150 | def __init__(self, channels): |
| 151 | super(Decoder, self).__init__() |
| 152 | self.config = Config() |
| 153 | DecoderBlock = eval(self.config.dec_blk) |
| 154 | LateralBlock = eval(self.config.lat_blk) |
| 155 | |
| 156 | if self.config.dec_ipt: |
| 157 | self.split = self.config.dec_ipt_split |
| 158 | N_dec_ipt = 64 |
| 159 | DBlock = SimpleConvs |
| 160 | ic = 64 |
| 161 | ipt_cha_opt = 1 |
| 162 | self.ipt_blk5 = DBlock(2**10*3 if self.split else 3, [N_dec_ipt, channels[0]//8][ipt_cha_opt], inter_channels=ic) |
| 163 | self.ipt_blk4 = DBlock(2**8*3 if self.split else 3, [N_dec_ipt, channels[0]//8][ipt_cha_opt], inter_channels=ic) |
| 164 | self.ipt_blk3 = DBlock(2**6*3 if self.split else 3, [N_dec_ipt, channels[1]//8][ipt_cha_opt], inter_channels=ic) |
| 165 | self.ipt_blk2 = DBlock(2**4*3 if self.split else 3, [N_dec_ipt, channels[2]//8][ipt_cha_opt], inter_channels=ic) |
| 166 | self.ipt_blk1 = DBlock(2**0*3 if self.split else 3, [N_dec_ipt, channels[3]//8][ipt_cha_opt], inter_channels=ic) |
| 167 | else: |
| 168 | self.split = None |
| 169 | |
| 170 | self.decoder_block4 = DecoderBlock(channels[0]+([N_dec_ipt, channels[0]//8][ipt_cha_opt] if self.config.dec_ipt else 0), channels[1]) |
| 171 | self.decoder_block3 = DecoderBlock(channels[1]+([N_dec_ipt, channels[0]//8][ipt_cha_opt] if self.config.dec_ipt else 0), channels[2]) |
| 172 | self.decoder_block2 = DecoderBlock(channels[2]+([N_dec_ipt, channels[1]//8][ipt_cha_opt] if self.config.dec_ipt else 0), channels[3]) |
| 173 | self.decoder_block1 = DecoderBlock(channels[3]+([N_dec_ipt, channels[2]//8][ipt_cha_opt] if self.config.dec_ipt else 0), channels[3]//2) |
| 174 | self.conv_out1 = nn.Sequential(nn.Conv2d(channels[3]//2+([N_dec_ipt, channels[3]//8][ipt_cha_opt] if self.config.dec_ipt else 0), 1, 1, 1, 0)) |
| 175 | |
| 176 | self.lateral_block4 = LateralBlock(channels[1], channels[1]) |
| 177 | self.lateral_block3 = LateralBlock(channels[2], channels[2]) |
| 178 | self.lateral_block2 = LateralBlock(channels[3], channels[3]) |
| 179 | |
| 180 | if self.config.ms_supervision: |
| 181 | self.conv_ms_spvn_4 = nn.Conv2d(channels[1], 1, 1, 1, 0) |
| 182 | self.conv_ms_spvn_3 = nn.Conv2d(channels[2], 1, 1, 1, 0) |
| 183 | self.conv_ms_spvn_2 = nn.Conv2d(channels[3], 1, 1, 1, 0) |
| 184 | |
| 185 | if self.config.out_ref: |
| 186 | _N = 16 |
| 187 | self.gdt_convs_4 = nn.Sequential(nn.Conv2d(channels[1], _N, 3, 1, 1), nn.BatchNorm2d(_N) if self.config.batch_size > 1 else nn.Identity(), nn.ReLU(inplace=True)) |
| 188 | self.gdt_convs_3 = nn.Sequential(nn.Conv2d(channels[2], _N, 3, 1, 1), nn.BatchNorm2d(_N) if self.config.batch_size > 1 else nn.Identity(), nn.ReLU(inplace=True)) |
| 189 | self.gdt_convs_2 = nn.Sequential(nn.Conv2d(channels[3], _N, 3, 1, 1), nn.BatchNorm2d(_N) if self.config.batch_size > 1 else nn.Identity(), nn.ReLU(inplace=True)) |
| 190 | |
| 191 | self.gdt_convs_pred_4 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) |
| 192 | self.gdt_convs_pred_3 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) |
| 193 | self.gdt_convs_pred_2 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) |
| 194 | |
| 195 | self.gdt_convs_attn_4 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) |
| 196 | self.gdt_convs_attn_3 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) |
| 197 | self.gdt_convs_attn_2 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) |
| 198 | |
| 199 | |
| 200 | def get_patches_batch(self, x, p): |