Upload network_fbcnn.py

network_fbcnn.py

@@ -0,0 +1,337 @@
+from collections import OrderedDict
+import torch
+import torch.nn as nn
+import numpy as np
+import torch.nn.functional as F
+import torchvision.models as models
+
+'''
+# --------------------------------------------
+# Advanced nn.Sequential
+# https://github.com/xinntao/BasicSR
+# --------------------------------------------
+'''
+
+
+def sequential(*args):
+    """Advanced nn.Sequential.
+
+    Args:
+        nn.Sequential, nn.Module
+
+    Returns:
+        nn.Sequential
+    """
+    if len(args) == 1:
+        if isinstance(args[0], OrderedDict):
+            raise NotImplementedError('sequential does not support OrderedDict input.')
+        return args[0]  # No sequential is needed.
+    modules = []
+    for module in args:
+        if isinstance(module, nn.Sequential):
+            for submodule in module.children():
+                modules.append(submodule)
+        elif isinstance(module, nn.Module):
+            modules.append(module)
+    return nn.Sequential(*modules)
+
+# --------------------------------------------
+# return nn.Sequantial of (Conv + BN + ReLU)
+# --------------------------------------------
+def conv(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1, bias=True, mode='CBR', negative_slope=0.2):
+    L = []
+    for t in mode:
+        if t == 'C':
+            L.append(nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride, padding=padding, bias=bias))
+        elif t == 'T':
+            L.append(nn.ConvTranspose2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride, padding=padding, bias=bias))
+        elif t == 'B':
+            L.append(nn.BatchNorm2d(out_channels, momentum=0.9, eps=1e-04, affine=True))
+        elif t == 'I':
+            L.append(nn.InstanceNorm2d(out_channels, affine=True))
+        elif t == 'R':
+            L.append(nn.ReLU(inplace=True))
+        elif t == 'r':
+            L.append(nn.ReLU(inplace=False))
+        elif t == 'L':
+            L.append(nn.LeakyReLU(negative_slope=negative_slope, inplace=True))
+        elif t == 'l':
+            L.append(nn.LeakyReLU(negative_slope=negative_slope, inplace=False))
+        elif t == '2':
+            L.append(nn.PixelShuffle(upscale_factor=2))
+        elif t == '3':
+            L.append(nn.PixelShuffle(upscale_factor=3))
+        elif t == '4':
+            L.append(nn.PixelShuffle(upscale_factor=4))
+        elif t == 'U':
+            L.append(nn.Upsample(scale_factor=2, mode='nearest'))
+        elif t == 'u':
+            L.append(nn.Upsample(scale_factor=3, mode='nearest'))
+        elif t == 'v':
+            L.append(nn.Upsample(scale_factor=4, mode='nearest'))
+        elif t == 'M':
+            L.append(nn.MaxPool2d(kernel_size=kernel_size, stride=stride, padding=0))
+        elif t == 'A':
+            L.append(nn.AvgPool2d(kernel_size=kernel_size, stride=stride, padding=0))
+        else:
+            raise NotImplementedError('Undefined type: '.format(t))
+    return sequential(*L)
+
+# --------------------------------------------
+# Res Block: x + conv(relu(conv(x)))
+# --------------------------------------------
+class ResBlock(nn.Module):
+    def __init__(self, in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1, bias=True, mode='CRC', negative_slope=0.2):
+        super(ResBlock, self).__init__()
+
+        assert in_channels == out_channels, 'Only support in_channels==out_channels.'
+        if mode[0] in ['R', 'L']:
+            mode = mode[0].lower() + mode[1:]
+
+        self.res = conv(in_channels, out_channels, kernel_size, stride, padding, bias, mode, negative_slope)
+
+    def forward(self, x):
+        res = self.res(x)
+        return x + res
+
+# --------------------------------------------
+# conv + subp (+ relu)
+# --------------------------------------------
+def upsample_pixelshuffle(in_channels=64, out_channels=3, kernel_size=3, stride=1, padding=1, bias=True, mode='2R', negative_slope=0.2):
+    assert len(mode)<4 and mode[0] in ['2', '3', '4'], 'mode examples: 2, 2R, 2BR, 3, ..., 4BR.'
+    up1 = conv(in_channels, out_channels * (int(mode[0]) ** 2), kernel_size, stride, padding, bias, mode='C'+mode, negative_slope=negative_slope)
+    return up1
+
+
+# --------------------------------------------
+# nearest_upsample + conv (+ R)
+# --------------------------------------------
+def upsample_upconv(in_channels=64, out_channels=3, kernel_size=3, stride=1, padding=1, bias=True, mode='2R', negative_slope=0.2):
+    assert len(mode)<4 and mode[0] in ['2', '3', '4'], 'mode examples: 2, 2R, 2BR, 3, ..., 4BR'
+    if mode[0] == '2':
+        uc = 'UC'
+    elif mode[0] == '3':
+        uc = 'uC'
+    elif mode[0] == '4':
+        uc = 'vC'
+    mode = mode.replace(mode[0], uc)
+    up1 = conv(in_channels, out_channels, kernel_size, stride, padding, bias, mode=mode, negative_slope=negative_slope)
+    return up1
+
+
+# --------------------------------------------
+# convTranspose (+ relu)
+# --------------------------------------------
+def upsample_convtranspose(in_channels=64, out_channels=3, kernel_size=2, stride=2, padding=0, bias=True, mode='2R', negative_slope=0.2):
+    assert len(mode)<4 and mode[0] in ['2', '3', '4'], 'mode examples: 2, 2R, 2BR, 3, ..., 4BR.'
+    kernel_size = int(mode[0])
+    stride = int(mode[0])
+    mode = mode.replace(mode[0], 'T')
+    up1 = conv(in_channels, out_channels, kernel_size, stride, padding, bias, mode, negative_slope)
+    return up1
+
+
+'''
+# --------------------------------------------
+# Downsampler
+# Kai Zhang, https://github.com/cszn/KAIR
+# --------------------------------------------
+# downsample_strideconv
+# downsample_maxpool
+# downsample_avgpool
+# --------------------------------------------
+'''
+
+
+# --------------------------------------------
+# strideconv (+ relu)
+# --------------------------------------------
+def downsample_strideconv(in_channels=64, out_channels=64, kernel_size=2, stride=2, padding=0, bias=True, mode='2R', negative_slope=0.2):
+    assert len(mode)<4 and mode[0] in ['2', '3', '4'], 'mode examples: 2, 2R, 2BR, 3, ..., 4BR.'
+    kernel_size = int(mode[0])
+    stride = int(mode[0])
+    mode = mode.replace(mode[0], 'C')
+    down1 = conv(in_channels, out_channels, kernel_size, stride, padding, bias, mode, negative_slope)
+    return down1
+
+
+# --------------------------------------------
+# maxpooling + conv (+ relu)
+# --------------------------------------------
+def downsample_maxpool(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=0, bias=True, mode='2R', negative_slope=0.2):
+    assert len(mode)<4 and mode[0] in ['2', '3'], 'mode examples: 2, 2R, 2BR, 3, ..., 3BR.'
+    kernel_size_pool = int(mode[0])
+    stride_pool = int(mode[0])
+    mode = mode.replace(mode[0], 'MC')
+    pool = conv(kernel_size=kernel_size_pool, stride=stride_pool, mode=mode[0], negative_slope=negative_slope)
+    pool_tail = conv(in_channels, out_channels, kernel_size, stride, padding, bias, mode=mode[1:], negative_slope=negative_slope)
+    return sequential(pool, pool_tail)
+
+
+# --------------------------------------------
+# averagepooling + conv (+ relu)
+# --------------------------------------------
+def downsample_avgpool(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1, bias=True, mode='2R', negative_slope=0.2):
+    assert len(mode)<4 and mode[0] in ['2', '3'], 'mode examples: 2, 2R, 2BR, 3, ..., 3BR.'
+    kernel_size_pool = int(mode[0])
+    stride_pool = int(mode[0])
+    mode = mode.replace(mode[0], 'AC')
+    pool = conv(kernel_size=kernel_size_pool, stride=stride_pool, mode=mode[0], negative_slope=negative_slope)
+    pool_tail = conv(in_channels, out_channels, kernel_size, stride, padding, bias, mode=mode[1:], negative_slope=negative_slope)
+    return sequential(pool, pool_tail)
+
+
+
+class QFAttention(nn.Module):
+    def __init__(self, in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1, bias=True, mode='CRC', negative_slope=0.2):
+        super(QFAttention, self).__init__()
+
+        assert in_channels == out_channels, 'Only support in_channels==out_channels.'
+        if mode[0] in ['R', 'L']:
+            mode = mode[0].lower() + mode[1:]
+
+        self.res = conv(in_channels, out_channels, kernel_size, stride, padding, bias, mode, negative_slope)
+
+    def forward(self, x, gamma, beta):
+        gamma = gamma.unsqueeze(-1).unsqueeze(-1)
+        beta = beta.unsqueeze(-1).unsqueeze(-1)
+        res = (gamma)*self.res(x) + beta
+        return x + res
+
+
+class FBCNN(nn.Module):
+    def __init__(self, in_nc=3, out_nc=3, nc=[64, 128, 256, 512], nb=4, act_mode='R', downsample_mode='strideconv',
+                 upsample_mode='convtranspose'):
+        super(FBCNN, self).__init__()
+
+        self.m_head = conv(in_nc, nc[0], bias=True, mode='C')
+        self.nb = nb
+        self.nc = nc
+        # downsample
+        if downsample_mode == 'avgpool':
+            downsample_block = downsample_avgpool
+        elif downsample_mode == 'maxpool':
+            downsample_block = downsample_maxpool
+        elif downsample_mode == 'strideconv':
+            downsample_block = downsample_strideconv
+        else:
+            raise NotImplementedError('downsample mode [{:s}] is not found'.format(downsample_mode))
+
+        self.m_down1 = sequential(
+            *[ResBlock(nc[0], nc[0], bias=True, mode='C' + act_mode + 'C') for _ in range(nb)],
+            downsample_block(nc[0], nc[1], bias=True, mode='2'))
+        self.m_down2 = sequential(
+            *[ResBlock(nc[1], nc[1], bias=True, mode='C' + act_mode + 'C') for _ in range(nb)],
+            downsample_block(nc[1], nc[2], bias=True, mode='2'))
+        self.m_down3 = sequential(
+            *[ResBlock(nc[2], nc[2], bias=True, mode='C' + act_mode + 'C') for _ in range(nb)],
+            downsample_block(nc[2], nc[3], bias=True, mode='2'))
+
+        self.m_body_encoder = sequential(
+            *[ResBlock(nc[3], nc[3], bias=True, mode='C' + act_mode + 'C') for _ in range(nb)])
+
+        self.m_body_decoder = sequential(
+            *[ResBlock(nc[3], nc[3], bias=True, mode='C' + act_mode + 'C') for _ in range(nb)])
+
+        # upsample
+        if upsample_mode == 'upconv':
+            upsample_block = upsample_upconv
+        elif upsample_mode == 'pixelshuffle':
+            upsample_block = upsample_pixelshuffle
+        elif upsample_mode == 'convtranspose':
+            upsample_block = upsample_convtranspose
+        else:
+            raise NotImplementedError('upsample mode [{:s}] is not found'.format(upsample_mode))
+
+        self.m_up3 = nn.ModuleList([upsample_block(nc[3], nc[2], bias=True, mode='2'),
+                                  *[QFAttention(nc[2], nc[2], bias=True, mode='C' + act_mode + 'C') for _ in range(nb)]])
+
+        self.m_up2 = nn.ModuleList([upsample_block(nc[2], nc[1], bias=True, mode='2'),
+                                  *[QFAttention(nc[1], nc[1], bias=True, mode='C' + act_mode + 'C') for _ in range(nb)]])
+
+        self.m_up1 = nn.ModuleList([upsample_block(nc[1], nc[0], bias=True, mode='2'),
+                                  *[QFAttention(nc[0], nc[0], bias=True, mode='C' + act_mode + 'C') for _ in range(nb)]])
+
+
+        self.m_tail = conv(nc[0], out_nc, bias=True, mode='C')
+
+
+        self.qf_pred = sequential(*[ResBlock(nc[3], nc[3], bias=True, mode='C' + act_mode + 'C') for _ in range(nb)],
+                                  torch.nn.AdaptiveAvgPool2d((1,1)),
+                                  torch.nn.Flatten(),
+                                  torch.nn.Linear(512, 512), 
+                                  nn.ReLU(),
+                                  torch.nn.Linear(512, 512),
+                                  nn.ReLU(),
+                                  torch.nn.Linear(512, 1),
+                                  nn.Sigmoid()
+                                )
+
+        self.qf_embed = sequential(torch.nn.Linear(1, 512),
+                                  nn.ReLU(),
+                                  torch.nn.Linear(512, 512),
+                                  nn.ReLU(),
+                                  torch.nn.Linear(512, 512),
+                                  nn.ReLU()
+                                )
+
+        self.to_gamma_3 = sequential(torch.nn.Linear(512, nc[2]),nn.Sigmoid())
+        self.to_beta_3 =  sequential(torch.nn.Linear(512, nc[2]),nn.Tanh())
+        self.to_gamma_2 = sequential(torch.nn.Linear(512, nc[1]),nn.Sigmoid())
+        self.to_beta_2 =  sequential(torch.nn.Linear(512, nc[1]),nn.Tanh())
+        self.to_gamma_1 = sequential(torch.nn.Linear(512, nc[0]),nn.Sigmoid())
+        self.to_beta_1 =  sequential(torch.nn.Linear(512, nc[0]),nn.Tanh())
+
+
+    def forward(self, x, qf_input=None):
+
+        h, w = x.size()[-2:]
+        paddingBottom = int(np.ceil(h / 8) * 8 - h)
+        paddingRight = int(np.ceil(w / 8) * 8 - w)
+        x = nn.ReplicationPad2d((0, paddingRight, 0, paddingBottom))(x)
+
+        x1 = self.m_head(x)
+        x2 = self.m_down1(x1)
+        x3 = self.m_down2(x2)
+        x4 = self.m_down3(x3)
+        x = self.m_body_encoder(x4)
+        qf = self.qf_pred(x)
+        x = self.m_body_decoder(x)
+        qf_embedding = self.qf_embed(qf_input) if qf_input is not None else self.qf_embed(qf)
+        gamma_3 = self.to_gamma_3(qf_embedding)
+        beta_3 = self.to_beta_3(qf_embedding)
+
+        gamma_2 = self.to_gamma_2(qf_embedding)
+        beta_2 = self.to_beta_2(qf_embedding)
+
+        gamma_1 = self.to_gamma_1(qf_embedding)
+        beta_1 = self.to_beta_1(qf_embedding)
+
+
+        x = x + x4
+        x = self.m_up3[0](x)
+        for i in range(self.nb):
+            x = self.m_up3[i+1](x, gamma_3,beta_3)
+
+        x = x + x3
+
+        x = self.m_up2[0](x)
+        for i in range(self.nb):
+            x = self.m_up2[i+1](x, gamma_2, beta_2)
+        x = x + x2
+
+        x = self.m_up1[0](x)
+        for i in range(self.nb):
+            x = self.m_up1[i+1](x, gamma_1, beta_1)
+
+        x = x + x1
+        x = self.m_tail(x)
+        x = x[..., :h, :w]
+
+        return x, qf
+
+if __name__ == "__main__":
+    x = torch.randn(1, 3, 96, 96)#.cuda()#.to(torch.device('cuda'))
+    fbar=FBAR()
+    y,qf = fbar(x)
+    print(y.shape,qf.shape)