def __init__(self, channel_num): super(CARBBlock, self).__init__() self.conv1 = M.Sequential( M.Conv2d(channel_num, channel_num, kernel_size=3, padding=1, stride=1), M.LeakyReLU(), M.Conv2d(channel_num, channel_num, kernel_size=3, padding=1, stride=1), ) # self.global_average_pooling = nn.AdaptiveAvgPool2d((1,1)) # B,C,H,W -> B,C,1,1 self.linear = M.Sequential(M.Linear(channel_num, channel_num // 2), M.LeakyReLU(), M.Linear(channel_num // 2, channel_num), M.Sigmoid()) self.conv2 = M.Conv2d(channel_num * 2, channel_num, kernel_size=1, padding=0, stride=1) self.lrelu = M.LeakyReLU()
def __init__(self, ch=128, nframes=7, input_nc=3, output_nc=3, upscale_factor=4, use_cost_volume=False): super(MUCAN, self).__init__() self.nframes = nframes self.upscale_factor = upscale_factor # 每个LR搞三个尺度 self.feature_encoder_carb = M.Sequential( M.Conv2d(input_nc, ch, kernel_size=3, stride=1, padding=1), M.LeakyReLU(negative_slope=0.05), CARBBlocks(channel_num=ch, block_num=4)) self.fea_L1_conv1 = M.Conv2d(ch, ch, 3, 2, 1) self.fea_L1_conv2 = M.Conv2d(ch, ch, 3, 1, 1) self.fea_L2_conv1 = M.Conv2d(ch, ch, 3, 2, 1) self.fea_L2_conv2 = M.Conv2d(ch, ch, 3, 1, 1) self.lrelu = M.LeakyReLU(negative_slope=0.05) if use_cost_volume: self.AU0 = AU_CV(K=6, d=5, ch=ch) self.AU1 = AU_CV(K=5, d=3, ch=ch) self.AU2 = AU_CV(K=4, d=3, ch=ch) else: self.AU0 = AU(ch=ch) self.AU1 = AU(ch=ch) self.AU2 = AU(ch=ch) self.UP0 = M.ConvTranspose2d(ch, ch, kernel_size=4, stride=2, padding=1) self.UP1 = M.ConvTranspose2d(ch, ch, kernel_size=4, stride=2, padding=1) self.aggre = M.Conv2d(ch * self.nframes, ch, kernel_size=3, stride=1, padding=1) self.main_conv = M.Sequential( CARBBlocks(channel_num=ch, block_num=10), M.ConvTranspose2d(ch, ch // 2, kernel_size=4, stride=2, padding=1), M.LeakyReLU(), M.Conv2d(ch // 2, ch // 2, kernel_size=3, stride=1, padding=1), M.LeakyReLU(), M.ConvTranspose2d(ch // 2, ch // 4, kernel_size=4, stride=2, padding=1), M.LeakyReLU(), M.Conv2d(ch // 4, output_nc, kernel_size=3, stride=1, padding=1))
def __init__(self, ch): super(AU, self).__init__() self.conv = M.Sequential( M.Conv2d(2*ch, ch, kernel_size=3, stride=1, padding=1), M.LeakyReLU(), M.Conv2d(ch, ch, kernel_size=3, stride=1, padding=1), M.LeakyReLU(), )
def conv(in_planes, out_planes, kernel_size=3, stride=1, dilation=1, isReLU=True, if_IN=False, IN_affine=False, if_BN=False): if isReLU: if if_IN: return nn.Sequential( nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, dilation=dilation, padding=((kernel_size - 1) * dilation) // 2, bias=True), nn.LeakyReLU(0.1), nn.InstanceNorm(out_planes, affine=IN_affine)) elif if_BN: return nn.Sequential( nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, dilation=dilation, padding=((kernel_size - 1) * dilation) // 2, bias=True), nn.LeakyReLU(0.1), nn.BatchNorm2d(out_planes, affine=IN_affine)) else: return nn.Sequential( nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, dilation=dilation, padding=((kernel_size - 1) * dilation) // 2, bias=True), nn.LeakyReLU(0.1)) else: if if_IN: return nn.Sequential( nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, dilation=dilation, padding=((kernel_size - 1) * dilation) // 2, bias=True), nn.InstanceNorm(out_planes, affine=IN_affine)) elif if_BN: return nn.Sequential( nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, dilation=dilation, padding=((kernel_size - 1) * dilation) // 2, bias=True), nn.BatchNorm2d(out_planes, affine=IN_affine)) else: return nn.Sequential( nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, dilation=dilation, padding=((kernel_size - 1) * dilation) // 2, bias=True))
def __init__(self, ch=128, nframes = 7, input_nc = 3, output_nc = 3, upscale_factor=4, blocknums1 = 5, blocknums2 = 15, non_local = True): super(MUCANV2, self).__init__() self.nframes = nframes self.upscale_factor = upscale_factor # 每个LR搞三个尺度 self.feature_encoder_carb = M.Sequential( M.Conv2d(input_nc, ch, kernel_size=3, stride=1, padding=1), M.LeakyReLU(negative_slope=0.05), CARBBlocks(channel_num=ch, block_num=blocknums1) ) self.fea_L1_conv1 = M.Conv2d(ch, ch, 3, 2, 1) self.fea_L1_conv2 = M.Conv2d(ch, ch, 3, 1, 1) self.fea_L2_conv1 = M.Conv2d(ch, ch, 3, 2, 1) self.fea_L2_conv2 = M.Conv2d(ch, ch, 3, 1, 1) self.lrelu = M.LeakyReLU(negative_slope=0.05) self.AU0 = AU(ch = ch) self.AU1 = AU(ch = ch) self.AU2 = AU(ch = ch) self.UP0 = M.ConvTranspose2d(ch, ch, kernel_size=4, stride=2, padding=1) self.UP1 = M.ConvTranspose2d(ch, ch, kernel_size=4, stride=2, padding=1) if non_local: self.non_local = Separate_non_local(ch, nframes) else: self.non_local = Identi() self.aggre = M.Conv2d(ch * self.nframes, ch, kernel_size=3, stride=1, padding=1) self.carbs = M.Sequential( CARBBlocks(channel_num=ch, block_num=blocknums2), ) self.main_conv = M.Sequential( M.Conv2d(ch, ch*4, kernel_size=3, stride=1, padding=1), M.LeakyReLU(), PixelShuffle(scale=2), # 128 M.Conv2d(ch, ch*2, kernel_size=3, stride=1, padding=1), M.LeakyReLU(), PixelShuffle(scale=2), # 64 M.Conv2d(ch//2, ch//2, kernel_size=3, stride=1, padding=1), M.LeakyReLU(), M.Conv2d(ch//2, 3, kernel_size=3, stride=1, padding=1) )
def __init__(self, channel_nums): super(SDBlock, self).__init__() self.netS = M.Sequential( Conv2d(channel_nums, channel_nums, 3, 1, 1), M.LeakyReLU(negative_slope=0.05), Conv2d(channel_nums, channel_nums, 3, 1, 1) ) self.netD = M.Sequential( Conv2d(channel_nums, channel_nums, 3, 1, 1), M.LeakyReLU(negative_slope=0.05), Conv2d(channel_nums, channel_nums, 3, 1, 1) )
def __init__(self, in_size, out_size, downsample, relu_slope): super(UNetConvBlock, self).__init__() self.block = nn.Sequential( nn.Conv2d(in_size, out_size, kernel_size=3, padding=1, bias=True), nn.LeakyReLU(relu_slope), nn.Conv2d(out_size, out_size, kernel_size=3, padding=1, bias=True), nn.LeakyReLU(relu_slope)) self.downsample = downsample if downsample: self.downsample = conv_down(out_size, out_size, bias=False) self.shortcut = nn.Conv2d(in_size, out_size, kernel_size=1, bias=True)
def __init__(self, K, d, ch): super(AU_CV, self).__init__() self.K = K self.d = d self.conv = M.Sequential( M.Conv2d(ch * self.K, ch, kernel_size=3, stride=1, padding=1), M.LeakyReLU())
def __init__(self, hidden): super(Upsample, self).__init__() self.shirking = M.Conv2d(4 * hidden, hidden, kernel_size=1, stride=1, padding=0) self.sel = SEL(hidden) self.reconstruction = IMDModule(in_channels=hidden) self.conv_last = M.Conv2d(hidden, 3, kernel_size=3, stride=1, padding=1) self.conv_hr1 = M.Conv2d(hidden, hidden, kernel_size=3, stride=1, padding=1) self.conv_hr2 = M.Conv2d(hidden, hidden, kernel_size=3, stride=1, padding=1) self.lrelu = M.LeakyReLU(negative_slope=0.1) self.init_weights()
def get_activation(name="silu"): if name == "silu": module = SiLU() elif name == "relu": module = M.ReLU() elif name == "lrelu": module = M.LeakyReLU(0.1) else: raise AttributeError("Unsupported act type: {}".format(name)) return module
def __init__(self, in_channels, c1=20, c2=12, c3=4, distillation_rate=0.5): super(IMDModule, self).__init__() self.distilled_channels = int(in_channels * distillation_rate) self.remaining_channels = int(in_channels - self.distilled_channels) self.c1 = M.Conv2d(in_channels, c1, 3, stride=1, padding=1) self.c2 = M.Conv2d(c1 - self.distilled_channels, c2, 3, stride=1, padding=1) self.c3 = M.Conv2d(c2 - self.distilled_channels, c3, 3, stride=1, padding=1) self.act = M.LeakyReLU(negative_slope=0.1) self.c5 = M.Conv2d(2 * self.distilled_channels + c3 + in_channels, in_channels, 1, stride=1, padding=0) self.sel = SEL(in_channels) # self.gct = GCT(2*self.distilled_channels + c3 + in_channels) self.init_weights()
def __init__(self, in_channels=3, out_channels=3, mid_channels = 160, ch = 24, blocknums1 = 3, blocknums2 = 3, upscale_factor=4, hsa = False, pixel_shuffle = True, window_size = 5): super(RSDNV4, self).__init__() if hsa: raise NotImplementedError("") else: self.hsa = Identi() self.window_size = window_size self.blocknums1 = blocknums1 self.blocknums2 = blocknums2 # 每个LR搞三个尺度 self.feature_encoder_carb = M.Sequential( M.Conv2d(in_channels, ch, kernel_size=3, stride=1, padding=1), M.LeakyReLU(negative_slope=0.05), CARBBlocks(channel_num=ch, block_num=self.blocknums1) ) self.fea_L1_conv1 = M.Conv2d(ch, ch, 3, 2, 1) self.fea_L1_conv2 = M.Conv2d(ch, ch, 3, 1, 1) self.fea_L2_conv1 = M.Conv2d(ch, ch, 3, 2, 1) self.fea_L2_conv2 = M.Conv2d(ch, ch, 3, 1, 1) self.lrelu = M.LeakyReLU(negative_slope=0.05) self.AU0 = AU(ch = ch) self.AU1 = AU(ch = ch) self.AU2 = AU(ch = ch) self.UP0 = M.ConvTranspose2d(ch, ch, kernel_size=4, stride=2, padding=1) self.UP1 = M.ConvTranspose2d(ch, ch, kernel_size=4, stride=2, padding=1) self.pre_SD_S = M.Sequential( Conv2d(48 + mid_channels + self.window_size*ch, mid_channels, 3, 1, 1), M.LeakyReLU() ) self.convs = M.Sequential( CARBBlocks(channel_num=mid_channels, block_num=self.blocknums2), ) self.hidden = M.Sequential( Conv2d(2*mid_channels, mid_channels, 3, 1, 1), M.LeakyReLU(), CARBBlocks(channel_num=mid_channels, block_num=3), ) self.tail = M.Sequential( CARBBlocks(channel_num=mid_channels, block_num=3), Conv2d(mid_channels, 48, 3, 1, 1) ) if pixel_shuffle: self.trans_HR = PixelShuffle(upscale_factor) else: self.trans_HR = ConvTranspose2d(48, 3, 4, 4, 0, bias=True)
def __init__(self, in_channels=3, out_channels=3, mid_channels = 128, hidden_channels = 3 * 4 * 4, ch = 24, blocknums = 5, upscale_factor=4, hsa = False, pixel_shuffle = False, window_size = 5): super(RSDNV2, self).__init__() if hsa: raise NotImplementedError("") else: self.hsa = Identi() self.window_size = window_size self.blocknums = blocknums self.hidden_channels = hidden_channels # 每个LR搞三个尺度(同时适用于S和D) self.feature_encoder_carb = M.Sequential( M.Conv2d(in_channels, ch, kernel_size=3, stride=1, padding=1), M.LeakyReLU(negative_slope=0.05), CARBBlocks(channel_num=ch, block_num=blocknums) ) self.fea_L1_conv1 = M.Conv2d(ch, ch, 3, 2, 1) self.fea_L1_conv2 = M.Conv2d(ch, ch, 3, 1, 1) self.fea_L2_conv1 = M.Conv2d(ch, ch, 3, 2, 1) self.fea_L2_conv2 = M.Conv2d(ch, ch, 3, 1, 1) self.lrelu = M.LeakyReLU(negative_slope=0.05) self.AU0 = AU(ch = ch) self.AU1 = AU(ch = ch) self.AU2 = AU(ch = ch) self.UP0 = M.ConvTranspose2d(ch, ch, kernel_size=4, stride=2, padding=1) self.UP1 = M.ConvTranspose2d(ch, ch, kernel_size=4, stride=2, padding=1) SDBlocks = [] for _ in range(blocknums): SDBlocks.append(SDBlock(mid_channels)) self.SDBlocks = M.Sequential(*SDBlocks) self.pre_SD_S = M.Sequential( Conv2d(hidden_channels*2 + self.window_size*2*ch, mid_channels, 3, 1, 1), M.LeakyReLU() ) self.pre_SD_D = M.Sequential( Conv2d(hidden_channels*2 + self.window_size*2*ch, mid_channels, 3, 1, 1), M.LeakyReLU() ) self.conv_SD = M.Sequential( Conv2d(mid_channels, hidden_channels, 3, 1, 1), M.LeakyReLU() ) self.convS = Conv2d(mid_channels, hidden_channels, 3, 1, 1) self.convD = Conv2d(mid_channels, hidden_channels, 3, 1, 1) self.convHR = Conv2d(2 * hidden_channels, hidden_channels, 3, 1, 1) if pixel_shuffle: self.trans_S = PixelShuffle(upscale_factor) self.trans_D = PixelShuffle(upscale_factor) self.trans_HR = PixelShuffle(upscale_factor) else: self.trans_S = ConvTranspose2d(hidden_channels, 3, 4, 4, 0, bias=False) self.trans_D = ConvTranspose2d(hidden_channels, 3, 4, 4, 0, bias=False) self.trans_HR = ConvTranspose2d(hidden_channels, 3, 4, 4, 0, bias=False)
def __init__(self, in_channels, out_channels, hidden_channels, init_nums, blocknums, reconstruction_blocks, upscale_factor, pretrained_optical_flow_path, flownet_layers=5, blocktype="resblock", Lambda=0.002): super(BasicVSR, self).__init__() assert blocktype in ("resblock", "shuffleblock", "MobileNeXt") self.in_channels = in_channels self.out_channels = out_channels self.hidden_channels = hidden_channels self.blocknums = blocknums self.upscale_factor = upscale_factor self.reconstruction_blocknums = reconstruction_blocks self.Lambda = Lambda self.flownet = Spynet(num_layers=flownet_layers, pretrain_ckpt_path=pretrained_optical_flow_path, blocktype=blocktype) self.conv1 = M.ConvRelu2d(in_channels, hidden_channels, kernel_size=3, stride=1, padding=1) # need init self.conv2 = ResBlocks(channel_num=hidden_channels, resblock_num=init_nums, blocktype=blocktype) self.border_mode = "REPLICATE" self.feature_extracter = ResBlocks(channel_num=hidden_channels, resblock_num=blocknums, blocktype=blocktype) self.reconstruction = ResBlocks(channel_num=hidden_channels, resblock_num=reconstruction_blocks, blocktype=blocktype) self.conv3 = M.ConvRelu2d(2 * hidden_channels, hidden_channels, kernel_size=3, stride=1, padding=1) # need init self.conv4 = M.ConvRelu2d(2 * hidden_channels, hidden_channels, kernel_size=3, stride=1, padding=1) # need init self.upsample1 = PixelShufflePack(hidden_channels, hidden_channels, 2, upsample_kernel=3) self.upsample2 = PixelShufflePack(hidden_channels, 64, 2, upsample_kernel=3) self.conv_hr = M.Conv2d(64, 64, 3, 1, 1) # need init self.conv_last = M.Conv2d(64, out_channels, 3, 1, 1) self.lrelu = M.LeakyReLU(negative_slope=0.01)