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)
