def __init__(self,
in_channels,
out_channels,
mid_channels=64,
num_frames=5,
deform_groups=8,
num_blocks_extraction=5,
num_blocks_reconstruction=10,
center_frame_idx=2,
with_tsa=True):
super(EDVRNet, self).__init__()
self.center_frame_idx = center_frame_idx
self.with_tsa = with_tsa
act_cfg = dict(type='LeakyReLU', negative_slope=0.1)
self.conv_first = nn.Conv2d(in_channels, mid_channels, 3, 1, 1)
self.feature_extraction = make_layer(
ResidualBlockNoBN,
num_blocks_extraction,
mid_channels=mid_channels)
# generate pyramid features
self.feat_l2_conv1 = ConvModule(
mid_channels, mid_channels, 3, 2, 1, act_cfg=act_cfg)
self.feat_l2_conv2 = ConvModule(
mid_channels, mid_channels, 3, 1, 1, act_cfg=act_cfg)
self.feat_l3_conv1 = ConvModule(
mid_channels, mid_channels, 3, 2, 1, act_cfg=act_cfg)
self.feat_l3_conv2 = ConvModule(
mid_channels, mid_channels, 3, 1, 1, act_cfg=act_cfg)
# pcd alignment
self.pcd_alignment = PCDAlignment(
mid_channels=mid_channels, deform_groups=deform_groups)
# fusion
if self.with_tsa:
self.fusion = TSAFusion(
mid_channels=mid_channels,
num_frames=num_frames,
center_frame_idx=self.center_frame_idx)
else:
self.fusion = nn.Conv2d(num_frames * mid_channels, mid_channels, 1,
1)
# reconstruction
self.reconstruction = make_layer(
ResidualBlockNoBN,
num_blocks_reconstruction,
mid_channels=mid_channels)
# upsample
self.upsample1 = PixelShufflePack(
mid_channels, mid_channels, 2, upsample_kernel=3)
self.upsample2 = PixelShufflePack(
mid_channels, 64, 2, upsample_kernel=3)
# we fix the output channels in the last few layers to 64.
self.conv_hr = nn.Conv2d(64, 64, 3, 1, 1)
self.conv_last = nn.Conv2d(64, 3, 3, 1, 1)
self.img_upsample = nn.Upsample(
scale_factor=4, mode='bilinear', align_corners=False)
# activation function
self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=True)
def __init__(self,
in_channels,
out_channels,
mid_channels=64,
num_blocks=16,
upscale_factor=4,
res_scale=1,
rgb_mean=(0.4488, 0.4371, 0.4040),
rgb_std=(1.0, 1.0, 1.0)):
super(EDSR, self).__init__()
self.in_channels = in_channels
self.out_channels = out_channels
self.mid_channels = mid_channels
self.num_blocks = num_blocks
self.upscale_factor = upscale_factor
self.mean = torch.Tensor(rgb_mean).view(1, 3, 1, 1)
self.std = torch.Tensor(rgb_std).view(1, 3, 1, 1)
self.conv_first = nn.Conv2d(in_channels, mid_channels, 3, padding=1)
self.body = make_layer(ResidualBlockNoBN,
num_blocks,
mid_channels=mid_channels,
res_scale=res_scale)
self.conv_after_body = nn.Conv2d(mid_channels, mid_channels, 3, 1, 1)
self.upsample = UpsampleModule(upscale_factor, mid_channels)
self.conv_last = nn.Conv2d(mid_channels,
out_channels,
3,
1,
1,
bias=True)
def __init__(self,
in_channels=3,
kernel_size=3,
num_block_groups=5,
num_block_layers=12,
depth=3,
reduction=16,
norm=None,
padding=7,
act=nn.LeakyReLU(0.2, True)):
super().__init__()
mid_channels = in_channels * (4**depth)
self.scale = 2**depth
self.padding = padding
self.conv_first = nn.Conv2d(mid_channels * 2, mid_channels,
kernel_size, 1, 1)
self.body = make_layer(ResidualGroup,
num_block_groups,
block_layer=ResidualChannelAttention,
num_block_layers=num_block_layers,
mid_channels=mid_channels,
kernel_size=kernel_size,
reduction=reduction,
norm=norm,
act=act)
self.conv_last = nn.Conv2d(mid_channels, mid_channels, kernel_size, 1,
1)
def __init__(self,
in_channels,
out_channels,
mid_channels=64,
num_blocks=23,
growth_channels=32,
upscale_factor=4):
super().__init__()
if upscale_factor in self._supported_upscale_factors:
in_channels = in_channels * ((4 // upscale_factor)**2)
else:
raise ValueError(f'Unsupported scale factor {upscale_factor}. '
f'Currently supported ones are '
f'{self._supported_upscale_factors}.')
self.upscale_factor = upscale_factor
self.conv_first = nn.Conv2d(in_channels, mid_channels, 3, 1, 1)
self.body = make_layer(RRDB,
num_blocks,
mid_channels=mid_channels,
growth_channels=growth_channels)
self.conv_body = nn.Conv2d(mid_channels, mid_channels, 3, 1, 1)
# upsample
self.conv_up1 = nn.Conv2d(mid_channels, mid_channels, 3, 1, 1)
self.conv_up2 = nn.Conv2d(mid_channels, mid_channels, 3, 1, 1)
self.conv_hr = nn.Conv2d(mid_channels, mid_channels, 3, 1, 1)
self.conv_last = nn.Conv2d(mid_channels, out_channels, 3, 1, 1)
self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
def __init__(self,
in_channels=3,
mid_channels=64,
out_channels=3,
num_blocks_before_align=5,
num_blocks_after_align=10):
super().__init__()
self.feat_extract = nn.Sequential(
ConvModule(in_channels, mid_channels, 3, padding=1),
make_layer(ResidualBlockNoBN,
num_blocks_before_align,
mid_channels=mid_channels))
self.feat_aggregate = nn.Sequential(
nn.Conv2d(mid_channels * 2, mid_channels, 3, padding=1, bias=True),
DeformConv2dPack(mid_channels,
mid_channels,
3,
padding=1,
deform_groups=8),
DeformConv2dPack(mid_channels,
mid_channels,
3,
padding=1,
deform_groups=8))
self.align_1 = AugmentedDeformConv2dPack(mid_channels,
mid_channels,
3,
padding=1,
deform_groups=8)
self.align_2 = DeformConv2dPack(mid_channels,
mid_channels,
3,
padding=1,
deform_groups=8)
self.to_rgb = nn.Conv2d(mid_channels, 3, 3, padding=1, bias=True)
self.reconstruct = nn.Sequential(
ConvModule(in_channels * 5, mid_channels, 3, padding=1),
make_layer(ResidualBlockNoBN,
num_blocks_after_align,
mid_channels=mid_channels),
PixelShufflePack(mid_channels, mid_channels, 2, upsample_kernel=3),
PixelShufflePack(mid_channels, mid_channels, 2, upsample_kernel=3),
nn.Conv2d(mid_channels, out_channels, 3, 1, 1, bias=False))
def __init__(self, in_channels, mid_channels, num_blocks, res_scale):
super().__init__()
self.num_blocks = num_blocks
self.conv_first = _conv3x3_layer(in_channels, mid_channels)
self.body = make_layer(ResidualBlockNoBN,
num_blocks,
mid_channels=mid_channels,
res_scale=res_scale)
self.conv_last = _conv3x3_layer(mid_channels, mid_channels)
def __init__(self,
in_channels=3,
mid_channels=64,
num_blocks=23,
growth_channels=32):
super().__init__()
self.conv_first = nn.Conv2d(in_channels, mid_channels, 3, 1, 1)
self.body = make_layer(RRDB,
num_blocks,
mid_channels=mid_channels,
growth_channels=growth_channels)
self.conv_body = nn.Conv2d(mid_channels, mid_channels, 3, 1, 1)
def __init__(self, in_channels, out_channels=64, num_blocks=30):
super().__init__()
main = []
# a convolution used to match the channels of the residual blocks
main.append(nn.Conv2d(in_channels, out_channels, 3, 1, 1, bias=True))
main.append(nn.LeakyReLU(negative_slope=0.1, inplace=True))
# residual blocks
main.append(
make_layer(
ResidualBlockNoBN, num_blocks, mid_channels=out_channels))
self.main = nn.Sequential(*main)
def __init__(self,
in_channels,
out_channels,
mid_channels=64,
num_blocks=23,
growth_channels=32):
super().__init__()
self.conv_first = nn.Conv2d(in_channels, mid_channels, 3, 1, 1)
self.body = make_layer(RRDB,
num_blocks,
mid_channels=mid_channels,
growth_channels=growth_channels)
self.conv_body = nn.Conv2d(mid_channels, mid_channels, 3, 1, 1)
# upsample
self.conv_up1 = nn.Conv2d(mid_channels, mid_channels, 3, 1, 1)
self.conv_up2 = nn.Conv2d(mid_channels, mid_channels, 3, 1, 1)
self.conv_hr = nn.Conv2d(mid_channels, mid_channels, 3, 1, 1)
self.conv_last = nn.Conv2d(mid_channels, out_channels, 3, 1, 1)
self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
def __init__(self,
block_layer,
num_block_layers,
mid_channels,
kernel_size,
reduction,
act=nn.LeakyReLU(0.2, True),
norm=None):
super().__init__()
self.body = make_layer(block_layer,
num_block_layers,
mid_channels=mid_channels,
kernel_size=kernel_size,
reduction=reduction,
norm=norm,
act=act)
self.conv_after_body = ConvNormWithReflectionPad(mid_channels,
mid_channels,
kernel_size,
norm=norm)
def __init__(self,
in_channels,
num_heatmaps,
num_blocks,
mid_channels=None):
super().__init__()
self.num_heatmaps = num_heatmaps
res_block_channel = in_channels * num_heatmaps
if mid_channels is None:
self.mid_channels = num_heatmaps * in_channels
else:
self.mid_channels = mid_channels
self.conv_first = nn.Sequential(
nn.Conv2d(in_channels, res_block_channel, kernel_size=1),
nn.LeakyReLU(negative_slope=0.2, inplace=True))
self.body = make_layer(GroupResBlock,
num_blocks,
in_channels=res_block_channel,
out_channels=res_block_channel,
mid_channels=self.mid_channels,
groups=num_heatmaps)
def __init__(self,
in_channels,
out_channels,
mid_channels=64,
num_blocks=16,
upscale_factor=4):
super().__init__()
self.in_channels = in_channels
self.out_channels = out_channels
self.mid_channels = mid_channels
self.num_blocks = num_blocks
self.upscale_factor = upscale_factor
self.conv_first = nn.Conv2d(in_channels,
mid_channels,
3,
1,
1,
bias=True)
self.trunk_net = make_layer(ResidualBlockNoBN,
num_blocks,
mid_channels=mid_channels)
# upsampling
if self.upscale_factor in [2, 3]:
self.upsample1 = PixelShufflePack(mid_channels,
mid_channels,
self.upscale_factor,
upsample_kernel=3)
elif self.upscale_factor == 4:
self.upsample1 = PixelShufflePack(mid_channels,
mid_channels,
2,
upsample_kernel=3)
self.upsample2 = PixelShufflePack(mid_channels,
mid_channels,
2,
upsample_kernel=3)
else:
raise ValueError(
f'Unsupported scale factor {self.upscale_factor}. '
f'Currently supported ones are '
f'{self._supported_upscale_factors}.')
self.conv_hr = nn.Conv2d(mid_channels,
mid_channels,
3,
1,
1,
bias=True)
self.conv_last = nn.Conv2d(mid_channels,
out_channels,
3,
1,
1,
bias=True)
self.img_upsampler = nn.Upsample(scale_factor=self.upscale_factor,
mode='bilinear',
align_corners=False)
# activation function
self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=True)
def __init__(self,
in_channels,
out_channels,
mid_channels=64,
texture_channels=64,
num_blocks=(16, 16, 8, 4),
res_scale=1.0):
super().__init__()
self.texture_channels = texture_channels
self.sfe = SFE(in_channels, mid_channels, num_blocks[0], res_scale)
# stage 1
self.conv_first1 = _conv3x3_layer(4 * texture_channels + mid_channels,
mid_channels)
self.res_block1 = make_layer(ResidualBlockNoBN,
num_blocks[1],
mid_channels=mid_channels,
res_scale=res_scale)
self.conv_last1 = _conv3x3_layer(mid_channels, mid_channels)
# up-sampling 1 -> 2
self.up1 = PixelShufflePack(in_channels=mid_channels,
out_channels=mid_channels,
scale_factor=2,
upsample_kernel=3)
# stage 2
self.conv_first2 = _conv3x3_layer(2 * texture_channels + mid_channels,
mid_channels)
self.csfi2 = CSFI2(mid_channels)
self.res_block2_1 = make_layer(ResidualBlockNoBN,
num_blocks[2],
mid_channels=mid_channels,
res_scale=res_scale)
self.res_block2_2 = make_layer(ResidualBlockNoBN,
num_blocks[2],
mid_channels=mid_channels,
res_scale=res_scale)
self.conv_last2_1 = _conv3x3_layer(mid_channels, mid_channels)
self.conv_last2_2 = _conv3x3_layer(mid_channels, mid_channels)
# up-sampling 2 -> 3
self.up2 = PixelShufflePack(in_channels=mid_channels,
out_channels=mid_channels,
scale_factor=2,
upsample_kernel=3)
# stage 3
self.conv_first3 = _conv3x3_layer(texture_channels + mid_channels,
mid_channels)
self.csfi3 = CSFI3(mid_channels)
self.res_block3_1 = make_layer(ResidualBlockNoBN,
num_blocks[3],
mid_channels=mid_channels,
res_scale=res_scale)
self.res_block3_2 = make_layer(ResidualBlockNoBN,
num_blocks[3],
mid_channels=mid_channels,
res_scale=res_scale)
self.res_block3_3 = make_layer(ResidualBlockNoBN,
num_blocks[3],
mid_channels=mid_channels,
res_scale=res_scale)
self.conv_last3_1 = _conv3x3_layer(mid_channels, mid_channels)
self.conv_last3_2 = _conv3x3_layer(mid_channels, mid_channels)
self.conv_last3_3 = _conv3x3_layer(mid_channels, mid_channels)
# end, merge features
self.merge_features = MergeFeatures(mid_channels, out_channels)
def __init__(self,
in_channels=3,
out_channel=3,
mid_channels=64,
num_frames=5,
deform_groups=8,
num_blocks_extraction=5,
num_blocks_reconstruction=10,
center_frame_idx=2,
with_tsa=True,
pretrained=None):
super().__init__()
self.center_frame_idx = center_frame_idx
self.with_tsa = with_tsa
act_cfg = dict(type='LeakyReLU', negative_slope=0.1)
self.conv_first = nn.Conv2d(in_channels, mid_channels, 3, 1, 1)
self.feature_extraction = make_layer(ResidualBlockNoBN,
num_blocks_extraction,
mid_channels=mid_channels)
# generate pyramid features
self.feat_l2_conv1 = ConvModule(mid_channels,
mid_channels,
3,
2,
1,
act_cfg=act_cfg)
self.feat_l2_conv2 = ConvModule(mid_channels,
mid_channels,
3,
1,
1,
act_cfg=act_cfg)
self.feat_l3_conv1 = ConvModule(mid_channels,
mid_channels,
3,
2,
1,
act_cfg=act_cfg)
self.feat_l3_conv2 = ConvModule(mid_channels,
mid_channels,
3,
1,
1,
act_cfg=act_cfg)
# pcd alignment
self.pcd_alignment = PCDAlignment(mid_channels=mid_channels,
deform_groups=deform_groups)
# fusion
if self.with_tsa:
self.fusion = TSAFusion(mid_channels=mid_channels,
num_frames=num_frames,
center_frame_idx=self.center_frame_idx)
else:
self.fusion = nn.Conv2d(num_frames * mid_channels, mid_channels, 1,
1)
# activation function
self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=True)
if isinstance(pretrained, str):
logger = get_root_logger()
load_checkpoint(self, pretrained, strict=True, logger=logger)
elif pretrained is not None:
raise TypeError(f'"pretrained" must be a str or None. '
f'But received {type(pretrained)}.')
def __init__(self,
in_channels=3,
out_channels=3,
mid_channels=64,
num_frames=5,
deform_groups=8,
num_blocks_extraction=5,
num_blocks_reconstruction=10,
center_frame_idx=2,
hr_in=False,
with_predeblur=False,
with_tsa=True,
with_nonlocal=False,
with_car=False):
super(EDVRNet_WoPre, self).__init__()
self.center_frame_idx = center_frame_idx
self.hr_in = hr_in
self.with_predeblur = with_predeblur
self.with_tsa = with_tsa
self.with_nonlocal = with_nonlocal
self.with_car = with_car
act_cfg = dict(type='LeakyReLU', negative_slope=0.1)
# extract features for each frame
if self.with_predeblur:
self.predeblur = PredeblurModule(
mid_channels=mid_channels, hr_in=self.hr_in)
self.conv_1x1 = nn.Conv2d(mid_channels, mid_channels, 1, 1)
else:
self.conv_first = nn.Conv2d(in_channels, mid_channels, 3, 1, 1)
# extract pyramid features
if self.with_car:
self.feature_extraction = make_layer(
CAResidualBlockNoBN,
num_blocks_extraction,
mid_channels=mid_channels)
else:
self.feature_extraction = make_layer(
ResidualBlockNoBN,
num_blocks_extraction,
mid_channels=mid_channels)
self.feat_l2_conv1 = ConvModule(
mid_channels, mid_channels, 3, 2, 1, act_cfg=act_cfg)
self.feat_l2_conv2 = ConvModule(
mid_channels, mid_channels, 3, 1, 1, act_cfg=act_cfg)
self.feat_l3_conv1 = ConvModule(
mid_channels, mid_channels, 3, 2, 1, act_cfg=act_cfg)
self.feat_l3_conv2 = ConvModule(
mid_channels, mid_channels, 3, 1, 1, act_cfg=act_cfg)
# pcd alignment
self.pcd_alignment = PCDAlignment(
mid_channels=mid_channels, deform_groups=deform_groups)
# tsa fusion
if self.with_tsa:
self.fusion = TSAFusion(
mid_channels=mid_channels,
num_frames=num_frames,
center_frame_idx=self.center_frame_idx)
else:
self.fusion = nn.Conv2d(num_frames * mid_channels, mid_channels, 1,
1)
# non local module
if self.with_nonlocal:
self.non_local = NonLocalModule(
mid_channels=mid_channels,
num_frames=num_frames,
center_frame_idx=self.center_frame_idx)
# reconstruction
if self.with_car:
self.reconstruction = make_layer(
CAResidualBlockNoBN,
num_blocks_reconstruction,
mid_channels=mid_channels)
else:
self.reconstruction = make_layer(
ResidualBlockNoBN,
num_blocks_reconstruction,
mid_channels=mid_channels)
# upsample
self.upsample1 = PixelShufflePack(
mid_channels, mid_channels, 2, upsample_kernel=3)
self.upsample2 = PixelShufflePack(
mid_channels, mid_channels, 2, upsample_kernel=3)
# we fix the output channels in the last few layers to 64.
self.conv_hr = nn.Conv2d(mid_channels, 64, 3, 1, 1)
self.conv_last = nn.Conv2d(64, 3, 3, 1, 1)
self.img_upsample = nn.Upsample(
scale_factor=4, mode='bilinear', align_corners=False)
# activation function
self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=True)
