| 104 | |
| 105 | |
| 106 | class PCBActiv(nn.Module): |
| 107 | def __init__(self, in_ch, out_ch, bn=True, sample='none-3', activ='relu', |
| 108 | conv_bias=False): |
| 109 | super().__init__() |
| 110 | if sample == 'down-5': |
| 111 | self.conv = PartialConv(in_ch, out_ch, 5, 2, 2, bias=conv_bias) |
| 112 | elif sample == 'down-7': |
| 113 | self.conv = PartialConv(in_ch, out_ch, 7, 2, 3, bias=conv_bias) |
| 114 | elif sample == 'down-3': |
| 115 | self.conv = PartialConv(in_ch, out_ch, 3, 2, 1, bias=conv_bias) |
| 116 | else: |
| 117 | self.conv = PartialConv(in_ch, out_ch, 3, 1, 1, bias=conv_bias) |
| 118 | |
| 119 | if bn: |
| 120 | self.bn = nn.BatchNorm2d(out_ch) |
| 121 | if activ == 'relu': |
| 122 | self.activation = nn.ReLU() |
| 123 | elif activ == 'leaky': |
| 124 | self.activation = nn.LeakyReLU(negative_slope=0.2) |
| 125 | |
| 126 | def forward(self, input, input_mask): |
| 127 | h, h_mask = self.conv(input, input_mask) |
| 128 | if hasattr(self, 'bn'): |
| 129 | h = self.bn(h) |
| 130 | if hasattr(self, 'activation'): |
| 131 | h = self.activation(h) |
| 132 | return h, h_mask |
| 133 | |
| 134 | class Inpaint_Depth_Net(nn.Module): |
| 135 | def __init__(self, layer_size=7, upsampling_mode='nearest'): |