def __init__(self,
num_in_ch=3,
num_out_ch=3,
num_feat=64,
num_frame=5,
deformable_groups=8,
num_extract_block=5,
num_reconstruct_block=10,
center_frame_idx=2,
hr_in=False,
with_predeblur=False,
with_tsa=True):
super(EDVR, self).__init__()
if center_frame_idx is None:
self.center_frame_idx = num_frame // 2
else:
self.center_frame_idx = center_frame_idx
self.hr_in = hr_in
self.with_predeblur = with_predeblur
self.with_tsa = with_tsa
# extract features for each frame
if self.with_predeblur:
self.predeblur = PredeblurModule(
num_feat=num_feat, hr_in=self.hr_in)
self.conv_1x1 = nn.Conv2d(num_feat, num_feat, 1, 1)
else:
self.conv_first = nn.Conv2d(num_in_ch, num_feat, 3, 1, 1)
# extrat pyramid features
self.feature_extraction = make_layer(
ResidualBlockNoBN, num_extract_block, num_feat=num_feat)
self.conv_l2_1 = nn.Conv2d(num_feat, num_feat, 3, 2, 1)
self.conv_l2_2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
self.conv_l3_1 = nn.Conv2d(num_feat, num_feat, 3, 2, 1)
self.conv_l3_2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
# pcd and tsa module
self.pcd_align = PCDAlignment(
num_feat=num_feat, deformable_groups=deformable_groups)
if self.with_tsa:
self.fusion = TSAFusion(
num_feat=num_feat,
num_frame=num_frame,
center_frame_idx=self.center_frame_idx)
else:
self.fusion = nn.Conv2d(num_frame * num_feat, num_feat, 1, 1)
# reconstruction
self.reconstruction = make_layer(
ResidualBlockNoBN, num_reconstruct_block, num_feat=num_feat)
# upsample
self.upconv1 = nn.Conv2d(num_feat, num_feat * 4, 3, 1, 1)
self.upconv2 = nn.Conv2d(num_feat, 64 * 4, 3, 1, 1)
self.pixel_shuffle = nn.PixelShuffle(2)
self.conv_hr = nn.Conv2d(64, 64, 3, 1, 1)
self.conv_last = nn.Conv2d(64, 3, 3, 1, 1)
# activation function
self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=True)
def __init__(self,
num_in_ch,
num_out_ch,
num_feat=64,
num_group=10,
num_block=16,
squeeze_factor=16,
upscale=4,
res_scale=1,
img_range=255.,
rgb_mean=(0.4488, 0.4371, 0.4040)):
super(RCAN, self).__init__()
self.img_range = img_range
self.mean = torch.Tensor(rgb_mean).view(1, 3, 1, 1)
self.conv_first = nn.Conv2d(num_in_ch, num_feat, 3, 1, 1)
self.body = make_layer(ResidualGroup,
num_group,
num_feat=num_feat,
num_block=num_block,
squeeze_factor=squeeze_factor,
res_scale=res_scale)
self.conv_after_body = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
self.upsample = Upsample(upscale, num_feat)
self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)
def __init__(self,
num_in_ch=3,
num_out_ch=3,
num_feat=64,
num_block=16,
upscale=4):
super(MSRResNet, self).__init__()
self.upscale = upscale
self.conv_first = nn.Conv2d(num_in_ch, num_feat, 3, 1, 1)
self.body = arch_util.make_layer(
arch_util.ResidualBlockNoBN, num_block, num_feat=num_feat)
# upsampling
if self.upscale in [2, 3]:
self.upconv1 = nn.Conv2d(num_feat,
num_feat * self.upscale * self.upscale, 3,
1, 1)
self.pixel_shuffle = nn.PixelShuffle(self.upscale)
elif self.upscale == 4:
self.upconv1 = nn.Conv2d(num_feat, num_feat * 4, 3, 1, 1)
self.upconv2 = nn.Conv2d(num_feat, num_feat * 4, 3, 1, 1)
self.pixel_shuffle = nn.PixelShuffle(2)
self.conv_hr = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)
# activation function
self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=True)
# initialization
arch_util.default_init_weights(
[self.conv_first, self.upconv1, self.conv_hr, self.conv_last], 0.1)
if self.upscale == 4:
arch_util.default_init_weights(self.upconv2, 0.1)
def __init__(self, num_feat, num_block, squeeze_factor=16, res_scale=1):
super(ResidualGroup, self).__init__()
self.residual_group = make_layer(RCAB,
num_block,
num_feat=num_feat,
squeeze_factor=squeeze_factor,
res_scale=res_scale)
self.conv = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
def __init__(self, in_feat, out_feat=64, num_block=30):
super(ConvResBlock, self).__init__()
conv_resblock = []
conv_resblock.append(
nn.Conv2d(in_feat,
out_feat,
kernel_size=3,
stride=1,
padding=1,
bias=True))
conv_resblock.append(nn.LeakyReLU(negative_slope=0.1, inplace=True))
conv_resblock.append(
make_layer(ResidualBlockNoBN, num_block, num_feat=out_feat))
self.conv_resblock = nn.Sequential(*conv_resblock)
def __init__(self,
num_in_ch=3,
num_out_ch=3,
num_feat=64,
num_block=16,
upscale=4):
super(MSRResNet_WT_Pixel, self).__init__()
self.upscale = upscale
self.conv_first = nn.Conv2d(num_in_ch, num_feat, 3, 1, 1)
self.body = arch_util.make_layer(
arch_util.ResidualBlockNoBN, num_block, num_feat=num_feat)
# upsampling
if self.upscale in [2, 3]:
self.upconv1 = nn.Conv2d(num_feat,
num_feat * self.upscale * self.upscale, 3,
1, 1)
self.pixel_shuffle = nn.PixelShuffle(self.upscale)
elif self.upscale == 4:
self.upconv1 = nn.Conv2d(num_feat, num_feat * 4, 3, 1, 1)
self.upconv2 = nn.Conv2d(num_feat, num_feat * 4, 3, 1, 1)
self.pixel_shuffle = nn.PixelShuffle(2)
self.conv_hr = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)
# activation function
self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=True)
# initialization
arch_util.default_init_weights(
[self.conv_first, self.upconv1, self.conv_hr, self.conv_last], 0.1)
if self.upscale == 4:
arch_util.default_init_weights(self.upconv2, 0.1)
# WT filter
inv_filters = arch_util.create_inv_filters()
self.register_buffer('inv_filters', inv_filters)
# Normalization buffers
self.register_buffer(
'shift',
torch.Tensor([3.0]))
self.register_buffer(
'scale',
torch.Tensor([10.0]))
def __init__(self,
num_in_ch,
num_out_ch,
num_feat=80,
num_block=23,
num_grow_ch=40):
super(RRDBNet_x10, self).__init__()
self.conv_first = nn.Conv2d(num_in_ch, num_feat, 3, 1, 1)
self.body = make_layer(
RRDB_x10, num_block, num_feat=num_feat, num_grow_ch=num_grow_ch)
self.conv_body = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
# upsample
self.conv_up1 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
self.conv_up2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
self.conv_hr = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)
self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
def __init__(
self,
num_in_ch=3,
num_out_ch=3,
num_feat=64,
num_block=16,
):
super().__init__()
self.conv_first = nn.Conv2d(num_in_ch, num_feat, 3, 1, 1)
self.body = arch_util.make_layer(arch_util.ResidualBlockNoBN,
num_block,
num_feat=num_feat)
self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)
# activation function
self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=True)
# initialization
arch_util.default_init_weights([self.conv_first, self.conv_last], 0.1)
def __init__(self,
in_channels,
mid_channels,
out_channels,
num_block=4,
img_range=255.,
rgb_mean=(0.4488, 0.4371, 0.4040),
rgb_std=(1.0, 1.0, 1.0)):
super(RIDNet, self).__init__()
self.sub_mean = MeanShift(img_range, rgb_mean, rgb_std)
self.add_mean = MeanShift(img_range, rgb_mean, rgb_std, 1)
self.head = nn.Conv2d(in_channels, mid_channels, 3, 1, 1)
self.body = make_layer(EAM,
num_block,
in_channels=mid_channels,
mid_channels=mid_channels,
out_channels=mid_channels)
self.tail = nn.Conv2d(mid_channels, out_channels, 3, 1, 1)
self.relu = nn.ReLU(inplace=True)
def __init__(self,
num_in_ch,
num_out_ch,
num_feat=64,
num_block=16,
upscale=4,
res_scale=1,
img_range=255.,
rgb_mean=(0.4488, 0.4371, 0.4040)):
super(EDSR, self).__init__()
self.img_range = img_range
self.mean = torch.Tensor(rgb_mean).view(1, 3, 1, 1)
self.conv_first = nn.Conv2d(num_in_ch, num_feat, 3, 1, 1)
self.body = make_layer(ResidualBlockNoBN,
num_block,
num_feat=num_feat,
res_scale=res_scale,
pytorch_init=True)
self.conv_after_body = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
self.upsample = Upsample(upscale, num_feat)
self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)
def __init__(self,
num_in_ch=3,
num_out_ch=3,
num_feat=64,
num_frame=5,
deformable_groups=8,
num_extract_block=5,
num_reconstruct_block=10,
center_frame_idx=2,
hr_in=False,
with_predeblur=False,
with_tsa=True,
dropout=0.0):
super(EDVR, self).__init__()
self.hidden_dim = hdim = 128
self.context_dim = cdim = 128
if center_frame_idx is None:
self.center_frame_idx = num_frame // 2
else:
self.center_frame_idx = center_frame_idx
self.hr_in = hr_in
self.dropout = dropout
# RAFT Parameter
self.corr_radius = 4
self.num_levels = 4
self.mixed_precision = False
self.iters = 6
# extrat v features
self.feature_extractor_v = make_layer(ResidualBlockNoBN,
num_extract_block,
num_feat=num_feat)
# RAFT Network
# extrat qk/c features
self.feature_extractor_qk = BasicEncoder(output_dim=256,
norm_fn='instance',
dropout=self.dropout)
self.feature_extractor_c = BasicEncoder(output_dim=hdim + cdim,
norm_fn='batch',
dropout=self.dropout)
# update flow
self.update_block = BasicUpdateBlock(self.num_levels,
self.corr_radius,
hidden_dim=hdim)
# pcd and tsa module
self.fusion = nn.Conv2d(num_frame * num_feat, num_feat, 1, 1)
# reconstruction
self.reconstruction = make_layer(ResidualBlockNoBN,
num_reconstruct_block,
num_feat=num_feat)
# upsample
self.upconv1 = nn.Conv2d(num_feat, num_feat * 4, 3, 1, 1)
self.upconv2 = nn.Conv2d(num_feat, 64 * 4, 3, 1, 1)
self.upconv3 = nn.Conv2d(num_feat, 64 * 4, 3, 1, 1)
self.pixel_shuffle = nn.PixelShuffle(2)
self.conv_hr = nn.Conv2d(64, 64, 3, 1, 1)
self.conv_last = nn.Conv2d(64, 3, 3, 1, 1)
# activation function
self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=True)
# kaiming_normal_
for m in self.modules():
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight,
mode='fan_out',
nonlinearity='relu')
elif isinstance(m,
(nn.BatchNorm2d, nn.InstanceNorm2d, nn.GroupNorm)):
if m.weight is not None:
nn.init.constant_(m.weight, 1)
if m.bias is not None:
nn.init.constant_(m.bias, 0)
def __init__(self,
num_in_ch=3,
num_out_ch=3,
num_feat=64,
num_frame=5,
deformable_groups=8,
num_extract_block=5,
center_frame_idx=None,
hr_in=None,
with_predeblur=False,
with_tsa=True):
super(EDVRExtractor, self).__init__()
if center_frame_idx is None:
self.center_frame_idx = num_frame // 2
else:
self.center_frame_idx = center_frame_idx
self.hr_in = hr_in
self.with_predeblur = with_predeblur
self.with_tsa = with_tsa
# extract features for each frame
if self.with_predeblur:
self.pre_deblur = PredeblurModule(num_feat=num_feat,
hr_in=self.hr_in)
self.conv_1x1 = nn.Conv2d(num_feat,
num_feat,
kernel_size=1,
stride=1,
padding=0,
bias=True)
else:
self.conv_first = nn.Conv2d(num_in_ch,
num_feat,
kernel_size=3,
stride=1,
padding=1)
# extract pyramid features
self.feature_extraction = make_layer(ResidualBlockNoBN,
num_extract_block,
num_feat=num_feat)
self.conv_l2_1 = nn.Conv2d(num_feat,
num_feat,
kernel_size=3,
stride=2,
padding=1)
self.conv_l2_2 = nn.Conv2d(num_feat,
num_feat,
kernel_size=3,
stride=1,
padding=1)
self.conv_l3_1 = nn.Conv2d(num_feat,
num_feat,
kernel_size=3,
stride=2,
padding=1)
self.conv_l3_2 = nn.Conv2d(num_feat,
num_feat,
kernel_size=3,
stride=1,
padding=1)
# pcd and tsa module
self.pcd_align = PCDAlignment(num_feat=num_feat,
deformable_groups=deformable_groups)
if self.with_tsa:
self.fusion = TSAFusion(num_feat=num_feat,
num_frame=num_frame,
center_frame_idx=self.center_frame_idx)
else:
self.fusion = nn.Conv2d(num_frame * num_feat, num_feat, 1, 1)
# activation function
self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=True)
def __init__(self,
num_in_ch=3,
num_out_ch=3,
num_feat=64,
num_frame=5,
deformable_groups=8,
num_extract_block=5,
num_reconstruct_block=10,
center_frame_idx=2,
hr_in=False,
with_predeblur=False,
with_tsa=True):
super(EDVR, self).__init__()
if center_frame_idx is None:
self.center_frame_idx = num_frame // 2
else:
self.center_frame_idx = center_frame_idx
self.hr_in = hr_in
self.with_predeblur = with_predeblur
self.with_tsa = with_tsa
# extract features for each frame
if self.with_predeblur:
self.predeblur = PredeblurModule(num_feat=num_feat,
hr_in=self.hr_in)
self.conv_1x1 = nn.Conv2d(num_feat, num_feat, 1, 1)
else:
self.conv_first = nn.Conv2d(num_in_ch, num_feat, 3, 1, 1)
# extrat pyramid features
self.feature_extraction = make_layer(ResidualBlockNoBN,
num_extract_block,
num_feat=num_feat)
self.conv_l2_1 = nn.Conv2d(num_feat, num_feat, 3, 2, 1)
self.conv_l2_2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
self.conv_l3_1 = nn.Conv2d(num_feat, num_feat, 3, 2, 1)
self.conv_l3_2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
# pcd and tsa module
# self.pcd_align = PCDAlignment(
# num_feat=num_feat, deformable_groups=deformable_groups)
self.pcd_align_3d = PCDAlignment(num_feat=num_feat,
deformable_groups=deformable_groups)
if self.with_tsa:
self.fusion = TSAFusion(num_feat=num_feat,
num_frame=num_frame,
center_frame_idx=self.center_frame_idx)
else:
self.fusion = nn.Conv2d(num_frame * num_feat, num_feat, 1, 1)
# reconstruction
self.reconstruction = make_layer(ResidualBlockNoBN,
num_reconstruct_block,
num_feat=num_feat)
# upsample
self.upconv1 = nn.Conv2d(num_feat, num_feat * 4, 3, 1, 1)
self.upconv2 = nn.Conv2d(num_feat, num_feat * 4, 3, 1, 1)
self.pixel_shuffle = nn.PixelShuffle(2)
self.conv_hr = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
self.conv_last = nn.Conv2d(num_feat, 3, 3, 1, 1)
# self.attn1 = Self_Attn( 128, 'relu')
# self.attn2 = Self_Attn( 64, 'relu')
# activation function
self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=True)
# init parameter
for m in self.modules():
print('m:', m)
# assert(0)
if isinstance(m, (nn.Conv2d, nn.Linear)):
# print(m.weight)
nn.init.xavier_uniform_(m.weight,
gain=nn.init.calculate_gain('relu'))
