def __init__(self,
in_channels_r, out_channels_r, in_channels_u, out_channels,
kernel_size_in, kernel_size_out,
up_stride_in, stride_out,
up_padding_in, padding_out, output_padding=0,
activation_in='lrelu', activation_out='lrelu',
norm_in='bn', norm_out='none'):
"""
u {l-1} and r {l} new should be of equal sizes to be concatenated
:param in_channels_r: in channels for r {l}
:param out_channels_r: out channels for r {l-1}
:param in_channels_u: in channels for u {l-1}
:param out_channels: out channels for (u {l-1}, r {l}) -> r {l-1}
:param kernel_size_in: kernel size for r {l}) -> r {l} new
:param kernel_size_out: kernel size for (u {l-1}, r {l}) -> r {l-1}
:param up_stride_in: stride for transposed conv r {l}) -> r {l} new
:param stride_out: stride for (u {l-1}, r {l}) -> r {l-1}
:param up_padding_in: padding for transposed convolution r {l} -> r {l} new
:param output_padding: output padding for transposed convolution r {l} -> r {l} new
:param padding_out: padding for (u {l-1}, r {l}) -> r {l-1}
:param activation_in: activation layer for r {l}) -> r {l} new
:param activation_out: activation layer for (u {l-1}, r {l}) -> r {l-1}
:param norm_in: normalization layer for r {l}) -> r {l} new
:param norm_out: normalization layer for (u {l-1}, r {l}) -> r {l-1}
"""
super(RecovecyBlock, self).__init__()
self.in_upconv = upconv_block(
in_channels=in_channels_r,
out_channels=out_channels_r,
kernel_size=kernel_size_in,
stride=up_stride_in,
padding=up_padding_in,
output_padding=output_padding,
norm=norm_in,
activation=activation_in
)
self.out_conv = conv_block(
in_channels=out_channels_r + in_channels_u,
out_channels=out_channels,
kernel_size=kernel_size_out,
stride=stride_out,
padding=padding_out,
norm=norm_out,
activation=activation_out
)
def __init__(self,
in_channels_1, in_channels_2, out_channels,
kernel_size_1, kernel_size_2,
stride_1, up_stride_2,
padding_1, up_padding_2, output_padding=0,
activation_in='relu', activation_out='lrelu',
norm_in='bn', norm_out='none'):
"""
1 convolution - from f {l-1}{r-1}
2 transposed convolution - from f {l} {k}
3 convolution - the one to be summed with result of f {l}{k} (2) convolution
4 convolution - the one to be producted with result of f {l}{k} (2) convolution
:param in_channels_1: Input channels of the 1st convolution layer
:param in_channels_2: Input channels of the 2st convolution layer
:param out_channels: Output channels of the 1st, 2st, 3rd, 4th convolution layers
:param kernel_size_1: Kernel size of 1st convolution layer
:param kernel_size_2: Kernel size of 2nd transposed convolution layer
:param stride_1: Stride of 1st convolution layer
:param up_stride_2: Stride of 2nd transposed convolution layer
:param padding_1: Padding of 1st convolution layer
:param up_padding_2: Padding of 2nd transposed convolution layer
:param output_padding: Output padding of 2nd transposed convolution layer
:param activation_in: Activation layer of 1st convolution layer
:param activation_out: Activation layer 2nd transposed convolution layer
:param norm_in: Normalization layer of 1st convolution layer
:param norm_out: Normalization layer of 2nd transposed convolution layer
"""
super(RefinementBlock, self).__init__()
self.conv_1 = conv_block(
in_channels=in_channels_1,
out_channels=out_channels,
kernel_size=kernel_size_1,
stride=stride_1,
padding=padding_1,
norm='none',
activation=activation_in
)
self.upconv_2 = upconv_block(
in_channels=in_channels_2,
out_channels=out_channels,
kernel_size=kernel_size_2,
stride=up_stride_2,
padding=up_padding_2,
output_padding=output_padding,
norm=norm_in,
activation='none'
)
self.conv_3 = conv_block(
in_channels=out_channels,
out_channels=out_channels,
kernel_size=3,
stride=1,
padding=1,
norm='none',
activation='none'
)
self.conv_4 = conv_block(
in_channels=out_channels,
out_channels=out_channels,
kernel_size=3,
stride=1,
padding=1,
norm='none',
activation='none'
)
self.out_act = _activation(act_type=activation_out)
self.out_norm = _norm(norm_type=norm_out, channels=out_channels)
def __init__(self,
in_nc=4,
out_nc=3,
res_blocks_num=8,
c_num=48,
dmfb_blocks_num=8,
norm='in',
activation='relu',
device=None):
super(FineGenerator, self).__init__()
self.encoder = nn.ModuleList([
# [4, 256, 256]
conv_block(in_channels=in_nc,
out_channels=c_num,
kernel_size=5,
stride=1,
padding=2,
norm=norm,
activation=activation),
# [c_num, 256, 256] -> decoder
conv_block(in_channels=c_num,
out_channels=2 * c_num,
kernel_size=3,
stride=2,
padding=1,
norm=norm,
activation=activation),
# [c_num * 2, 128, 128]
conv_block(in_channels=2 * c_num,
out_channels=2 * c_num,
kernel_size=3,
stride=1,
padding=1,
norm=norm,
activation=activation),
# [c_num * 2, 128, 128] -> decoder
conv_block(in_channels=2 * c_num,
out_channels=4 * c_num,
kernel_size=3,
stride=2,
padding=1,
norm=norm,
activation=activation),
# [c_num * 4, 64, 64]
conv_block(in_channels=4 * c_num,
out_channels=4 * c_num,
kernel_size=3,
stride=1,
padding=1,
norm=norm,
activation=activation),
# [c_num * 4, 64, 64]
conv_block(in_channels=4 * c_num,
out_channels=4 * c_num,
kernel_size=3,
stride=1,
padding=1,
norm=norm,
activation=activation),
# [c_num * 4, 64, 64]
])
self.fine_bn = FineBottleneck(c_num=4 * c_num,
dmfb_blocks_num=dmfb_blocks_num,
res_blocks_num=res_blocks_num,
device=device)
self.decoder = nn.ModuleList([
# [c_num * 4, 64, 64]
conv_block(in_channels=c_num * 4,
out_channels=c_num * 4,
kernel_size=3,
stride=1,
padding=1),
# [c_num * 4, 64, 64]
conv_block(in_channels=c_num * 4,
out_channels=c_num * 4,
kernel_size=3,
stride=1,
padding=1,
norm=norm,
activation=activation),
# [c_num * 4, 64, 64]
upconv_block(in_channels=c_num * 4,
out_channels=c_num * 2,
kernel_size=3,
stride=2,
output_padding=1,
padding=get_pad_tp(64, 64, [1, 1], [3, 3], [2, 2],
[1, 1]),
norm=norm,
activation=activation),
# encoder 3 + skip [c_num * 2, 128, 128] -> decoder 4
# [c_num * 4, 128, 128]
conv_block(in_channels=c_num * 4,
out_channels=c_num * 2,
kernel_size=3,
stride=1,
padding=1,
norm=norm,
activation=activation),
# [c_num * 2, 128, 128]
upconv_block(in_channels=c_num * 2,
out_channels=c_num,
kernel_size=3,
stride=2,
output_padding=1,
padding=get_pad_tp(128, 128, [1, 1], [3, 3], [2, 2],
[1, 1]),
norm=norm,
activation=activation),
# encoder 1 + skip [c_num, 256, 256] -> decoder 6
# [c_num, 256, 256]
conv_block(in_channels=c_num * 2,
out_channels=c_num // 2,
kernel_size=3,
stride=1,
padding=1,
norm=norm,
activation=activation),
# [c_num//2, 256, 256]
conv_block(in_channels=c_num // 2,
out_channels=out_nc,
kernel_size=3,
stride=1,
padding=1,
norm='none',
activation='tanh'),
# [out_nc, 256, 256]
])
def __init__(self,
in_nc=4,
c_num=48,
out_nc=3,
res_blocks_num=8,
norm="in",
activation="relu"):
super(CoarseGenerator, self).__init__()
self.encoder_coarse = nn.ModuleList([
# [4, 256, 256]
conv_block(in_nc,
c_num,
kernel_size=5,
stride=1,
padding=2,
norm=norm,
activation=activation),
# [c_num, 256, 256] -> decoder_6
conv_block(c_num,
c_num * 2,
kernel_size=3,
stride=2,
padding=1,
norm=norm,
activation=activation),
# [c_num * 2, 128, 128]
conv_block(c_num * 2,
c_num * 2,
kernel_size=3,
stride=1,
padding=1,
norm=norm,
activation=activation),
# [c_num * 2, 128, 128] -> decoder_4
conv_block(c_num * 2,
c_num * 4,
kernel_size=3,
stride=2,
padding=1,
norm=norm,
activation=activation),
# [c_num * 4, 64, 64]
conv_block(c_num * 4,
c_num * 4,
kernel_size=3,
stride=1,
padding=1,
norm=norm,
activation=activation),
# [c_num * 4, 64, 64]
conv_block(c_num * 4,
c_num * 4,
kernel_size=3,
stride=1,
padding=1,
norm=norm,
activation=activation)
# [c_num * 4, 64, 64]]
])
blocks = []
for _ in range(res_blocks_num):
block = B.ResConv(4 * c_num, kernel_size=3, dilation=2,
padding=2) # [192, 64, 64]
blocks.append(block)
self.coarse_bn = nn.Sequential(*blocks)
self.decoder_coarse = nn.ModuleList([
conv_block(c_num * 4,
c_num * 4,
kernel_size=3,
stride=1,
padding=1,
norm=norm,
activation=activation),
# [c_num * 4, 64, 64]
conv_block(c_num * 4,
c_num * 4,
kernel_size=3,
stride=1,
padding=1,
norm=norm,
activation=activation),
# [c_num * 4, 64, 64]
upconv_block(c_num * 4,
c_num * 2,
kernel_size=3,
stride=2,
output_padding=1,
padding=get_pad_tp(64, 64, [1, 1], [3, 3], [2, 2],
[1, 1]),
norm=norm,
activation=activation),
# [c_num * 2, 128, 128] + skip [c_num * 2, 128, 128]
conv_block(c_num * 4,
c_num * 2,
kernel_size=3,
stride=1,
padding=1,
norm=norm,
activation=activation),
# [c_num * 2, 128, 128]
upconv_block(c_num * 2,
c_num,
kernel_size=3,
stride=2,
output_padding=1,
padding=get_pad_tp(128, 128, [1, 1], [3, 3], [2, 2],
[1, 1]),
norm=norm,
activation=activation),
# [c_num, 256, 256] + skip [c_num, 256, 256]
conv_block(c_num * 2,
c_num,
kernel_size=3,
stride=1,
padding=1,
norm=norm,
activation=activation),
# [c_num, 256, 256]
conv_block(c_num,
c_num // 2,
kernel_size=3,
stride=1,
padding=1,
norm=norm,
activation=activation)
# [c_num // 2, 256, 256]
])
self.out_coarse = nn.Sequential(
conv_block(c_num // 2,
out_nc,
3,
stride=1,
padding=1,
norm='none',
activation='tanh'))
def __init__(self,
in_nc=4,
out_nc=3,
nf=48,
norm="in",
activation="relu",
res_blocks_num=8,
dmfb_block_num=4,
device=torch.device('cuda')):
"""
:param in_nc: in channels number
:param out_nc: out channels number
:param nf: number of intermediete features
:param res_blocks_num: number of Residual blocks in fine and coarse roots
:param dmfb_block_num: nunumber of DMFB blocks in fine route
"""
super(InpaintingGenerator, self).__init__()
# [4, 256, 256]
self.encoder_coarse = nn.ModuleList([
# [48, 256, 256] -> decoder_6
conv_block(in_nc,
nf,
kernel_size=5,
stride=1,
padding=2,
norm=norm,
activation=activation),
# [96, 128, 128]
conv_block(nf,
nf * 2,
kernel_size=3,
stride=2,
padding=1,
norm=norm,
activation=activation),
# [96, 128, 128] -> decoder_4
conv_block(nf * 2,
nf * 2,
kernel_size=3,
stride=1,
padding=1,
norm=norm,
activation=activation),
# [192, 64, 64]
conv_block(nf * 2,
nf * 4,
kernel_size=3,
stride=2,
padding=1,
norm=norm,
activation=activation),
# [192, 64, 64]
conv_block(nf * 4,
nf * 4,
kernel_size=3,
stride=1,
padding=1,
norm=norm,
activation=activation),
# [192, 64, 64]
conv_block(nf * 4,
nf * 4,
kernel_size=3,
stride=1,
padding=1,
norm=norm,
activation=activation)
])
self.encoder_fine = nn.ModuleList([
# [48, 256, 256] -> decoder_6
conv_block(in_nc,
nf,
kernel_size=5,
stride=1,
padding=2,
norm=norm,
activation=activation),
# [96, 128, 128]
conv_block(nf,
nf * 2,
kernel_size=3,
stride=2,
padding=1,
norm=norm,
activation=activation),
# [96, 128, 128] -> decoder_4
conv_block(nf * 2,
nf * 2,
kernel_size=3,
stride=1,
padding=1,
norm=norm,
activation=activation),
# [192, 64, 64]
conv_block(nf * 2,
nf * 4,
kernel_size=3,
stride=2,
padding=1,
norm=norm,
activation=activation),
# [192, 64, 64]
conv_block(nf * 4,
nf * 4,
kernel_size=3,
stride=1,
padding=1,
norm=norm,
activation=activation),
# [192, 64, 64]
conv_block(nf * 4,
nf * 4,
kernel_size=3,
stride=1,
padding=1,
norm=norm,
activation=activation)
])
blocks = []
for _ in range(res_blocks_num):
block = B.ResConv(4 * nf, kernel_size=3, dilation=2,
padding=2) # [192, 64, 64]
blocks.append(block)
self.coarse = nn.Sequential(*blocks)
self.fine = FineBottleneck(4 * nf,
res_blocks_num,
dmfb_block_num,
attention_mode="self")
self.decoder_coarse = nn.ModuleList([
# [192, 64, 64]
conv_block(nf * 4,
nf * 4,
kernel_size=3,
stride=1,
padding=1,
norm=norm,
activation=activation),
# [192, 64, 64]
conv_block(nf * 4,
nf * 4,
kernel_size=3,
stride=1,
padding=1,
norm=norm,
activation=activation),
# [96, 128, 128]
upconv_block(nf * 4,
nf * 2,
kernel_size=3,
stride=2,
output_padding=1,
padding=get_pad_tp(64, 64, [1, 1], [3, 3], [2, 2],
[1, 1]),
norm=norm,
activation=activation),
# [96, 128, 128]
conv_block(nf * 4,
nf * 2,
kernel_size=3,
stride=1,
padding=1,
norm=norm,
activation=activation),
# [48, 256, 256]
upconv_block(nf * 2,
nf,
kernel_size=3,
stride=2,
output_padding=1,
padding=get_pad_tp(128, 128, [1, 1], [3, 3], [2, 2],
[1, 1]),
norm=norm,
activation=activation),
# [48, 256, 256]
conv_block(nf * 2,
nf,
kernel_size=3,
stride=1,
padding=1,
norm=norm,
activation=activation),
# [24, 256, 256]
conv_block(nf,
nf // 2,
kernel_size=3,
stride=1,
padding=1,
norm=norm,
activation=activation)
])
self.decoder_fine = nn.ModuleList([
# [192, 64, 64]
conv_block(nf * 4,
nf * 4,
kernel_size=3,
stride=1,
padding=1,
norm=norm,
activation=activation),
# [192, 64, 64]
conv_block(nf * 4,
nf * 4,
kernel_size=3,
stride=1,
padding=1,
norm=norm,
activation=activation),
# [96, 128, 128]
upconv_block(nf * 4,
nf * 2,
kernel_size=3,
stride=2,
output_padding=1,
padding=get_pad_tp(64, 64, [1, 1], [3, 3], [2, 2],
[1, 1]),
norm=norm,
activation=activation),
# [96, 128, 128]
conv_block(nf * 4,
nf * 2,
kernel_size=3,
stride=1,
padding=1,
norm=norm,
activation=activation),
# [48, 256, 256]
upconv_block(nf * 2,
nf,
kernel_size=3,
stride=2,
output_padding=1,
padding=get_pad_tp(128, 128, [1, 1], [3, 3], [2, 2],
[1, 1]),
norm=norm,
activation=activation),
# [48, 256, 256]
conv_block(nf * 2,
nf,
kernel_size=3,
stride=1,
padding=1,
norm=norm,
activation=activation),
# [24, 256, 256]
conv_block(nf,
nf // 2,
kernel_size=3,
stride=1,
padding=1,
norm=norm,
activation=activation)
])
self.out_coarse = nn.Sequential(
conv_block(nf // 2,
out_nc,
3,
stride=1,
padding=1,
norm='none',
activation='tanh'))
self.out_fine = nn.Sequential(
conv_block(nf // 2,
out_nc,
3,
stride=1,
padding=1,
norm='none',
activation='tanh'))
self.device = device
def __init__(self, in_nc=4, c_num=32, device=None):
super(FineGenerator, self).__init__()
self.conv_seq = nn.Sequential(
# [4, 256, 256]
conv_block(in_channels=in_nc,
out_channels=c_num,
kernel_size=5,
stride=1,
padding=2),
# [cnum, 256, 256]
conv_block(in_channels=c_num,
out_channels=2 * c_num,
kernel_size=3,
stride=2,
padding=1),
# [cnum * 2, 128, 128]
conv_block(in_channels=2 * c_num,
out_channels=2 * c_num,
kernel_size=3,
stride=2,
padding=1),
# [cnum * 2, 64, 64]
conv_block(in_channels=2 * c_num,
out_channels=4 * c_num,
kernel_size=3,
stride=1,
padding=1),
# [cnum * 4, 64, 64]
conv_block(in_channels=4 * c_num,
out_channels=4 * c_num,
kernel_size=3,
stride=1,
padding=1),
# [cnum * 4, 64, 64]
# dilation
conv_block(in_channels=4 * c_num,
out_channels=4 * c_num,
kernel_size=3,
stride=1,
padding=2,
dilation=2),
# [cnum * 4, 64, 64]
conv_block(in_channels=4 * c_num,
out_channels=4 * c_num,
kernel_size=3,
stride=1,
padding=4,
dilation=4),
# [cnum * 4, 64, 64]
conv_block(in_channels=4 * c_num,
out_channels=4 * c_num,
kernel_size=3,
stride=1,
padding=8,
dilation=8),
# [cnum * 4, 64, 64]
conv_block(in_channels=4 * c_num,
out_channels=4 * c_num,
kernel_size=3,
stride=1,
padding=16,
dilation=16)
# [cnum * 4, 64, 64]
)
self.before_attn = nn.Sequential(
# [4, 256, 256]
conv_block(in_channels=in_nc,
out_channels=c_num,
kernel_size=5,
stride=1,
padding=2),
# [cnum, 256, 256]
conv_block(in_channels=c_num,
out_channels=2 * c_num,
kernel_size=3,
stride=2,
padding=1),
# [cnum * 2, 128, 128]
conv_block(in_channels=2 * c_num,
out_channels=2 * c_num,
kernel_size=3,
stride=2,
padding=1),
# [cnum * 2, 64, 64]
conv_block(in_channels=2 * c_num,
out_channels=4 * c_num,
kernel_size=3,
stride=1,
padding=1),
# [cnum * 4, 64, 64]
conv_block(in_channels=4 * c_num,
out_channels=4 * c_num,
kernel_size=3,
stride=1,
padding=1),
# [cnum * 4, 64, 64]
)
# Contextual attention
self.contextual_attention = ContextualAttention(ksize=3,
stride=1,
rate=2,
fuse_k=3,
softmax_scale=10,
fuse=True,
device=device)
# [cnum * 4, 64, 64]
self.after_attn = nn.Sequential(
conv_block(in_channels=4 * c_num,
out_channels=4 * c_num,
kernel_size=3,
stride=1,
padding=1),
# [cnum * 4, 64, 64]
conv_block(in_channels=4 * c_num,
out_channels=4 * c_num,
kernel_size=3,
stride=1,
padding=1),
# [cnum * 4, 64, 64]
)
self.all_conv = nn.Sequential(
# concatenated input [cnum * 8, 64, 64]
conv_block(in_channels=8 * c_num,
out_channels=4 * c_num,
kernel_size=3,
stride=1,
padding=1),
# [cnum * 4, 64, 64]
conv_block(in_channels=4 * c_num,
out_channels=4 * c_num,
kernel_size=3,
stride=1,
padding=1),
# [cnum * 4, 64, 64]
upconv_block(in_channels=4 * c_num,
out_channels=2 * c_num,
kernel_size=3,
stride=2,
padding=1,
norm='in',
output_padding=1),
# [cnum * 2, 128, 128]
conv_block(in_channels=2 * c_num,
out_channels=2 * c_num,
kernel_size=3,
stride=1,
padding=1),
# [cnum * 2, 128, 128]
upconv_block(in_channels=2 * c_num,
out_channels=c_num,
kernel_size=3,
stride=2,
padding=1,
norm='in',
output_padding=1),
# [cnum, 256, 256]
conv_block(in_channels=c_num,
out_channels=c_num // 2,
kernel_size=3,
stride=1,
padding=1),
# [cnum * 8, 256, 256]
conv_block(in_channels=c_num // 2,
out_channels=3,
kernel_size=3,
stride=1,
padding=1,
activation='tanh'),
)