| 106 | |
| 107 | |
| 108 | class UpSampleModule(BaseModule): |
| 109 | def __init__(self, nc_in, nc_out, nf, use_bias, norm, conv_by, conv_type, |
| 110 | use_skip_connection=False): |
| 111 | super().__init__(conv_type) |
| 112 | # Upsample 1 |
| 113 | self.deconv1 = self.DeconvBlock( |
| 114 | nc_in * 2 if use_skip_connection else nc_in, |
| 115 | nf * 2, kernel_size=(3, 3, 3), stride=1, padding=1, |
| 116 | bias=use_bias, norm=norm, conv_by=conv_by) |
| 117 | self.conv9 = self.ConvBlock( |
| 118 | nf * 2, nf * 2, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=1, |
| 119 | bias=use_bias, norm=norm, conv_by=conv_by) |
| 120 | # Upsample 2 |
| 121 | self.deconv2 = self.DeconvBlock( |
| 122 | nf * 4 if use_skip_connection else nf * 2, |
| 123 | nf * 1, kernel_size=(3, 3, 3), stride=1, padding=1, |
| 124 | bias=use_bias, norm=norm, conv_by=conv_by) |
| 125 | self.conv10 = self.ConvBlock( |
| 126 | nf * 1, nf // 2, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=1, |
| 127 | bias=use_bias, norm=norm, conv_by=conv_by) |
| 128 | self.conv11 = self.ConvBlock( |
| 129 | nf // 2, nc_out, kernel_size=(3, 3, 3), stride=(1, 1, 1), |
| 130 | padding=1, bias=use_bias, norm=None, activation=None, conv_by=conv_by) |
| 131 | self.use_skip_connection = use_skip_connection |
| 132 | |
| 133 | def concat_feature(self, ca, cb): |
| 134 | if self.conv_type == 'partial': |
| 135 | ca_feature, ca_mask = ca |
| 136 | cb_feature, cb_mask = cb |
| 137 | feature_cat = torch.cat((ca_feature, cb_feature), 1) |
| 138 | # leave only the later mask |
| 139 | return feature_cat, ca_mask |
| 140 | else: |
| 141 | return torch.cat((ca, cb), 1) |
| 142 | |
| 143 | def forward(self, inp): |
| 144 | c8, c4, c2 = inp |
| 145 | if self.use_skip_connection: |
| 146 | d1 = self.deconv1(self.concat_feature(c8, c4)) |
| 147 | c9 = self.conv9(d1) |
| 148 | d2 = self.deconv2(self.concat_feature(c9, c2)) |
| 149 | else: |
| 150 | d1 = self.deconv1(c8) |
| 151 | c9 = self.conv9(d1) |
| 152 | d2 = self.deconv2(c9) |
| 153 | |
| 154 | c10 = self.conv10(d2) |
| 155 | c11 = self.conv11(c10) |
| 156 | return c11 |
| 157 | |
| 158 | |
| 159 | class UpSampleResNetSkip(BaseModule): |