Class Inpaint_Edge_Net

inpaint/networks.py:237–330

        return output

class Inpaint_Edge_Net(BaseNetwork):
    def __init__(self, residual_blocks=8, init_weights=True):
        super(Inpaint_Edge_Net, self).__init__()
        in_channels = 7
        out_channels = 1
        self.encoder = []
        # 0
        self.encoder_0 = nn.Sequential(
                            nn.ReflectionPad2d(3),
                            spectral_norm(nn.Conv2d(in_channels=in_channels, out_channels=64, kernel_size=7, padding=0), True),
                            nn.InstanceNorm2d(64, track_running_stats=False),
                            nn.ReLU(True))
        # 1
        self.encoder_1 = nn.Sequential(
                            spectral_norm(nn.Conv2d(in_channels=64, out_channels=128, kernel_size=4, stride=2, padding=1), True),
                            nn.InstanceNorm2d(128, track_running_stats=False),
                            nn.ReLU(True))
        # 2
        self.encoder_2 = nn.Sequential(
                            spectral_norm(nn.Conv2d(in_channels=128, out_channels=256, kernel_size=4, stride=2, padding=1), True),
                            nn.InstanceNorm2d(256, track_running_stats=False),
                            nn.ReLU(True))
        # 3
        blocks = []
        for _ in range(residual_blocks):
            block = ResnetBlock(256, 2)
            blocks.append(block)

        self.middle = nn.Sequential(*blocks)
        # + 3
        self.decoder_0 = nn.Sequential(
                            spectral_norm(nn.ConvTranspose2d(in_channels=256+256, out_channels=128, kernel_size=4, stride=2, padding=1), True),
                            nn.InstanceNorm2d(128, track_running_stats=False),
                            nn.ReLU(True))
        # + 2
        self.decoder_1 = nn.Sequential(
                            spectral_norm(nn.ConvTranspose2d(in_channels=128+128, out_channels=64, kernel_size=4, stride=2, padding=1), True),
                            nn.InstanceNorm2d(64, track_running_stats=False),
                            nn.ReLU(True))
        # + 1
        self.decoder_2 = nn.Sequential(
                            nn.ReflectionPad2d(3),
                            nn.Conv2d(in_channels=64+64, out_channels=out_channels, kernel_size=7, padding=0),
                            )

        if init_weights:
            self.init_weights()

    def add_border(self, input, channel_pad_1=None):
        h = input.shape[-2]
        w = input.shape[-1]
        require_len_unit = 16
        residual_h = int(np.ceil(h / float(require_len_unit)) * require_len_unit - h) # + 2*require_len_unit
        residual_w = int(np.ceil(w / float(require_len_unit)) * require_len_unit - w) # + 2*require_len_unit
        enlarge_input = torch.zeros((input.shape[0], input.shape[1], h + residual_h, w + residual_w)).to(input.device)
        if channel_pad_1 is not None:
            for channel in channel_pad_1:
                enlarge_input[:, channel] = 1