Class Inpaint_Depth_Net

inpaint/networks.py:134–235

        return h, h_mask

class Inpaint_Depth_Net(nn.Module):
    def __init__(self, layer_size=7, upsampling_mode='nearest'):
        super().__init__()
        in_channels = 4
        out_channels = 1
        self.freeze_enc_bn = False
        self.upsampling_mode = upsampling_mode
        self.layer_size = layer_size
        self.enc_1 = PCBActiv(in_channels, 64, bn=False, sample='down-7', conv_bias=True)
        self.enc_2 = PCBActiv(64, 128, sample='down-5', conv_bias=True)
        self.enc_3 = PCBActiv(128, 256, sample='down-5')
        self.enc_4 = PCBActiv(256, 512, sample='down-3')
        for i in range(4, self.layer_size):
            name = 'enc_{:d}'.format(i + 1)
            setattr(self, name, PCBActiv(512, 512, sample='down-3'))

        for i in range(4, self.layer_size):
            name = 'dec_{:d}'.format(i + 1)
            setattr(self, name, PCBActiv(512 + 512, 512, activ='leaky'))
        self.dec_4 = PCBActiv(512 + 256, 256, activ='leaky')
        self.dec_3 = PCBActiv(256 + 128, 128, activ='leaky')
        self.dec_2 = PCBActiv(128 + 64, 64, activ='leaky')
        self.dec_1 = PCBActiv(64 + in_channels, out_channels,
                              bn=False, activ=None, conv_bias=True)
    def add_border(self, input, mask_flag, PCONV=True):
        with torch.no_grad():
            h = input.shape[-2]
            w = input.shape[-1]
            require_len_unit = 2 ** self.layer_size
            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 mask_flag:
                if PCONV is False:
                    enlarge_input += 1.0
                enlarge_input = enlarge_input.clamp(0.0, 1.0)
            else:
                enlarge_input[:, 2, ...] = 0.0
            anchor_h = residual_h//2
            anchor_w = residual_w//2
            enlarge_input[..., anchor_h:anchor_h+h, anchor_w:anchor_w+w] = input

        return enlarge_input, [anchor_h, anchor_h+h, anchor_w, anchor_w+w]

    def forward_3P(self, mask, context, depth, edge, unit_length=128, cuda=None):
        with torch.no_grad():
            input = torch.cat((depth, edge, context, mask), dim=1)
            n, c, h, w = input.shape
            residual_h = int(np.ceil(h / float(unit_length)) * unit_length - h)
            residual_w = int(np.ceil(w / float(unit_length)) * unit_length - w)
            anchor_h = residual_h//2
            anchor_w = residual_w//2
            enlarge_input = torch.zeros((n, c, h + residual_h, w + residual_w)).to(cuda)
            enlarge_input[..., anchor_h:anchor_h+h, anchor_w:anchor_w+w] = input
            # enlarge_input[:, 3] = 1. - enlarge_input[:, 3]
            depth_output = self.forward(enlarge_input)
            depth_output = depth_output[..., anchor_h:anchor_h+h, anchor_w:anchor_w+w]
            # import pdb; pdb.set_trace()