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)