Upload 26 files

.gitattributes

@@ -33,3 +33,9 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/ARAD_1K_0001.mat filter=lfs diff=lfs merge=lfs -text
+assets/ARAD_1K_0002.mat filter=lfs diff=lfs merge=lfs -text
+assets/ARAD_1K_0003.mat filter=lfs diff=lfs merge=lfs -text
+assets/ARAD_1K_0004.mat filter=lfs diff=lfs merge=lfs -text
+assets/ARAD_1K_0005.mat filter=lfs diff=lfs merge=lfs -text
+assets/ARAD_1K_0006.mat filter=lfs diff=lfs merge=lfs -text

app.py

@@ -0,0 +1,125 @@
+import torch
+import numpy as np
+import gradio as gr
+import cv2
+import h5py
+from test_develop_code.architecture import model_generator
+from PIL import Image
+
+
+# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+device = torch.device("cpu")
+model = model_generator("mst_plus_plus", "mst_plus_plus.pth").to(device)
+model.eval()
+wavelengths = np.linspace(400, 700, 31)
+
+
+def wavelength_to_rgb(wl: float) -> tuple:
+    if 380 <= wl <= 440:
+        R = -(wl - 440) / (440 - 380)
+        G = 0.0
+        B = 1.0
+    elif 440 < wl <= 490:
+        R = 0.0
+        G = (wl - 440) / (490 - 440)
+        B = 1.0
+    elif 490 < wl <= 510:
+        R = 0.0
+        G = 1.0
+        B = -(wl - 510) / (510 - 490)
+    elif 510 < wl <= 580:
+        R = (wl - 510) / (580 - 510)
+        G = 1.0
+        B = 0.0
+    elif 580 < wl <= 645:
+        R = 1.0
+        G = -(wl - 645) / (645 - 580)
+        B = 0.0
+    elif 645 < wl <= 700:
+        R = 1.0
+        G = 0.0
+        B = 0.0
+    else:
+        R = G = B = 0.0
+    return (max(R, 0.0), max(G, 0.0), max(B, 0.0))
+
+
+def predict(img):
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+    img = img.astype(np.float32) / 255.0
+    img = np.transpose(img, (2, 0, 1))
+    img_tensor = torch.from_numpy(img).unsqueeze(0).to(device)
+    with torch.no_grad():
+        pred = model(img_tensor)
+    pred = pred.squeeze(0).cpu().numpy()
+    pred = np.clip(pred, 0, 1)
+    return pred
+
+
+def visualize_channel(cube: np.ndarray, index: int) -> Image.Image:
+    if cube is None:
+        return None
+    band = cube[index]
+    band = (band - band.min()) / (band.max() - band.min() + 1e-8)
+    color = wavelength_to_rgb(wavelengths[index])
+    rgb = np.stack([band * c for c in color], axis=-1)
+    rgb = (rgb * 255).astype(np.uint8)
+    return Image.fromarray(rgb)
+
+
+def load_mat(mat_file) -> np.ndarray:
+    with h5py.File(mat_file.name, "r") as f:
+        cube = np.array(f["cube"])
+    cube = np.transpose(cube, (0, 2, 1))
+    cube = np.clip(cube, 0, 1)
+    return cube
+
+
+def reset_all():
+    return None, None, None, None, 0
+
+
+with gr.Blocks() as demo:
+    gr.Markdown("## Spectral Reconstruction")
+
+    with gr.Row():
+        with gr.Column():
+            rgb_input = gr.Image(type="numpy", label="Upload RGB Image")
+            pred_state = gr.State()
+        with gr.Column():
+            pred_output = gr.Image(label="Prediction Visualization")
+            pred_slider = gr.Slider(minimum=0, maximum=30, step=1, label="Channel (Prediction)", value=0)
+
+    with gr.Row():
+        with gr.Column():
+            mat_input = gr.File(label="Upload .mat file (Ground Truth)")
+            gt_state = gr.State()
+        with gr.Column():
+            gt_output = gr.Image(label="Ground Truth Visualization")
+            gt_slider = gr.Slider(minimum=0, maximum=30, step=1, label="Channel (Ground Truth)", value=0)
+
+    clear_btn = gr.Button("Clear")
+    rgb_input.change(fn=predict, inputs=rgb_input, outputs=pred_state)
+    pred_slider.change(fn=visualize_channel, inputs=[pred_state, pred_slider], outputs=pred_output)
+
+    mat_input.change(fn=load_mat, inputs=mat_input, outputs=gt_state)
+    gt_slider.change(fn=visualize_channel, inputs=[gt_state, gt_slider], outputs=gt_output)
+
+    clear_btn.click(fn=reset_all, outputs=[rgb_input, pred_output, pred_state, gt_state, mat_input])
+
+    gr.Examples(
+        examples=[
+            ["assets/ARAD_1K_0001.jpg", 0, "assets/ARAD_1K_0001.mat", 0],
+            ["assets/ARAD_1K_0002.jpg", 0, "assets/ARAD_1K_0002.mat", 0],
+            ["assets/ARAD_1K_0003.jpg", 0, "assets/ARAD_1K_0003.mat", 0],
+            ["assets/ARAD_1K_0004.jpg", 0, "assets/ARAD_1K_0004.mat", 0],
+            ["assets/ARAD_1K_0005.jpg", 0, "assets/ARAD_1K_0005.mat", 0],
+        ],
+        inputs=[rgb_input, pred_slider, mat_input, gt_slider],
+        outputs=[pred_output, gt_output],
+        label="Try Examples"
+    )
+
+
+if __name__ == "__main__":
+    demo.launch()

assets/ARAD_1K_0001.mat

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d75d3689827d69fdb600ef49eceec4692615db5f0bb8882b41cc3cc0f29a139f
+size 21896837

assets/ARAD_1K_0002.mat

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5ee379c0574ac7a225cdd6cd39f942441f48d3401f4d1abe0e26a664ce8f1af3
+size 22920548

assets/ARAD_1K_0003.mat

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0729ab86549ca910803e2d0dcdc76c10087f9a0f4efbb252e08b564bd3a9a741
+size 21218912

assets/ARAD_1K_0004.mat

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:abe2de20dc3b76e1a1701292e7ad1774aef643a88366084eef58d92f7e280a6d
+size 22377073

assets/ARAD_1K_0005.mat

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6fd2f928230e076be0d2c707b657a5e6cdf84108cd53d86aab91450d8ce4222b
+size 21042006

assets/ARAD_1K_0006.mat

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:32d2e4ff2732cd16bf3a9103592c9af304598307dac6967990bd25f00b47d9f7
+size 22078636

mst_plus_plus.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d285430cc688d08434582eee71bae2d82661be7997af1a68d6636ec25f7a3421
+size 6580074

requirements.txt

@@ -0,0 +1,12 @@
+opencv-python==4.11.0.86
+einops==0.8.1
+torchvision==0.22.0
+torch==2.7.0
+scipy==1.15.2
+h5py==3.13.0
+hdf5storage==0.1.19
+tqdm==4.67.1
+gdown==5.2.0
+matplotlib==3.10.1
+gradio==5.29.0
+

test_challenge_code/architecture/HDNet.py

@@ -0,0 +1,397 @@
+import torch
+import torch.nn as nn
+def default_conv(in_channels, out_channels, kernel_size, bias=True):
+    return nn.Conv2d(
+        in_channels, out_channels, kernel_size,
+        padding=(kernel_size//2), bias=bias)
+
+class MeanShift(nn.Conv2d):
+    def __init__(self, rgb_range, rgb_mean, rgb_std, sign=-1):
+        super(MeanShift, self).__init__(3, 3, kernel_size=1)
+        std = torch.Tensor(rgb_std)
+        self.weight.data = torch.eye(3).view(3, 3, 1, 1)
+        self.weight.data.div_(std.view(3, 1, 1, 1))
+        self.bias.data = sign * rgb_range * torch.Tensor(rgb_mean)
+        self.bias.data.div_(std)
+        self.requires_grad = False
+
+class BasicBlock(nn.Sequential):
+    def __init__(
+        self, in_channels, out_channels, kernel_size, stride=1, bias=False,
+        bn=True, act=nn.ReLU(True)):
+
+        m = [nn.Conv2d(
+            in_channels, out_channels, kernel_size,
+            padding=(kernel_size//2), stride=stride, bias=bias)
+        ]
+        if bn: m.append(nn.BatchNorm2d(out_channels))
+        if act is not None: m.append(act)
+        super(BasicBlock, self).__init__(*m)
+
+class ResBlock(nn.Module):
+    def __init__(
+        self, conv=default_conv, n_feat=31, kernel_size=3,
+        bias=True, bn=False, act=nn.ReLU(True), res_scale=1):
+
+        super(ResBlock, self).__init__()
+        m = []
+        for i in range(2):
+            m.append(conv(n_feat, n_feat, kernel_size, bias=bias))
+            if bn: m.append(nn.BatchNorm2d(n_feat))
+            if i == 0: m.append(act)
+
+        self.body = nn.Sequential(*m)
+        self.res_scale = res_scale
+
+    def forward(self, x):
+        res = self.body(x).mul(self.res_scale)
+        res += x
+
+        return res
+
+class Upsampler(nn.Sequential):
+    def __init__(self, conv, scale, n_feat, bn=False, act=False, bias=True):
+
+        m = []
+        if (scale & (scale - 1)) == 0:    # Is scale = 2^n?
+            for _ in range(int(math.log(scale, 2))):
+                m.append(conv(n_feat, 4 * n_feat, 3, bias))
+                m.append(nn.PixelShuffle(2))
+                if bn: m.append(nn.BatchNorm2d(n_feat))
+                if act: m.append(act())
+        elif scale == 3:
+            m.append(conv(n_feat, 9 * n_feat, 3, bias))
+            m.append(nn.PixelShuffle(3))
+            if bn: m.append(nn.BatchNorm2d(n_feat))
+            if act: m.append(act())
+        else:
+            raise NotImplementedError
+
+        super(Upsampler, self).__init__(*m)
+
+## add SELayer
+class SELayer(nn.Module):
+    def __init__(self, channel, reduction=16):
+        super(SELayer, self).__init__()
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.conv_du = nn.Sequential(
+                nn.Conv2d(channel, channel // reduction, 1, padding=0, bias=True),
+                nn.ReLU(inplace=True),
+                nn.Conv2d(channel // reduction, channel, 1, padding=0, bias=True),
+                nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        y = self.avg_pool(x)
+        y = self.conv_du(y)
+        return x * y
+
+## add SEResBlock
+class SEResBlock(nn.Module):
+    def __init__(
+        self, conv, n_feat, kernel_size, reduction,
+        bias=True, bn=False, act=nn.ReLU(True), res_scale=1):
+
+        super(SEResBlock, self).__init__()
+        modules_body = []
+        for i in range(2):
+            modules_body.append(conv(n_feat, n_feat, kernel_size, bias=bias))
+            if bn: modules_body.append(nn.BatchNorm2d(n_feat))
+            if i == 0: modules_body.append(act)
+        modules_body.append(SELayer(n_feat, reduction))
+        self.body = nn.Sequential(*modules_body)
+        self.res_scale = res_scale
+
+    def forward(self, x):
+        res = self.body(x)
+        #res = self.body(x).mul(self.res_scale)
+        res += x
+
+        return res
+
+
+_NORM_BONE = False
+
+def constant_init(module, val, bias=0):
+    if hasattr(module, 'weight') and module.weight is not None:
+        nn.init.constant_(module.weight, val)
+    if hasattr(module, 'bias') and module.bias is not None:
+        nn.init.constant_(module.bias, bias)
+
+
+def kaiming_init(module,
+                 a=0,
+                 mode='fan_out',
+                 nonlinearity='relu',
+                 bias=0,
+                 distribution='normal'):
+    assert distribution in ['uniform', 'normal']
+    if distribution == 'uniform':
+        nn.init.kaiming_uniform_(
+            module.weight, a=a, mode=mode, nonlinearity=nonlinearity)
+    else:
+        nn.init.kaiming_normal_(
+            module.weight, a=a, mode=mode, nonlinearity=nonlinearity)
+    if hasattr(module, 'bias') and module.bias is not None:
+        nn.init.constant_(module.bias, bias)
+
+# depthwise-separable convolution (DSC)
+class DSC(nn.Module):
+
+    def __init__(self, nin: int) -> None:
+        super(DSC, self).__init__()
+        self.conv_dws = nn.Conv2d(
+            nin, nin, kernel_size=1, stride=1, padding=0, groups=nin
+        )
+        self.bn_dws = nn.BatchNorm2d(nin, momentum=0.9)
+        self.relu_dws = nn.ReLU(inplace=False)
+
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=1, padding=1)
+
+        self.conv_point = nn.Conv2d(
+            nin, 1, kernel_size=1, stride=1, padding=0, groups=1
+        )
+        self.bn_point = nn.BatchNorm2d(1, momentum=0.9)
+        self.relu_point = nn.ReLU(inplace=False)
+
+        self.softmax = nn.Softmax(dim=2)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        out = self.conv_dws(x)
+        out = self.bn_dws(out)
+        out = self.relu_dws(out)
+
+        out = self.maxpool(out)
+
+        out = self.conv_point(out)
+        out = self.bn_point(out)
+        out = self.relu_point(out)
+
+        m, n, p, q = out.shape
+        out = self.softmax(out.view(m, n, -1))
+        out = out.view(m, n, p, q)
+
+        out = out.expand(x.shape[0], x.shape[1], x.shape[2], x.shape[3])
+
+        out = torch.mul(out, x)
+
+        out = out + x
+
+        return out
+
+# Efficient Feature Fusion(EFF)
+class EFF(nn.Module):
+    def __init__(self, nin: int, nout: int, num_splits: int) -> None:
+        super(EFF, self).__init__()
+
+        assert nin % num_splits == 0
+
+        self.nin = nin
+        self.nout = nout
+        self.num_splits = num_splits
+        self.subspaces = nn.ModuleList(
+            [DSC(int(self.nin / self.num_splits)) for i in range(self.num_splits)]
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        sub_feat = torch.chunk(x, self.num_splits, dim=1)
+        out = []
+        for idx, l in enumerate(self.subspaces):
+            out.append(self.subspaces[idx](sub_feat[idx]))
+        out = torch.cat(out, dim=1)
+
+        return out
+
+
+# spatial-spectral domain attention learning(SDL)
+class SDL_attention(nn.Module):
+    def __init__(self, inplanes, planes, kernel_size=1, stride=1):
+        super(SDL_attention, self).__init__()
+
+        self.inplanes = inplanes
+        self.inter_planes = planes // 2
+        self.planes = planes
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.padding = (kernel_size-1)//2
+
+        self.conv_q_right = nn.Conv2d(self.inplanes, 1, kernel_size=1, stride=stride, padding=0, bias=False)
+        self.conv_v_right = nn.Conv2d(self.inplanes, self.inter_planes, kernel_size=1, stride=stride, padding=0, bias=False)
+        self.conv_up = nn.Conv2d(self.inter_planes, self.planes, kernel_size=1, stride=1, padding=0, bias=False)
+        self.softmax_right = nn.Softmax(dim=2)
+        self.sigmoid = nn.Sigmoid()
+
+        self.conv_q_left = nn.Conv2d(self.inplanes, self.inter_planes, kernel_size=1, stride=stride, padding=0, bias=False)   #g
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.conv_v_left = nn.Conv2d(self.inplanes, self.inter_planes, kernel_size=1, stride=stride, padding=0, bias=False)   #theta
+        self.softmax_left = nn.Softmax(dim=2)
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        kaiming_init(self.conv_q_right, mode='fan_in')
+        kaiming_init(self.conv_v_right, mode='fan_in')
+        kaiming_init(self.conv_q_left, mode='fan_in')
+        kaiming_init(self.conv_v_left, mode='fan_in')
+
+        self.conv_q_right.inited = True
+        self.conv_v_right.inited = True
+        self.conv_q_left.inited = True
+        self.conv_v_left.inited = True
+    # HR spatial attention
+    def spatial_attention(self, x):
+        input_x = self.conv_v_right(x)
+        batch, channel, height, width = input_x.size()
+
+        input_x = input_x.view(batch, channel, height * width)
+        context_mask = self.conv_q_right(x)
+        context_mask = context_mask.view(batch, 1, height * width)
+        context_mask = self.softmax_right(context_mask)
+
+        context = torch.matmul(input_x, context_mask.transpose(1,2))
+        context = context.unsqueeze(-1)
+        context = self.conv_up(context)
+
+        mask_ch = self.sigmoid(context)
+
+        out = x * mask_ch
+
+        return out
+    # HR spectral attention
+    def spectral_attention(self, x):
+
+        g_x = self.conv_q_left(x)
+        batch, channel, height, width = g_x.size()
+
+        avg_x = self.avg_pool(g_x)
+        batch, channel, avg_x_h, avg_x_w = avg_x.size()
+
+        avg_x = avg_x.view(batch, channel, avg_x_h * avg_x_w).permute(0, 2, 1)
+        theta_x = self.conv_v_left(x).view(batch, self.inter_planes, height * width)
+        context = torch.matmul(avg_x, theta_x)
+        context = self.softmax_left(context)
+        context = context.view(batch, 1, height, width)
+
+        mask_sp = self.sigmoid(context)
+
+        out = x * mask_sp
+
+        return out
+
+    def forward(self, x):
+        context_spectral = self.spectral_attention(x)
+        context_spatial = self.spatial_attention(x)
+        out = context_spatial + context_spectral
+        return out
+
+
+class HDNet(nn.Module):
+
+    def __init__(self, in_ch=3, out_ch=31, conv=default_conv):
+        super(HDNet, self).__init__()
+
+        n_resblocks = 32
+        n_feats = 48
+        kernel_size = 3
+        act = nn.ReLU(True)
+
+        # define head module
+        m_head = [conv(in_ch, n_feats, kernel_size)]
+
+        # define body module
+        m_body = [
+            ResBlock(
+                conv, n_feats, kernel_size, act=act, res_scale= 1
+            ) for _ in range(n_resblocks)
+        ]
+        m_body.append(SDL_attention(inplanes = n_feats, planes = n_feats))
+        m_body.append(EFF(nin=n_feats, nout=n_feats, num_splits=4))
+
+        for i in range(1, n_resblocks):
+            m_body.append(ResBlock(
+                conv, n_feats, kernel_size, act=act, res_scale= 1
+            ))
+
+        m_body.append(conv(n_feats, n_feats, kernel_size))
+
+        m_tail = [conv(n_feats, out_ch, kernel_size)]
+
+        self.head = nn.Sequential(*m_head)
+        self.body = nn.Sequential(*m_body)
+        self.tail = nn.Sequential(*m_tail)
+
+    def forward(self, x):
+        x = self.head(x)
+
+        res = self.body(x)
+        res += x
+
+        x = self.tail(res)
+
+        return x
+
+# frequency domain learning(FDL)
+class FDL(nn.Module):
+    def __init__(self, loss_weight=1.0, alpha=1.0, patch_factor=1, ave_spectrum=False, log_matrix=False, batch_matrix=False):
+        super(FDL, self).__init__()
+        self.loss_weight = loss_weight
+        self.alpha = alpha
+        self.patch_factor = patch_factor
+        self.ave_spectrum = ave_spectrum
+        self.log_matrix = log_matrix
+        self.batch_matrix = batch_matrix
+
+    def tensor2freq(self, x):
+        patch_factor = self.patch_factor
+        _, _, h, w = x.shape
+        assert h % patch_factor == 0 and w % patch_factor == 0, (
+            'Patch factor should be divisible by image height and width')
+        patch_list = []
+        patch_h = h // patch_factor
+        patch_w = w // patch_factor
+        for i in range(patch_factor):
+            for j in range(patch_factor):
+                patch_list.append(x[:, :, i * patch_h:(i + 1) * patch_h, j * patch_w:(j + 1) * patch_w])
+
+        y = torch.stack(patch_list, 1)
+
+        return torch.rfft(y, 2, onesided=False, normalized=True)
+
+    def loss_formulation(self, recon_freq, real_freq, matrix=None):
+        if matrix is not None:
+            weight_matrix = matrix.detach()
+        else:
+            matrix_tmp = (recon_freq - real_freq) ** 2
+            matrix_tmp = torch.sqrt(matrix_tmp[..., 0] + matrix_tmp[..., 1]) ** self.alpha
+            if self.log_matrix:
+                matrix_tmp = torch.log(matrix_tmp + 1.0)
+
+            if self.batch_matrix:
+                matrix_tmp = matrix_tmp / matrix_tmp.max()
+            else:
+                matrix_tmp = matrix_tmp / matrix_tmp.max(-1).values.max(-1).values[:, :, :, None, None]
+
+            matrix_tmp[torch.isnan(matrix_tmp)] = 0.0
+            matrix_tmp = torch.clamp(matrix_tmp, min=0.0, max=1.0)
+            weight_matrix = matrix_tmp.clone().detach()
+
+        assert weight_matrix.min().item() >= 0 and weight_matrix.max().item() <= 1, (
+            'The values of spectrum weight matrix should be in the range [0, 1], '
+            'but got Min: %.10f Max: %.10f' % (weight_matrix.min().item(), weight_matrix.max().item()))
+
+        tmp = (recon_freq - real_freq) ** 2
+        freq_distance = tmp[..., 0] + tmp[..., 1]
+
+        loss = weight_matrix * freq_distance
+        return torch.mean(loss)
+
+    def forward(self, pred, target, matrix=None, **kwargs):
+
+        pred_freq = self.tensor2freq(pred)
+        target_freq = self.tensor2freq(target)
+
+        if self.ave_spectrum:
+            pred_freq = torch.mean(pred_freq, 0, keepdim=True)
+            target_freq = torch.mean(target_freq, 0, keepdim=True)
+
+        return self.loss_formulation(pred_freq, target_freq, matrix) * self.loss_weight

test_challenge_code/architecture/HSCNN_Plus.py

@@ -0,0 +1,77 @@
+import torch.nn as nn
+import torch
+class dfus_block(nn.Module):
+    def __init__(self, dim):
+        super(dfus_block, self).__init__()
+        self.conv1 = nn.Conv2d(dim, 128, 1, 1, 0, bias=False)
+
+        self.conv_up1 = nn.Conv2d(128, 32, 3, 1, 1, bias=False)
+        self.conv_up2 = nn.Conv2d(32, 16, 1, 1, 0, bias=False)
+
+        self.conv_down1 = nn.Conv2d(128, 32, 3, 1, 1, bias=False)
+        self.conv_down2 = nn.Conv2d(32, 16, 1, 1, 0, bias=False)
+
+        self.conv_fution = nn.Conv2d(96, 32, 1, 1, 0, bias=False)
+
+        #### activation function
+        self.relu = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        """
+        x: [b,c,h,w]
+        return out:[b,c,h,w]
+        """
+        feat = self.relu(self.conv1(x))
+        feat_up1 = self.relu(self.conv_up1(feat))
+        feat_up2 = self.relu(self.conv_up2(feat_up1))
+        feat_down1 = self.relu(self.conv_down1(feat))
+        feat_down2 = self.relu(self.conv_down2(feat_down1))
+        feat_fution = torch.cat([feat_up1,feat_up2,feat_down1,feat_down2],dim=1)
+        feat_fution = self.relu(self.conv_fution(feat_fution))
+        out = torch.cat([x, feat_fution], dim=1)
+        return out
+
+class ddfn(nn.Module):
+    def __init__(self, dim, num_blocks=78):
+        super(ddfn, self).__init__()
+
+        self.conv_up1 = nn.Conv2d(dim, 32, 3, 1, 1, bias=False)
+        self.conv_up2 = nn.Conv2d(32, 32, 1, 1, 0, bias=False)
+
+        self.conv_down1 = nn.Conv2d(dim, 32, 3, 1, 1, bias=False)
+        self.conv_down2 = nn.Conv2d(32, 32, 1, 1, 0, bias=False)
+
+        dfus_blocks = [dfus_block(dim=128+32*i) for i in range(num_blocks)]
+        self.dfus_blocks = nn.Sequential(*dfus_blocks)
+
+        #### activation function
+        self.relu = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        """
+        x: [b,c,h,w]
+        return out:[b,c,h,w]
+        """
+        feat_up1 = self.relu(self.conv_up1(x))
+        feat_up2 = self.relu(self.conv_up2(feat_up1))
+        feat_down1 = self.relu(self.conv_down1(x))
+        feat_down2 = self.relu(self.conv_down2(feat_down1))
+        feat_fution = torch.cat([feat_up1,feat_up2,feat_down1,feat_down2],dim=1)
+        out = self.dfus_blocks(feat_fution)
+        return out
+
+class HSCNN_Plus(nn.Module):
+    def __init__(self, in_channels=3, out_channels=31, num_blocks=30):
+        super(HSCNN_Plus, self).__init__()
+
+        self.ddfn = ddfn(dim=in_channels, num_blocks=num_blocks)
+        self.conv_out = nn.Conv2d(128+32*num_blocks, out_channels, 1, 1, 0, bias=False)
+
+    def forward(self, x):
+        """
+        x: [b,c,h,w]
+        return out:[b,c,h,w]
+        """
+        fea = self.ddfn(x)
+        out =  self.conv_out(fea)
+        return out

test_challenge_code/architecture/MIRNet.py

@@ -0,0 +1,416 @@
+# --- Imports --- #
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+# from pdb import set_trace as stx
+
+def get_pad_layer(pad_type):
+    if(pad_type in ['refl','reflect']):
+        PadLayer = nn.ReflectionPad2d
+    elif(pad_type in ['repl','replicate']):
+        PadLayer = nn.ReplicationPad2d
+    elif(pad_type=='zero'):
+        PadLayer = nn.ZeroPad2d
+    else:
+        print('Pad type [%s] not recognized'%pad_type)
+    return PadLayer
+
+class downsamp(nn.Module):
+    def __init__(self, pad_type='reflect', filt_size=3, stride=2, channels=None, pad_off=0):
+        super(downsamp, self).__init__()
+        self.filt_size = filt_size
+        self.pad_off = pad_off
+        self.pad_sizes = [int(1.*(filt_size-1)/2), int(np.ceil(1.*(filt_size-1)/2)), int(1.*(filt_size-1)/2), int(np.ceil(1.*(filt_size-1)/2))]
+        self.pad_sizes = [pad_size+pad_off for pad_size in self.pad_sizes]
+        self.stride = stride
+        self.off = int((self.stride-1)/2.)
+        self.channels = channels
+
+        # print('Filter size [%i]'%filt_size)
+        if(self.filt_size==1):
+            a = np.array([1.,])
+        elif(self.filt_size==2):
+            a = np.array([1., 1.])
+        elif(self.filt_size==3):
+            a = np.array([1., 2., 1.])
+        elif(self.filt_size==4):
+            a = np.array([1., 3., 3., 1.])
+        elif(self.filt_size==5):
+            a = np.array([1., 4., 6., 4., 1.])
+        elif(self.filt_size==6):
+            a = np.array([1., 5., 10., 10., 5., 1.])
+        elif(self.filt_size==7):
+            a = np.array([1., 6., 15., 20., 15., 6., 1.])
+
+        filt = torch.Tensor(a[:,None]*a[None,:])
+        filt = filt/torch.sum(filt)
+        self.register_buffer('filt', filt[None,None,:,:].repeat((self.channels,1,1,1)))
+
+        self.pad = get_pad_layer(pad_type)(self.pad_sizes)
+
+    def forward(self, inp):
+        if(self.filt_size==1):
+            if(self.pad_off==0):
+                return inp[:,:,::self.stride,::self.stride]
+            else:
+                return self.pad(inp)[:,:,::self.stride,::self.stride]
+        else:
+            return F.conv2d(self.pad(inp), self.filt, stride=self.stride, groups=inp.shape[1])
+
+##########################################################################
+
+def conv(in_channels, out_channels, kernel_size, bias=False, padding=1, stride=1):
+    return nn.Conv2d(
+        in_channels, out_channels, kernel_size,
+        padding=(kernel_size // 2), bias=bias, stride=stride)
+
+
+##########################################################################
+##---------- Selective Kernel Feature Fusion (SKFF) ----------
+class SKFF(nn.Module):
+    def __init__(self, in_channels, height=3, reduction=8, bias=False):
+        super(SKFF, self).__init__()
+
+        self.height = height
+        d = max(int(in_channels / reduction), 4)
+
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.conv_du = nn.Sequential(nn.Conv2d(in_channels, d, 1, padding=0, bias=bias), nn.PReLU())
+
+        self.fcs = nn.ModuleList([])
+        for i in range(self.height):
+            self.fcs.append(nn.Conv2d(d, in_channels, kernel_size=1, stride=1, bias=bias))
+
+        self.softmax = nn.Softmax(dim=1)
+
+    def forward(self, inp_feats):
+        batch_size = inp_feats[0].shape[0]
+        n_feats = inp_feats[0].shape[1]
+
+        inp_feats = torch.cat(inp_feats, dim=1)
+        inp_feats = inp_feats.view(batch_size, self.height, n_feats, inp_feats.shape[2], inp_feats.shape[3])
+
+        feats_U = torch.sum(inp_feats, dim=1)
+        feats_S = self.avg_pool(feats_U)
+        feats_Z = self.conv_du(feats_S)
+
+        attention_vectors = [fc(feats_Z) for fc in self.fcs]
+        attention_vectors = torch.cat(attention_vectors, dim=1)
+        attention_vectors = attention_vectors.view(batch_size, self.height, n_feats, 1, 1)
+        # stx()
+        attention_vectors = self.softmax(attention_vectors)
+
+        feats_V = torch.sum(inp_feats * attention_vectors, dim=1)
+
+        return feats_V
+
+    ##########################################################################
+
+
+##---------- Spatial Attention ----------
+class BasicConv(nn.Module):
+    def __init__(self, in_planes, out_planes, kernel_size, stride=1, padding=0, dilation=1, groups=1, relu=True,
+                 bn=False, bias=False):
+        super(BasicConv, self).__init__()
+        self.out_channels = out_planes
+        self.conv = nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=padding,
+                              dilation=dilation, groups=groups, bias=bias)
+        self.bn = nn.BatchNorm2d(out_planes, eps=1e-5, momentum=0.01, affine=True) if bn else None
+        self.relu = nn.ReLU() if relu else None
+
+    def forward(self, x):
+        x = self.conv(x)
+        if self.bn is not None:
+            x = self.bn(x)
+        if self.relu is not None:
+            x = self.relu(x)
+        return x
+
+
+class ChannelPool(nn.Module):
+    def forward(self, x):
+        return torch.cat((torch.max(x, 1)[0].unsqueeze(1), torch.mean(x, 1).unsqueeze(1)), dim=1)
+
+
+class spatial_attn_layer(nn.Module):
+    def __init__(self, kernel_size=5):
+        super(spatial_attn_layer, self).__init__()
+        self.compress = ChannelPool()
+        self.spatial = BasicConv(2, 1, kernel_size, stride=1, padding=(kernel_size - 1) // 2, relu=False)
+
+    def forward(self, x):
+        # import pdb;pdb.set_trace()
+        x_compress = self.compress(x)
+        x_out = self.spatial(x_compress)
+        scale = torch.sigmoid(x_out)  # broadcasting
+        return x * scale
+
+
+##########################################################################
+## ------ Channel Attention --------------
+class ca_layer(nn.Module):
+    def __init__(self, channel, reduction=8, bias=True):
+        super(ca_layer, self).__init__()
+        # global average pooling: feature --> point
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        # feature channel downscale and upscale --> channel weight
+        self.conv_du = nn.Sequential(
+            nn.Conv2d(channel, channel // reduction, 1, padding=0, bias=bias),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(channel // reduction, channel, 1, padding=0, bias=bias),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        y = self.avg_pool(x)
+        y = self.conv_du(y)
+        return x * y
+
+
+##########################################################################
+##---------- Dual Attention Unit (DAU) ----------
+class DAU(nn.Module):
+    def __init__(
+            self, n_feat, kernel_size=3, reduction=8,
+            bias=False, bn=False, act=nn.PReLU(), res_scale=1):
+        super(DAU, self).__init__()
+        modules_body = [conv(n_feat, n_feat, kernel_size, bias=bias), act, conv(n_feat, n_feat, kernel_size, bias=bias)]
+        self.body = nn.Sequential(*modules_body)
+
+        ## Spatial Attention
+        self.SA = spatial_attn_layer()
+
+        ## Channel Attention
+        self.CA = ca_layer(n_feat, reduction, bias=bias)
+
+        self.conv1x1 = nn.Conv2d(n_feat * 2, n_feat, kernel_size=1, bias=bias)
+
+    def forward(self, x):
+        res = self.body(x)
+        sa_branch = self.SA(res)
+        ca_branch = self.CA(res)
+        res = torch.cat([sa_branch, ca_branch], dim=1)
+        res = self.conv1x1(res)
+        res += x
+        return res
+
+
+##########################################################################
+##---------- Resizing Modules ----------
+class ResidualDownSample(nn.Module):
+    def __init__(self, in_channels, bias=False):
+        super(ResidualDownSample, self).__init__()
+
+        self.top = nn.Sequential(nn.Conv2d(in_channels, in_channels, 1, stride=1, padding=0, bias=bias),
+                                 nn.PReLU(),
+                                 nn.Conv2d(in_channels, in_channels, 3, stride=1, padding=1, bias=bias),
+                                 nn.PReLU(),
+                                 downsamp(channels=in_channels, filt_size=3, stride=2),
+                                 nn.Conv2d(in_channels, in_channels * 2, 1, stride=1, padding=0, bias=bias))
+
+        self.bot = nn.Sequential(downsamp(channels=in_channels, filt_size=3, stride=2),
+                                 nn.Conv2d(in_channels, in_channels * 2, 1, stride=1, padding=0, bias=bias))
+
+    def forward(self, x):
+        top = self.top(x)
+        bot = self.bot(x)
+        out = top + bot
+        return out
+
+
+class DownSample(nn.Module):
+    def __init__(self, in_channels, scale_factor, stride=2, kernel_size=3):
+        super(DownSample, self).__init__()
+        self.scale_factor = int(np.log2(scale_factor))
+
+        modules_body = []
+        for i in range(self.scale_factor):
+            modules_body.append(ResidualDownSample(in_channels))
+            in_channels = int(in_channels * stride)
+
+        self.body = nn.Sequential(*modules_body)
+
+    def forward(self, x):
+        x = self.body(x)
+        return x
+
+
+class ResidualUpSample(nn.Module):
+    def __init__(self, in_channels, bias=False):
+        super(ResidualUpSample, self).__init__()
+
+        self.top = nn.Sequential(nn.Conv2d(in_channels, in_channels, 1, stride=1, padding=0, bias=bias),
+                                 nn.PReLU(),
+                                 nn.ConvTranspose2d(in_channels, in_channels, 3, stride=2, padding=1, output_padding=1,
+                                                    bias=bias),
+                                 nn.PReLU(),
+                                 nn.Conv2d(in_channels, in_channels // 2, 1, stride=1, padding=0, bias=bias))
+
+        self.bot = nn.Sequential(nn.Upsample(scale_factor=2, mode='bilinear', align_corners=bias),
+                                 nn.Conv2d(in_channels, in_channels // 2, 1, stride=1, padding=0, bias=bias))
+
+    def forward(self, x):
+        top = self.top(x)
+        bot = self.bot(x)
+        out = top + bot
+        return out
+
+
+class UpSample(nn.Module):
+    def __init__(self, in_channels, scale_factor, stride=2, kernel_size=3):
+        super(UpSample, self).__init__()
+        self.scale_factor = int(np.log2(scale_factor))
+
+        modules_body = []
+        for i in range(self.scale_factor):
+            modules_body.append(ResidualUpSample(in_channels))
+            in_channels = int(in_channels // stride)
+
+        self.body = nn.Sequential(*modules_body)
+
+    def forward(self, x):
+        x = self.body(x)
+        return x
+
+
+##########################################################################
+##---------- Multi-Scale Resiudal Block (MSRB) ----------
+class MSRB(nn.Module):
+    def __init__(self, n_feat, height, width, stride, bias):
+        super(MSRB, self).__init__()
+
+        self.n_feat, self.height, self.width = n_feat, height, width
+        self.blocks = nn.ModuleList([nn.ModuleList([DAU(int(n_feat * stride ** i))] * width) for i in range(height)])
+
+        INDEX = np.arange(0, width, 2)
+        FEATS = [int((stride ** i) * n_feat) for i in range(height)]
+        SCALE = [2 ** i for i in range(1, height)]
+
+        self.last_up = nn.ModuleDict()
+        for i in range(1, height):
+            self.last_up.update({f'{i}': UpSample(int(n_feat * stride ** i), 2 ** i, stride)})
+
+        self.down = nn.ModuleDict()
+        self.up = nn.ModuleDict()
+
+        i = 0
+        SCALE.reverse()
+        for feat in FEATS:
+            for scale in SCALE[i:]:
+                self.down.update({f'{feat}_{scale}': DownSample(feat, scale, stride)})
+            i += 1
+
+        i = 0
+        FEATS.reverse()
+        for feat in FEATS:
+            for scale in SCALE[i:]:
+                self.up.update({f'{feat}_{scale}': UpSample(feat, scale, stride)})
+            i += 1
+
+        self.conv_out = nn.Conv2d(n_feat, n_feat, kernel_size=3, padding=1, bias=bias)
+
+        self.selective_kernel = nn.ModuleList([SKFF(n_feat * stride ** i, height) for i in range(height)])
+
+    def forward(self, x):
+        inp = x.clone()
+        # col 1 only
+        blocks_out = []
+        for j in range(self.height):
+            if j == 0:
+                inp = self.blocks[j][0](inp)
+            else:
+                inp = self.blocks[j][0](self.down[f'{inp.size(1)}_{2}'](inp))
+            blocks_out.append(inp)
+
+        # rest of grid
+        for i in range(1, self.width):
+            # Mesh
+            # Replace condition(i%2!=0) with True(Mesh) or False(Plain)
+            # if i%2!=0:
+            if True:
+                tmp = []
+                for j in range(self.height):
+                    TENSOR = []
+                    nfeats = (2 ** j) * self.n_feat
+                    for k in range(self.height):
+                        TENSOR.append(self.select_up_down(blocks_out[k], j, k))
+
+                    selective_kernel_fusion = self.selective_kernel[j](TENSOR)
+                    tmp.append(selective_kernel_fusion)
+            # Plain
+            else:
+                tmp = blocks_out
+            # Forward through either mesh or plain
+            for j in range(self.height):
+                blocks_out[j] = self.blocks[j][i](tmp[j])
+
+        # Sum after grid
+        out = []
+        for k in range(self.height):
+            out.append(self.select_last_up(blocks_out[k], k))
+
+        out = self.selective_kernel[0](out)
+
+        out = self.conv_out(out)
+        out = out + x
+
+        return out
+
+    def select_up_down(self, tensor, j, k):
+        if j == k:
+            return tensor
+        else:
+            diff = 2 ** np.abs(j - k)
+            if j < k:
+                return self.up[f'{tensor.size(1)}_{diff}'](tensor)
+            else:
+                return self.down[f'{tensor.size(1)}_{diff}'](tensor)
+
+    def select_last_up(self, tensor, k):
+        if k == 0:
+            return tensor
+        else:
+            return self.last_up[f'{k}'](tensor)
+
+
+##########################################################################
+##---------- Recursive Residual Group (RRG) ----------
+class RRG(nn.Module):
+    def __init__(self, n_feat, n_MSRB, height, width, stride, bias=False):
+        super(RRG, self).__init__()
+        modules_body = [MSRB(n_feat, height, width, stride, bias) for _ in range(n_MSRB)]
+        modules_body.append(conv(n_feat, n_feat, kernel_size=3))
+        self.body = nn.Sequential(*modules_body)
+
+    def forward(self, x):
+        res = self.body(x)
+        res += x
+        return res
+
+##########################################################################
+##---------- MIRNet  -----------------------
+class MIRNet(nn.Module):
+    def __init__(self, in_channels=3, out_channels=31, n_feat=31, kernel_size=3, stride=2, n_RRG=2, n_MSRB=1, height=3,
+                 width=1, bias=False):
+        super(MIRNet, self).__init__()
+
+        self.conv_in = nn.Conv2d(in_channels, n_feat, kernel_size=kernel_size, padding=(kernel_size - 1) // 2,
+                                 bias=bias)
+
+        modules_body = [RRG(n_feat, n_MSRB, height, width, stride, bias) for _ in range(n_RRG)]
+        self.body = nn.Sequential(*modules_body)
+
+        self.conv_out = nn.Conv2d(n_feat, out_channels, kernel_size=kernel_size, padding=(kernel_size - 1) // 2,
+                                  bias=bias)
+    def forward(self, x):
+        b, c, h_inp, w_inp = x.shape
+        hb, wb = 8, 8
+        pad_h = (hb - h_inp % hb) % hb
+        pad_w = (wb - w_inp % wb) % wb
+        x = F.pad(x, [0, pad_w, 0, pad_h], mode='reflect')
+        x = self.conv_in(x)
+        h = self.body(x)
+        h = self.conv_out(h)
+        h += x
+        return h[:, :, :h_inp, :w_inp]

test_challenge_code/architecture/MPRNet.py

@@ -0,0 +1,350 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+##########################################################################
+def conv(in_channels, out_channels, kernel_size, bias=False, stride = 1):
+    return nn.Conv2d(
+        in_channels, out_channels, kernel_size,
+        padding=(kernel_size//2), bias=bias, stride = stride)
+
+
+##########################################################################
+## Channel Attention Layer
+class CALayer(nn.Module):
+    def __init__(self, channel, reduction=16, bias=False):
+        super(CALayer, self).__init__()
+        # global average pooling: feature --> point
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        # feature channel downscale and upscale --> channel weight
+        self.conv_du = nn.Sequential(
+                nn.Conv2d(channel, channel // reduction, 1, padding=0, bias=bias),
+                nn.ReLU(inplace=True),
+                nn.Conv2d(channel // reduction, channel, 1, padding=0, bias=bias),
+                nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        y = self.avg_pool(x)
+        y = self.conv_du(y)
+        return x * y
+
+
+##########################################################################
+## Channel Attention Block (CAB)
+class CAB(nn.Module):
+    def __init__(self, n_feat, kernel_size, reduction, bias, act):
+        super(CAB, self).__init__()
+        modules_body = []
+        modules_body.append(conv(n_feat, n_feat, kernel_size, bias=bias))
+        modules_body.append(act)
+        modules_body.append(conv(n_feat, n_feat, kernel_size, bias=bias))
+
+        self.CA = CALayer(n_feat, reduction, bias=bias)
+        self.body = nn.Sequential(*modules_body)
+
+    def forward(self, x):
+        res = self.body(x)
+        res = self.CA(res)
+        res += x
+        return res
+
+##########################################################################
+## Supervised Attention Module
+class SAM(nn.Module):
+    def __init__(self, n_feat, kernel_size, bias):
+        super(SAM, self).__init__()
+        self.conv1 = conv(n_feat, n_feat, kernel_size, bias=bias)
+        self.conv2 = conv(n_feat, 31, kernel_size, bias=bias)
+        self.conv3 = conv(31, n_feat, kernel_size, bias=bias)
+
+    def forward(self, x, x_img):
+        x1 = self.conv1(x)
+        img = self.conv2(x) + x_img
+        x2 = torch.sigmoid(self.conv3(img))
+        x1 = x1*x2
+        x1 = x1+x
+        return x1, img
+
+##########################################################################
+## U-Net
+
+class Encoder(nn.Module):
+    def __init__(self, n_feat, kernel_size, reduction, act, bias, scale_unetfeats, csff):
+        super(Encoder, self).__init__()
+
+        self.encoder_level1 = [CAB(n_feat,                     kernel_size, reduction, bias=bias, act=act) for _ in range(2)]
+        self.encoder_level2 = [CAB(n_feat+scale_unetfeats,     kernel_size, reduction, bias=bias, act=act) for _ in range(2)]
+        self.encoder_level3 = [CAB(n_feat+(scale_unetfeats*2), kernel_size, reduction, bias=bias, act=act) for _ in range(2)]
+
+        self.encoder_level1 = nn.Sequential(*self.encoder_level1)
+        self.encoder_level2 = nn.Sequential(*self.encoder_level2)
+        self.encoder_level3 = nn.Sequential(*self.encoder_level3)
+
+        self.down12  = DownSample(n_feat, scale_unetfeats)
+        self.down23  = DownSample(n_feat+scale_unetfeats, scale_unetfeats)
+
+        # Cross Stage Feature Fusion (CSFF)
+        if csff:
+            self.csff_enc1 = nn.Conv2d(n_feat,                     n_feat,                     kernel_size=1, bias=bias)
+            self.csff_enc2 = nn.Conv2d(n_feat+scale_unetfeats,     n_feat+scale_unetfeats,     kernel_size=1, bias=bias)
+            self.csff_enc3 = nn.Conv2d(n_feat+(scale_unetfeats*2), n_feat+(scale_unetfeats*2), kernel_size=1, bias=bias)
+
+            self.csff_dec1 = nn.Conv2d(n_feat,                     n_feat,                     kernel_size=1, bias=bias)
+            self.csff_dec2 = nn.Conv2d(n_feat+scale_unetfeats,     n_feat+scale_unetfeats,     kernel_size=1, bias=bias)
+            self.csff_dec3 = nn.Conv2d(n_feat+(scale_unetfeats*2), n_feat+(scale_unetfeats*2), kernel_size=1, bias=bias)
+
+    def forward(self, x, encoder_outs=None, decoder_outs=None):
+        enc1 = self.encoder_level1(x)
+        if (encoder_outs is not None) and (decoder_outs is not None):
+            enc1 = enc1 + self.csff_enc1(encoder_outs[0]) + self.csff_dec1(decoder_outs[0])
+
+        x = self.down12(enc1)
+
+        enc2 = self.encoder_level2(x)
+        if (encoder_outs is not None) and (decoder_outs is not None):
+            enc2 = enc2 + self.csff_enc2(encoder_outs[1]) + self.csff_dec2(decoder_outs[1])
+
+        x = self.down23(enc2)
+
+        enc3 = self.encoder_level3(x)
+        if (encoder_outs is not None) and (decoder_outs is not None):
+            enc3 = enc3 + self.csff_enc3(encoder_outs[2]) + self.csff_dec3(decoder_outs[2])
+
+        return [enc1, enc2, enc3]
+
+class Decoder(nn.Module):
+    def __init__(self, n_feat, kernel_size, reduction, act, bias, scale_unetfeats):
+        super(Decoder, self).__init__()
+
+        self.decoder_level1 = [CAB(n_feat,                     kernel_size, reduction, bias=bias, act=act) for _ in range(2)]
+        self.decoder_level2 = [CAB(n_feat+scale_unetfeats,     kernel_size, reduction, bias=bias, act=act) for _ in range(2)]
+        self.decoder_level3 = [CAB(n_feat+(scale_unetfeats*2), kernel_size, reduction, bias=bias, act=act) for _ in range(2)]
+
+        self.decoder_level1 = nn.Sequential(*self.decoder_level1)
+        self.decoder_level2 = nn.Sequential(*self.decoder_level2)
+        self.decoder_level3 = nn.Sequential(*self.decoder_level3)
+
+        self.skip_attn1 = CAB(n_feat,                 kernel_size, reduction, bias=bias, act=act)
+        self.skip_attn2 = CAB(n_feat+scale_unetfeats, kernel_size, reduction, bias=bias, act=act)
+
+        self.up21  = SkipUpSample(n_feat, scale_unetfeats)
+        self.up32  = SkipUpSample(n_feat+scale_unetfeats, scale_unetfeats)
+
+    def forward(self, outs):
+        enc1, enc2, enc3 = outs
+        dec3 = self.decoder_level3(enc3)
+
+        x = self.up32(dec3, self.skip_attn2(enc2))
+        dec2 = self.decoder_level2(x)
+
+        x = self.up21(dec2, self.skip_attn1(enc1))
+        dec1 = self.decoder_level1(x)
+
+        return [dec1,dec2,dec3]
+
+##########################################################################
+##---------- Resizing Modules ----------
+class DownSample(nn.Module):
+    def __init__(self, in_channels,s_factor):
+        super(DownSample, self).__init__()
+        self.down = nn.Sequential(nn.Upsample(scale_factor=0.5, mode='bilinear', align_corners=False),
+                                  nn.Conv2d(in_channels, in_channels+s_factor, 1, stride=1, padding=0, bias=False))
+
+    def forward(self, x):
+        x = self.down(x)
+        return x
+
+class UpSample(nn.Module):
+    def __init__(self, in_channels,s_factor):
+        super(UpSample, self).__init__()
+        self.up = nn.Sequential(nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
+                                nn.Conv2d(in_channels+s_factor, in_channels, 1, stride=1, padding=0, bias=False))
+
+    def forward(self, x):
+        x = self.up(x)
+        return x
+
+class SkipUpSample(nn.Module):
+    def __init__(self, in_channels,s_factor):
+        super(SkipUpSample, self).__init__()
+        self.up = nn.Sequential(nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
+                                nn.Conv2d(in_channels+s_factor, in_channels, 1, stride=1, padding=0, bias=False))
+
+    def forward(self, x, y):
+        x = self.up(x)
+        x = x + y
+        return x
+
+##########################################################################
+## Original Resolution Block (ORB)
+class ORB(nn.Module):
+    def __init__(self, n_feat, kernel_size, reduction, act, bias, num_cab):
+        super(ORB, self).__init__()
+        modules_body = []
+        modules_body = [CAB(n_feat, kernel_size, reduction, bias=bias, act=act) for _ in range(num_cab)]
+        modules_body.append(conv(n_feat, n_feat, kernel_size))
+        self.body = nn.Sequential(*modules_body)
+
+    def forward(self, x):
+        res = self.body(x)
+        res += x
+        return res
+
+##########################################################################
+class ORSNet(nn.Module):
+    def __init__(self, n_feat, scale_orsnetfeats, kernel_size, reduction, act, bias, scale_unetfeats, num_cab):
+        super(ORSNet, self).__init__()
+
+        self.orb1 = ORB(n_feat+scale_orsnetfeats, kernel_size, reduction, act, bias, num_cab)
+        self.orb2 = ORB(n_feat+scale_orsnetfeats, kernel_size, reduction, act, bias, num_cab)
+        self.orb3 = ORB(n_feat+scale_orsnetfeats, kernel_size, reduction, act, bias, num_cab)
+
+        self.up_enc1 = UpSample(n_feat, scale_unetfeats)
+        self.up_dec1 = UpSample(n_feat, scale_unetfeats)
+
+        self.up_enc2 = nn.Sequential(UpSample(n_feat+scale_unetfeats, scale_unetfeats), UpSample(n_feat, scale_unetfeats))
+        self.up_dec2 = nn.Sequential(UpSample(n_feat+scale_unetfeats, scale_unetfeats), UpSample(n_feat, scale_unetfeats))
+
+        self.conv_enc1 = nn.Conv2d(n_feat, n_feat+scale_orsnetfeats, kernel_size=1, bias=bias)
+        self.conv_enc2 = nn.Conv2d(n_feat, n_feat+scale_orsnetfeats, kernel_size=1, bias=bias)
+        self.conv_enc3 = nn.Conv2d(n_feat, n_feat+scale_orsnetfeats, kernel_size=1, bias=bias)
+
+        self.conv_dec1 = nn.Conv2d(n_feat, n_feat+scale_orsnetfeats, kernel_size=1, bias=bias)
+        self.conv_dec2 = nn.Conv2d(n_feat, n_feat+scale_orsnetfeats, kernel_size=1, bias=bias)
+        self.conv_dec3 = nn.Conv2d(n_feat, n_feat+scale_orsnetfeats, kernel_size=1, bias=bias)
+
+    def forward(self, x, encoder_outs, decoder_outs):
+        x = self.orb1(x)
+        x = x + self.conv_enc1(encoder_outs[0]) + self.conv_dec1(decoder_outs[0])
+
+        x = self.orb2(x)
+        x = x + self.conv_enc2(self.up_enc1(encoder_outs[1])) + self.conv_dec2(self.up_dec1(decoder_outs[1]))
+
+        x = self.orb3(x)
+        x = x + self.conv_enc3(self.up_enc2(encoder_outs[2])) + self.conv_dec3(self.up_dec2(decoder_outs[2]))
+
+        return x
+
+
+##########################################################################
+class MPRNet(nn.Module):
+    def __init__(self, in_c=31, out_c=31, n_feat=31, scale_unetfeats=31, scale_orsnetfeats=31, num_cab=4, kernel_size=3, reduction=1, bias=False):
+        super(MPRNet, self).__init__()
+
+        self.conv_in = nn.Conv2d(3, in_c, kernel_size=kernel_size, padding=(kernel_size - 1) // 2,
+                                 bias=bias)
+
+        act=nn.PReLU()
+        self.shallow_feat1 = nn.Sequential(conv(in_c, n_feat, kernel_size, bias=bias), CAB(n_feat,kernel_size, reduction, bias=bias, act=act))
+        self.shallow_feat2 = nn.Sequential(conv(in_c, n_feat, kernel_size, bias=bias), CAB(n_feat,kernel_size, reduction, bias=bias, act=act))
+        self.shallow_feat3 = nn.Sequential(conv(in_c, n_feat, kernel_size, bias=bias), CAB(n_feat,kernel_size, reduction, bias=bias, act=act))
+
+        # Cross Stage Feature Fusion (CSFF)
+        self.stage1_encoder = Encoder(n_feat, kernel_size, reduction, act, bias, scale_unetfeats, csff=False)
+        self.stage1_decoder = Decoder(n_feat, kernel_size, reduction, act, bias, scale_unetfeats)
+
+        self.stage2_encoder = Encoder(n_feat, kernel_size, reduction, act, bias, scale_unetfeats, csff=True)
+        self.stage2_decoder = Decoder(n_feat, kernel_size, reduction, act, bias, scale_unetfeats)
+
+        self.stage3_orsnet = ORSNet(n_feat, scale_orsnetfeats, kernel_size, reduction, act, bias, scale_unetfeats, num_cab)
+
+        self.sam12 = SAM(n_feat, kernel_size=1, bias=bias)
+        self.sam23 = SAM(n_feat, kernel_size=1, bias=bias)
+
+        self.concat12  = conv(n_feat*2, n_feat, kernel_size, bias=bias)
+        self.concat23  = conv(n_feat*2, n_feat+scale_orsnetfeats, kernel_size, bias=bias)
+        self.tail     = conv(n_feat+scale_orsnetfeats, out_c, kernel_size, bias=bias)
+
+    def forward(self, x3_img):
+        b, c, h_inp, w_inp = x3_img.shape
+        hb, wb = 8, 8
+        pad_h = (hb - h_inp % hb) % hb
+        pad_w = (wb - w_inp % wb) % wb
+        x3_img = F.pad(x3_img, [0, pad_w, 0, pad_h], mode='reflect')
+        x3_img = self.conv_in(x3_img)
+
+        # Original-resolution Image for Stage 3
+        H = x3_img.size(2)
+        W = x3_img.size(3)
+
+        # Multi-Patch Hierarchy: Split Image into four non-overlapping patches
+
+        # Two Patches for Stage 2
+        x2top_img  = x3_img[:,:,0:int(H/2),:]
+        x2bot_img  = x3_img[:,:,int(H/2):H,:]
+
+        # Four Patches for Stage 1
+        x1ltop_img = x2top_img[:,:,:,0:int(W/2)]
+        x1rtop_img = x2top_img[:,:,:,int(W/2):W]
+        x1lbot_img = x2bot_img[:,:,:,0:int(W/2)]
+        x1rbot_img = x2bot_img[:,:,:,int(W/2):W]
+
+        ##-------------------------------------------
+        ##-------------- Stage 1---------------------
+        ##-------------------------------------------
+        ## Compute Shallow Features
+        x1ltop = self.shallow_feat1(x1ltop_img)
+        x1rtop = self.shallow_feat1(x1rtop_img)
+        x1lbot = self.shallow_feat1(x1lbot_img)
+        x1rbot = self.shallow_feat1(x1rbot_img)
+
+        ## Process features of all 4 patches with Encoder of Stage 1
+        feat1_ltop = self.stage1_encoder(x1ltop)
+        feat1_rtop = self.stage1_encoder(x1rtop)
+        feat1_lbot = self.stage1_encoder(x1lbot)
+        feat1_rbot = self.stage1_encoder(x1rbot)
+
+        ## Concat deep features
+        feat1_top = [torch.cat((k,v), 3) for k,v in zip(feat1_ltop,feat1_rtop)]
+        feat1_bot = [torch.cat((k,v), 3) for k,v in zip(feat1_lbot,feat1_rbot)]
+
+        ## Pass features through Decoder of Stage 1
+        res1_top = self.stage1_decoder(feat1_top)
+        res1_bot = self.stage1_decoder(feat1_bot)
+
+        ## Apply Supervised Attention Module (SAM)
+        x2top_samfeats, stage1_img_top = self.sam12(res1_top[0], x2top_img)
+        x2bot_samfeats, stage1_img_bot = self.sam12(res1_bot[0], x2bot_img)
+
+        ## Output image at Stage 1
+        stage1_img = torch.cat([stage1_img_top, stage1_img_bot],2)
+        ##-------------------------------------------
+        ##-------------- Stage 2---------------------
+        ##-------------------------------------------
+        ## Compute Shallow Features
+        x2top  = self.shallow_feat2(x2top_img)
+        x2bot  = self.shallow_feat2(x2bot_img)
+
+        ## Concatenate SAM features of Stage 1 with shallow features of Stage 2
+        x2top_cat = self.concat12(torch.cat([x2top, x2top_samfeats], 1))
+        x2bot_cat = self.concat12(torch.cat([x2bot, x2bot_samfeats], 1))
+
+        ## Process features of both patches with Encoder of Stage 2
+        feat2_top = self.stage2_encoder(x2top_cat, feat1_top, res1_top)
+        feat2_bot = self.stage2_encoder(x2bot_cat, feat1_bot, res1_bot)
+
+        ## Concat deep features
+        feat2 = [torch.cat((k,v), 2) for k,v in zip(feat2_top,feat2_bot)]
+
+        ## Pass features through Decoder of Stage 2
+        res2 = self.stage2_decoder(feat2)
+
+        ## Apply SAM
+        x3_samfeats, stage2_img = self.sam23(res2[0], x3_img)
+
+
+        ##-------------------------------------------
+        ##-------------- Stage 3---------------------
+        ##-------------------------------------------
+        ## Compute Shallow Features
+        x3     = self.shallow_feat3(x3_img)
+
+        ## Concatenate SAM features of Stage 2 with shallow features of Stage 3
+        x3_cat = self.concat23(torch.cat([x3, x3_samfeats], 1))
+
+        x3_cat = self.stage3_orsnet(x3_cat, feat2, res2)
+
+        stage3_img = self.tail(x3_cat)
+
+        return (stage3_img + x3_img)[:, :, :h_inp, :w_inp]

test_challenge_code/architecture/MST.py

@@ -0,0 +1,313 @@
+import torch.nn as nn
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+import math
+import warnings
+from torch.nn.init import _calculate_fan_in_and_fan_out
+
+def _no_grad_trunc_normal_(tensor, mean, std, a, b):
+    def norm_cdf(x):
+        return (1. + math.erf(x / math.sqrt(2.))) / 2.
+
+    if (mean < a - 2 * std) or (mean > b + 2 * std):
+        warnings.warn("mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
+                      "The distribution of values may be incorrect.",
+                      stacklevel=2)
+    with torch.no_grad():
+        l = norm_cdf((a - mean) / std)
+        u = norm_cdf((b - mean) / std)
+        tensor.uniform_(2 * l - 1, 2 * u - 1)
+        tensor.erfinv_()
+        tensor.mul_(std * math.sqrt(2.))
+        tensor.add_(mean)
+        tensor.clamp_(min=a, max=b)
+        return tensor
+
+
+def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):
+    # type: (Tensor, float, float, float, float) -> Tensor
+    return _no_grad_trunc_normal_(tensor, mean, std, a, b)
+
+
+def variance_scaling_(tensor, scale=1.0, mode='fan_in', distribution='normal'):
+    fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor)
+    if mode == 'fan_in':
+        denom = fan_in
+    elif mode == 'fan_out':
+        denom = fan_out
+    elif mode == 'fan_avg':
+        denom = (fan_in + fan_out) / 2
+    variance = scale / denom
+    if distribution == "truncated_normal":
+        trunc_normal_(tensor, std=math.sqrt(variance) / .87962566103423978)
+    elif distribution == "normal":
+        tensor.normal_(std=math.sqrt(variance))
+    elif distribution == "uniform":
+        bound = math.sqrt(3 * variance)
+        tensor.uniform_(-bound, bound)
+    else:
+        raise ValueError(f"invalid distribution {distribution}")
+
+
+def lecun_normal_(tensor):
+    variance_scaling_(tensor, mode='fan_in', distribution='truncated_normal')
+
+
+class PreNorm(nn.Module):
+    def __init__(self, dim, fn):
+        super().__init__()
+        self.fn = fn
+        self.norm = nn.LayerNorm(dim)
+
+    def forward(self, x, *args, **kwargs):
+        x = self.norm(x)
+        return self.fn(x, *args, **kwargs)
+
+
+class GELU(nn.Module):
+    def forward(self, x):
+        return F.gelu(x)
+
+def conv(in_channels, out_channels, kernel_size, bias=False, padding = 1, stride = 1):
+    return nn.Conv2d(
+        in_channels, out_channels, kernel_size,
+        padding=(kernel_size//2), bias=bias, stride=stride)
+
+
+def shift_back(inputs,step=2):          # input [bs,28,256,310]  output [bs, 28, 256, 256]
+    [bs, nC, row, col] = inputs.shape
+    down_sample = 256//row
+    step = float(step)/float(down_sample*down_sample)
+    out_col = row
+    for i in range(nC):
+        inputs[:,i,:,:out_col] = \
+            inputs[:,i,:,int(step*i):int(step*i)+out_col]
+    return inputs[:, :, :, :out_col]
+
+class MaskGuidedMechanism(nn.Module):
+    def __init__(
+            self, n_feat):
+        super(MaskGuidedMechanism, self).__init__()
+
+        self.conv1 = nn.Conv2d(n_feat, n_feat, kernel_size=1, bias=True)
+        self.conv2 = nn.Conv2d(n_feat, n_feat, kernel_size=1, bias=True)
+        self.depth_conv = nn.Conv2d(n_feat, n_feat, kernel_size=5, padding=2, bias=True, groups=n_feat)
+
+    def forward(self, mask_shift):
+        # x: b,c,h,w
+        [bs, nC, row, col] = mask_shift.shape
+        mask_shift = self.conv1(mask_shift)
+        attn_map = torch.sigmoid(self.depth_conv(self.conv2(mask_shift)))
+        res = mask_shift * attn_map
+        mask_emb = res + mask_shift
+        return mask_emb
+
+class MS_MSA(nn.Module):
+    def __init__(
+            self,
+            dim,
+            dim_head,
+            heads,
+    ):
+        super().__init__()
+        self.num_heads = heads
+        self.dim_head = dim_head
+        self.to_q = nn.Linear(dim, dim_head * heads, bias=False)
+        self.to_k = nn.Linear(dim, dim_head * heads, bias=False)
+        self.to_v = nn.Linear(dim, dim_head * heads, bias=False)
+        self.rescale = nn.Parameter(torch.ones(heads, 1, 1))
+        self.proj = nn.Linear(dim_head * heads, dim, bias=True)
+        self.pos_emb = nn.Sequential(
+            nn.Conv2d(dim, dim, 3, 1, 1, bias=False, groups=dim),
+            GELU(),
+            nn.Conv2d(dim, dim, 3, 1, 1, bias=False, groups=dim),
+        )
+        self.mm = MaskGuidedMechanism(dim)
+        self.dim = dim
+
+    def forward(self, x_in, mask=None):
+        """
+        x_in: [b,h,w,c]
+        mask: [1,h,w,c]
+        return out: [b,h,w,c]
+        """
+        b, h, w, c = x_in.shape
+        x = x_in.reshape(b,h*w,c)
+        q_inp = self.to_q(x)
+        k_inp = self.to_k(x)
+        v_inp = self.to_v(x)
+        mask_attn = self.mm(mask.permute(0,3,1,2)).permute(0,2,3,1)
+        if b != 0:
+            mask_attn = (mask_attn[0, :, :, :]).expand([b, h, w, c])
+        q, k, v, mask_attn = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=self.num_heads),
+                                (q_inp, k_inp, v_inp, mask_attn.flatten(1, 2)))
+        v = v * mask_attn
+        # q: b,heads,hw,c
+        q = q.transpose(-2, -1)
+        k = k.transpose(-2, -1)
+        v = v.transpose(-2, -1)
+        q = F.normalize(q, dim=-1, p=2)
+        k = F.normalize(k, dim=-1, p=2)
+        attn = (k @ q.transpose(-2, -1))   # A = K^T*Q
+        attn = attn * self.rescale
+        attn = attn.softmax(dim=-1)
+        x = attn @ v   # b,heads,d,hw
+        x = x.permute(0, 3, 1, 2)    # Transpose
+        x = x.reshape(b, h * w, self.num_heads * self.dim_head)
+        out_c = self.proj(x).view(b, h, w, c)
+        out_p = self.pos_emb(v_inp.reshape(b,h,w,c).permute(0, 3, 1, 2)).permute(0, 2, 3, 1)
+        out = out_c + out_p
+
+        return out
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, mult=4):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Conv2d(dim, dim * mult, 1, 1, bias=False),
+            GELU(),
+            nn.Conv2d(dim * mult, dim * mult, 3, 1, 1, bias=False, groups=dim * mult),
+            GELU(),
+            nn.Conv2d(dim * mult, dim, 1, 1, bias=False),
+        )
+
+    def forward(self, x):
+        """
+        x: [b,h,w,c]
+        return out: [b,h,w,c]
+        """
+        out = self.net(x.permute(0, 3, 1, 2))
+        return out.permute(0, 2, 3, 1)
+
+class MSAB(nn.Module):
+    def __init__(
+            self,
+            dim,
+            dim_head,
+            heads,
+            num_blocks,
+    ):
+        super().__init__()
+        self.blocks = nn.ModuleList([])
+        for _ in range(num_blocks):
+            self.blocks.append(nn.ModuleList([
+                MS_MSA(dim=dim, dim_head=dim_head, heads=heads),
+                PreNorm(dim, FeedForward(dim=dim))
+            ]))
+
+    def forward(self, x, mask):
+        """
+        x: [b,c,h,w]
+        return out: [b,c,h,w]
+        """
+        x = x.permute(0, 2, 3, 1)
+        for (attn, ff) in self.blocks:
+            x = attn(x, mask=mask.permute(0, 2, 3, 1)) + x
+            x = ff(x) + x
+        out = x.permute(0, 3, 1, 2)
+        return out
+
+class MST(nn.Module):
+    def __init__(self, dim, stage, num_blocks):
+        super(MST, self).__init__()
+        self.dim = dim
+        self.stage = stage
+
+        # Input projection
+        self.embedding_1 = nn.Conv2d(3, self.dim, 3, 1, 1, bias=False)
+        self.embedding_2 = nn.Conv2d(3, self.dim, 3, 1, 1, bias=False)
+
+        # Encoder
+        self.encoder_layers = nn.ModuleList([])
+        dim_stage = dim
+        for i in range(stage):
+            self.encoder_layers.append(nn.ModuleList([
+                MSAB(
+                    dim=dim_stage, num_blocks=num_blocks[i], dim_head=dim, heads=dim_stage // dim),
+                nn.Conv2d(dim_stage, dim_stage * 2, 4, 2, 1, bias=False),
+                nn.Conv2d(dim_stage, dim_stage * 2, 4, 2, 1, bias=False)
+            ]))
+            dim_stage *= 2
+
+        # Bottleneck
+        self.bottleneck = MSAB(
+            dim=dim_stage, dim_head=dim, heads=dim_stage // dim, num_blocks=num_blocks[-1])
+
+        # Decoder
+        self.decoder_layers = nn.ModuleList([])
+        for i in range(stage):
+            self.decoder_layers.append(nn.ModuleList([
+                nn.ConvTranspose2d(dim_stage, dim_stage // 2, stride=2, kernel_size=2, padding=0, output_padding=0),
+                nn.Conv2d(dim_stage, dim_stage // 2, 1, 1, bias=False),
+                MSAB(
+                    dim=dim_stage // 2, num_blocks=num_blocks[stage - 1 - i], dim_head=dim,
+                    heads=(dim_stage // 2) // dim),
+            ]))
+            dim_stage //= 2
+
+        # Output projection
+        self.mapping = nn.Conv2d(self.dim, 31, 3, 1, 1, bias=False)
+
+        #### activation function
+        self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=True)
+
+    def forward(self, x):
+        """
+        x: [b,c,h,w]
+        return out:[b,c,h,w]
+        """
+        b, c, h_inp, w_inp = x.shape
+        hb, wb = 8, 8
+        pad_h = (hb - h_inp % hb) % hb
+        pad_w = (wb - w_inp % wb) % wb
+        x = F.pad(x, [0, pad_w, 0, pad_h], mode='reflect')
+
+        # Embedding
+        mask = self.lrelu(self.embedding_1(x))
+        x = self.lrelu(self.embedding_2(x))
+        fea = x
+
+        # Encoder
+        fea_encoder = []
+        masks = []
+        for (MSAB, FeaDownSample, MaskDownSample) in self.encoder_layers:
+            fea = MSAB(fea, mask)
+            masks.append(mask)
+            fea_encoder.append(fea)
+            fea = FeaDownSample(fea)
+            mask = MaskDownSample(mask)
+
+        # Bottleneck
+        fea = self.bottleneck(fea, mask)
+
+        # Decoder
+        for i, (FeaUpSample, Fution, LeWinBlcok) in enumerate(self.decoder_layers):
+            fea = FeaUpSample(fea)
+            fea = Fution(torch.cat([fea, fea_encoder[self.stage-1-i]], dim=1))
+            mask = masks[self.stage - 1 - i]
+            fea = LeWinBlcok(fea, mask)
+
+        # Mapping
+        out = self.mapping(fea) + x
+
+        return out[:, :, :h_inp, :w_inp]
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+

test_challenge_code/architecture/MST_Plus_Plus.py

@@ -0,0 +1,307 @@
+import torch.nn as nn
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+import math
+import warnings
+from torch.nn.init import _calculate_fan_in_and_fan_out
+
+def _no_grad_trunc_normal_(tensor, mean, std, a, b):
+    def norm_cdf(x):
+        return (1. + math.erf(x / math.sqrt(2.))) / 2.
+
+    if (mean < a - 2 * std) or (mean > b + 2 * std):
+        warnings.warn("mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
+                      "The distribution of values may be incorrect.",
+                      stacklevel=2)
+    with torch.no_grad():
+        l = norm_cdf((a - mean) / std)
+        u = norm_cdf((b - mean) / std)
+        tensor.uniform_(2 * l - 1, 2 * u - 1)
+        tensor.erfinv_()
+        tensor.mul_(std * math.sqrt(2.))
+        tensor.add_(mean)
+        tensor.clamp_(min=a, max=b)
+        return tensor
+
+
+def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):
+    # type: (Tensor, float, float, float, float) -> Tensor
+    return _no_grad_trunc_normal_(tensor, mean, std, a, b)
+
+
+def variance_scaling_(tensor, scale=1.0, mode='fan_in', distribution='normal'):
+    fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor)
+    if mode == 'fan_in':
+        denom = fan_in
+    elif mode == 'fan_out':
+        denom = fan_out
+    elif mode == 'fan_avg':
+        denom = (fan_in + fan_out) / 2
+    variance = scale / denom
+    if distribution == "truncated_normal":
+        trunc_normal_(tensor, std=math.sqrt(variance) / .87962566103423978)
+    elif distribution == "normal":
+        tensor.normal_(std=math.sqrt(variance))
+    elif distribution == "uniform":
+        bound = math.sqrt(3 * variance)
+        tensor.uniform_(-bound, bound)
+    else:
+        raise ValueError(f"invalid distribution {distribution}")
+
+
+def lecun_normal_(tensor):
+    variance_scaling_(tensor, mode='fan_in', distribution='truncated_normal')
+
+
+class PreNorm(nn.Module):
+    def __init__(self, dim, fn):
+        super().__init__()
+        self.fn = fn
+        self.norm = nn.LayerNorm(dim)
+
+    def forward(self, x, *args, **kwargs):
+        x = self.norm(x)
+        return self.fn(x, *args, **kwargs)
+
+
+class GELU(nn.Module):
+    def forward(self, x):
+        return F.gelu(x)
+
+def conv(in_channels, out_channels, kernel_size, bias=False, padding = 1, stride = 1):
+    return nn.Conv2d(
+        in_channels, out_channels, kernel_size,
+        padding=(kernel_size//2), bias=bias, stride=stride)
+
+
+def shift_back(inputs,step=2):          # input [bs,28,256,310]  output [bs, 28, 256, 256]
+    [bs, nC, row, col] = inputs.shape
+    down_sample = 256//row
+    step = float(step)/float(down_sample*down_sample)
+    out_col = row
+    for i in range(nC):
+        inputs[:,i,:,:out_col] = \
+            inputs[:,i,:,int(step*i):int(step*i)+out_col]
+    return inputs[:, :, :, :out_col]
+
+class MS_MSA(nn.Module):
+    def __init__(
+            self,
+            dim,
+            dim_head,
+            heads,
+    ):
+        super().__init__()
+        self.num_heads = heads
+        self.dim_head = dim_head
+        self.to_q = nn.Linear(dim, dim_head * heads, bias=False)
+        self.to_k = nn.Linear(dim, dim_head * heads, bias=False)
+        self.to_v = nn.Linear(dim, dim_head * heads, bias=False)
+        self.rescale = nn.Parameter(torch.ones(heads, 1, 1))
+        self.proj = nn.Linear(dim_head * heads, dim, bias=True)
+        self.pos_emb = nn.Sequential(
+            nn.Conv2d(dim, dim, 3, 1, 1, bias=False, groups=dim),
+            GELU(),
+            nn.Conv2d(dim, dim, 3, 1, 1, bias=False, groups=dim),
+        )
+        self.dim = dim
+
+    def forward(self, x_in):
+        """
+        x_in: [b,h,w,c]
+        return out: [b,h,w,c]
+        """
+        b, h, w, c = x_in.shape
+        x = x_in.reshape(b,h*w,c)
+        q_inp = self.to_q(x)
+        k_inp = self.to_k(x)
+        v_inp = self.to_v(x)
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=self.num_heads),
+                                (q_inp, k_inp, v_inp))
+        v = v
+        # q: b,heads,hw,c
+        q = q.transpose(-2, -1)
+        k = k.transpose(-2, -1)
+        v = v.transpose(-2, -1)
+        q = F.normalize(q, dim=-1, p=2)
+        k = F.normalize(k, dim=-1, p=2)
+        attn = (k @ q.transpose(-2, -1))   # A = K^T*Q
+        attn = attn * self.rescale
+        attn = attn.softmax(dim=-1)
+        x = attn @ v   # b,heads,d,hw
+        x = x.permute(0, 3, 1, 2)    # Transpose
+        x = x.reshape(b, h * w, self.num_heads * self.dim_head)
+        out_c = self.proj(x).view(b, h, w, c)
+        out_p = self.pos_emb(v_inp.reshape(b,h,w,c).permute(0, 3, 1, 2)).permute(0, 2, 3, 1)
+        out = out_c + out_p
+
+        return out
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, mult=4):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Conv2d(dim, dim * mult, 1, 1, bias=False),
+            GELU(),
+            nn.Conv2d(dim * mult, dim * mult, 3, 1, 1, bias=False, groups=dim * mult),
+            GELU(),
+            nn.Conv2d(dim * mult, dim, 1, 1, bias=False),
+        )
+
+    def forward(self, x):
+        """
+        x: [b,h,w,c]
+        return out: [b,h,w,c]
+        """
+        out = self.net(x.permute(0, 3, 1, 2))
+        return out.permute(0, 2, 3, 1)
+
+class MSAB(nn.Module):
+    def __init__(
+            self,
+            dim,
+            dim_head,
+            heads,
+            num_blocks,
+    ):
+        super().__init__()
+        self.blocks = nn.ModuleList([])
+        for _ in range(num_blocks):
+            self.blocks.append(nn.ModuleList([
+                MS_MSA(dim=dim, dim_head=dim_head, heads=heads),
+                PreNorm(dim, FeedForward(dim=dim))
+            ]))
+
+    def forward(self, x):
+        """
+        x: [b,c,h,w]
+        return out: [b,c,h,w]
+        """
+        x = x.permute(0, 2, 3, 1)
+        for (attn, ff) in self.blocks:
+            x = attn(x) + x
+            x = ff(x) + x
+        out = x.permute(0, 3, 1, 2)
+        return out
+
+class MST(nn.Module):
+    def __init__(self, in_dim=31, out_dim=31, dim=31, stage=2, num_blocks=[2,4,4]):
+        super(MST, self).__init__()
+        self.dim = dim
+        self.stage = stage
+
+        # Input projection
+        self.embedding = nn.Conv2d(in_dim, self.dim, 3, 1, 1, bias=False)
+
+        # Encoder
+        self.encoder_layers = nn.ModuleList([])
+        dim_stage = dim
+        for i in range(stage):
+            self.encoder_layers.append(nn.ModuleList([
+                MSAB(
+                    dim=dim_stage, num_blocks=num_blocks[i], dim_head=dim, heads=dim_stage // dim),
+                nn.Conv2d(dim_stage, dim_stage * 2, 4, 2, 1, bias=False),
+            ]))
+            dim_stage *= 2
+
+        # Bottleneck
+        self.bottleneck = MSAB(
+            dim=dim_stage, dim_head=dim, heads=dim_stage // dim, num_blocks=num_blocks[-1])
+
+        # Decoder
+        self.decoder_layers = nn.ModuleList([])
+        for i in range(stage):
+            self.decoder_layers.append(nn.ModuleList([
+                nn.ConvTranspose2d(dim_stage, dim_stage // 2, stride=2, kernel_size=2, padding=0, output_padding=0),
+                nn.Conv2d(dim_stage, dim_stage // 2, 1, 1, bias=False),
+                MSAB(
+                    dim=dim_stage // 2, num_blocks=num_blocks[stage - 1 - i], dim_head=dim,
+                    heads=(dim_stage // 2) // dim),
+            ]))
+            dim_stage //= 2
+
+        # Output projection
+        self.mapping = nn.Conv2d(self.dim, out_dim, 3, 1, 1, bias=False)
+
+        #### activation function
+        self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=True)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def forward(self, x):
+        """
+        x: [b,c,h,w]
+        return out:[b,c,h,w]
+        """
+
+        # Embedding
+        fea = self.embedding(x)
+
+        # Encoder
+        fea_encoder = []
+        for (MSAB, FeaDownSample) in self.encoder_layers:
+            fea = MSAB(fea)
+            fea_encoder.append(fea)
+            fea = FeaDownSample(fea)
+
+        # Bottleneck
+        fea = self.bottleneck(fea)
+
+        # Decoder
+        for i, (FeaUpSample, Fution, LeWinBlcok) in enumerate(self.decoder_layers):
+            fea = FeaUpSample(fea)
+            fea = Fution(torch.cat([fea, fea_encoder[self.stage-1-i]], dim=1))
+            fea = LeWinBlcok(fea)
+
+        # Mapping
+        out = self.mapping(fea) + x
+
+        return out
+
+class MST_Plus_Plus(nn.Module):
+    def __init__(self, in_channels=3, out_channels=31, n_feat=31, stage=3):
+        super(MST_Plus_Plus, self).__init__()
+        self.stage = stage
+        self.conv_in = nn.Conv2d(in_channels, n_feat, kernel_size=3, padding=(3 - 1) // 2,bias=False)
+        modules_body = [MST(dim=31, stage=2, num_blocks=[1,1,1]) for _ in range(stage)]
+        self.body = nn.Sequential(*modules_body)
+        self.conv_out = nn.Conv2d(n_feat, out_channels, kernel_size=3, padding=(3 - 1) // 2,bias=False)
+
+    def forward(self, x):
+        """
+        x: [b,c,h,w]
+        return out:[b,c,h,w]
+        """
+        b, c, h_inp, w_inp = x.shape
+        hb, wb = 8, 8
+        pad_h = (hb - h_inp % hb) % hb
+        pad_w = (wb - w_inp % wb) % wb
+        x = F.pad(x, [0, pad_w, 0, pad_h], mode='reflect')
+        x = self.conv_in(x)
+        h = self.body(x)
+        h = self.conv_out(h)
+        h += x
+        return h[:, :, :h_inp, :w_inp]
+
+
+
+
+
+
+
+
+
+
+
+
+
+

test_challenge_code/architecture/Restormer.py

@@ -0,0 +1,320 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numbers
+from einops import rearrange
+
+
+##########################################################################
+## Layer Norm
+
+def to_3d(x):
+    return rearrange(x, 'b c h w -> b (h w) c')
+
+
+def to_4d(x, h, w):
+    return rearrange(x, 'b (h w) c -> b c h w', h=h, w=w)
+
+
+class BiasFree_LayerNorm(nn.Module):
+    def __init__(self, normalized_shape):
+        super(BiasFree_LayerNorm, self).__init__()
+        if isinstance(normalized_shape, numbers.Integral):
+            normalized_shape = (normalized_shape,)
+        normalized_shape = torch.Size(normalized_shape)
+
+        assert len(normalized_shape) == 1
+
+        self.weight = nn.Parameter(torch.ones(normalized_shape))
+        self.normalized_shape = normalized_shape
+
+    def forward(self, x):
+        sigma = x.var(-1, keepdim=True, unbiased=False)
+        return x / torch.sqrt(sigma + 1e-5) * self.weight
+
+
+class WithBias_LayerNorm(nn.Module):
+    def __init__(self, normalized_shape):
+        super(WithBias_LayerNorm, self).__init__()
+        if isinstance(normalized_shape, numbers.Integral):
+            normalized_shape = (normalized_shape,)
+        normalized_shape = torch.Size(normalized_shape)
+
+        assert len(normalized_shape) == 1
+
+        self.weight = nn.Parameter(torch.ones(normalized_shape))
+        self.bias = nn.Parameter(torch.zeros(normalized_shape))
+        self.normalized_shape = normalized_shape
+
+    def forward(self, x):
+        mu = x.mean(-1, keepdim=True)
+        sigma = x.var(-1, keepdim=True, unbiased=False)
+        return (x - mu) / torch.sqrt(sigma + 1e-5) * self.weight + self.bias
+
+
+class LayerNorm(nn.Module):
+    def __init__(self, dim, LayerNorm_type):
+        super(LayerNorm, self).__init__()
+        if LayerNorm_type == 'BiasFree':
+            self.body = BiasFree_LayerNorm(dim)
+        else:
+            self.body = WithBias_LayerNorm(dim)
+
+    def forward(self, x):
+        h, w = x.shape[-2:]
+        return to_4d(self.body(to_3d(x)), h, w)
+
+
+##########################################################################
+## Gated-Dconv Feed-Forward Network (GDFN)
+class FeedForward(nn.Module):
+    def __init__(self, dim, ffn_expansion_factor, bias):
+        super(FeedForward, self).__init__()
+
+        hidden_features = int(dim * ffn_expansion_factor)
+
+        self.project_in = nn.Conv2d(dim, hidden_features * 2, kernel_size=1, bias=bias)
+
+        self.dwconv = nn.Conv2d(hidden_features * 2, hidden_features * 2, kernel_size=3, stride=1, padding=1,
+                                groups=hidden_features * 2, bias=bias)
+
+        self.project_out = nn.Conv2d(hidden_features, dim, kernel_size=1, bias=bias)
+
+    def forward(self, x):
+        x = self.project_in(x)
+        x1, x2 = self.dwconv(x).chunk(2, dim=1)
+        x = F.gelu(x1) * x2
+        x = self.project_out(x)
+        return x
+
+
+##########################################################################
+## Multi-DConv Head Transposed Self-Attention (MDTA)
+class Attention(nn.Module):
+    def __init__(self, dim, num_heads, bias):
+        super(Attention, self).__init__()
+        self.num_heads = num_heads
+        self.temperature = nn.Parameter(torch.ones(num_heads, 1, 1))
+
+        self.qkv = nn.Conv2d(dim, dim * 3, kernel_size=1, bias=bias)
+        self.qkv_dwconv = nn.Conv2d(dim * 3, dim * 3, kernel_size=3, stride=1, padding=1, groups=dim * 3, bias=bias)
+        self.project_out = nn.Conv2d(dim, dim, kernel_size=1, bias=bias)
+
+    def forward(self, x):
+        b, c, h, w = x.shape
+
+        qkv = self.qkv_dwconv(self.qkv(x))
+        q, k, v = qkv.chunk(3, dim=1)
+
+        q = rearrange(q, 'b (head c) h w -> b head c (h w)', head=self.num_heads)
+        k = rearrange(k, 'b (head c) h w -> b head c (h w)', head=self.num_heads)
+        v = rearrange(v, 'b (head c) h w -> b head c (h w)', head=self.num_heads)
+
+        q = torch.nn.functional.normalize(q, dim=-1)
+        k = torch.nn.functional.normalize(k, dim=-1)
+
+        attn = (q @ k.transpose(-2, -1)) * self.temperature
+        attn = attn.softmax(dim=-1)
+
+        out = (attn @ v)
+
+        out = rearrange(out, 'b head c (h w) -> b (head c) h w', head=self.num_heads, h=h, w=w)
+
+        out = self.project_out(out)
+        return out
+
+
+##########################################################################
+class TransformerBlock(nn.Module):
+    def __init__(self, dim, num_heads, ffn_expansion_factor, bias, LayerNorm_type):
+        super(TransformerBlock, self).__init__()
+
+        self.norm1 = LayerNorm(dim, LayerNorm_type)
+        self.attn = Attention(dim, num_heads, bias)
+        self.norm2 = LayerNorm(dim, LayerNorm_type)
+        self.ffn = FeedForward(dim, ffn_expansion_factor, bias)
+
+    def forward(self, x):
+        x = x + self.attn(self.norm1(x))
+        x = x + self.ffn(self.norm2(x))
+
+        return x
+
+
+##########################################################################
+## Overlapped image patch embedding with 3x3 Conv
+class OverlapPatchEmbed(nn.Module):
+    def __init__(self, in_c=3, embed_dim=48, bias=False):
+        super(OverlapPatchEmbed, self).__init__()
+
+        self.proj = nn.Conv2d(in_c, embed_dim, kernel_size=3, stride=1, padding=1, bias=bias)
+
+    def forward(self, x):
+        x = self.proj(x)
+
+        return x
+
+def pixel_unshuffle(input, downscale_factor):
+    '''
+    input: batchSize * c * k*w * k*h
+    downscale_factor: k
+    batchSize * c * k*w * k*h -> batchSize * k*k*c * w * h
+    '''
+    c = input.shape[1]
+    kernel = torch.zeros(size = [downscale_factor * downscale_factor * c, 1, downscale_factor, downscale_factor],
+                        device = input.device)
+    for y in range(downscale_factor):
+        for x in range(downscale_factor):
+            kernel[x + y * downscale_factor::downscale_factor * downscale_factor, 0, y, x] = 1
+    return F.conv2d(input, kernel, stride = downscale_factor, groups = c)
+
+class PixelUnShuffle(nn.Module):
+    def __init__(self, downscale_factor):
+        super(PixelUnShuffle, self).__init__()
+        self.downscale_factor = downscale_factor
+
+    def forward(self, input):
+        '''
+        input: batchSize * c * k*w * k*h
+        downscale_factor: k
+        batchSize * c * k*w * k*h -> batchSize * k*k*c * w * h
+        '''
+        return pixel_unshuffle(input, self.downscale_factor)
+
+##########################################################################
+## Resizing modules
+class Downsample(nn.Module):
+    def __init__(self, n_feat):
+        super(Downsample, self).__init__()
+
+        self.body = nn.Sequential(nn.Conv2d(n_feat, n_feat // 2, kernel_size=3, stride=1, padding=1, bias=False),
+                                  PixelUnShuffle(2))
+
+    def forward(self, x):
+        return self.body(x)
+
+
+class Upsample(nn.Module):
+    def __init__(self, n_feat):
+        super(Upsample, self).__init__()
+
+        self.body = nn.Sequential(nn.Conv2d(n_feat, n_feat * 2, kernel_size=3, stride=1, padding=1, bias=False),
+                                  nn.PixelShuffle(2))
+
+    def forward(self, x):
+        return self.body(x)
+
+
+##########################################################################
+##---------- Restormer -----------------------
+class Restormer(nn.Module):
+    def __init__(self,
+                 inp_channels=3,
+                 out_channels=31,
+                 dim=48,
+                 num_blocks=[2, 3, 3, 4],
+                 num_refinement_blocks=3,
+                 heads=[1, 2, 4, 8],
+                 ffn_expansion_factor=2.66,
+                 bias=False,
+                 LayerNorm_type='WithBias',  ## Other option 'BiasFree'
+                 dual_pixel_task=True  ## True for dual-pixel defocus deblurring only. Also set inp_channels=6
+                 ):
+
+        super(Restormer, self).__init__()
+
+        self.patch_embed = OverlapPatchEmbed(inp_channels, dim)
+
+        self.encoder_level1 = nn.Sequential(*[
+            TransformerBlock(dim=dim, num_heads=heads[0], ffn_expansion_factor=ffn_expansion_factor, bias=bias,
+                             LayerNorm_type=LayerNorm_type) for i in range(num_blocks[0])])
+
+        self.down1_2 = Downsample(dim)  ## From Level 1 to Level 2
+        self.encoder_level2 = nn.Sequential(*[
+            TransformerBlock(dim=int(dim * 2 ** 1), num_heads=heads[1], ffn_expansion_factor=ffn_expansion_factor,
+                             bias=bias, LayerNorm_type=LayerNorm_type) for i in range(num_blocks[1])])
+
+        self.down2_3 = Downsample(int(dim * 2 ** 1))  ## From Level 2 to Level 3
+        self.encoder_level3 = nn.Sequential(*[
+            TransformerBlock(dim=int(dim * 2 ** 2), num_heads=heads[2], ffn_expansion_factor=ffn_expansion_factor,
+                             bias=bias, LayerNorm_type=LayerNorm_type) for i in range(num_blocks[2])])
+
+        self.down3_4 = Downsample(int(dim * 2 ** 2))  ## From Level 3 to Level 4
+        self.latent = nn.Sequential(*[
+            TransformerBlock(dim=int(dim * 2 ** 3), num_heads=heads[3], ffn_expansion_factor=ffn_expansion_factor,
+                             bias=bias, LayerNorm_type=LayerNorm_type) for i in range(num_blocks[3])])
+
+        self.up4_3 = Upsample(int(dim * 2 ** 3))  ## From Level 4 to Level 3
+        self.reduce_chan_level3 = nn.Conv2d(int(dim * 2 ** 3), int(dim * 2 ** 2), kernel_size=1, bias=bias)
+        self.decoder_level3 = nn.Sequential(*[
+            TransformerBlock(dim=int(dim * 2 ** 2), num_heads=heads[2], ffn_expansion_factor=ffn_expansion_factor,
+                             bias=bias, LayerNorm_type=LayerNorm_type) for i in range(num_blocks[2])])
+
+        self.up3_2 = Upsample(int(dim * 2 ** 2))  ## From Level 3 to Level 2
+        self.reduce_chan_level2 = nn.Conv2d(int(dim * 2 ** 2), int(dim * 2 ** 1), kernel_size=1, bias=bias)
+        self.decoder_level2 = nn.Sequential(*[
+            TransformerBlock(dim=int(dim * 2 ** 1), num_heads=heads[1], ffn_expansion_factor=ffn_expansion_factor,
+                             bias=bias, LayerNorm_type=LayerNorm_type) for i in range(num_blocks[1])])
+
+        self.up2_1 = Upsample(int(dim * 2 ** 1))  ## From Level 2 to Level 1  (NO 1x1 conv to reduce channels)
+
+        self.decoder_level1 = nn.Sequential(*[
+            TransformerBlock(dim=int(dim * 2 ** 1), num_heads=heads[0], ffn_expansion_factor=ffn_expansion_factor,
+                             bias=bias, LayerNorm_type=LayerNorm_type) for i in range(num_blocks[0])])
+
+        self.refinement = nn.Sequential(*[
+            TransformerBlock(dim=int(dim * 2 ** 1), num_heads=heads[0], ffn_expansion_factor=ffn_expansion_factor,
+                             bias=bias, LayerNorm_type=LayerNorm_type) for i in range(num_refinement_blocks)])
+
+        #### For Dual-Pixel Defocus Deblurring Task ####
+        self.dual_pixel_task = dual_pixel_task
+        if self.dual_pixel_task:
+            self.skip_conv = nn.Conv2d(dim, int(dim * 2 ** 1), kernel_size=1, bias=bias)
+        ###########################
+
+        self.output = nn.Conv2d(int(dim * 2 ** 1), out_channels, kernel_size=3, stride=1, padding=1, bias=bias)
+
+    def forward(self, inp_img):
+        b, c, h_inp, w_inp = inp_img.shape
+        hb, wb = 8, 8
+        pad_h = (hb - h_inp % hb) % hb
+        pad_w = (wb - w_inp % wb) % wb
+        inp_img = F.pad(inp_img, [0, pad_w, 0, pad_h], mode='reflect')
+
+        inp_enc_level1 = self.patch_embed(inp_img)
+        out_enc_level1 = self.encoder_level1(inp_enc_level1)
+
+        inp_enc_level2 = self.down1_2(out_enc_level1)
+        out_enc_level2 = self.encoder_level2(inp_enc_level2)
+
+        inp_enc_level3 = self.down2_3(out_enc_level2)
+        out_enc_level3 = self.encoder_level3(inp_enc_level3)
+
+        inp_enc_level4 = self.down3_4(out_enc_level3)
+        latent = self.latent(inp_enc_level4)
+
+        inp_dec_level3 = self.up4_3(latent)
+        inp_dec_level3 = torch.cat([inp_dec_level3, out_enc_level3], 1)
+        inp_dec_level3 = self.reduce_chan_level3(inp_dec_level3)
+        out_dec_level3 = self.decoder_level3(inp_dec_level3)
+
+        inp_dec_level2 = self.up3_2(out_dec_level3)
+        inp_dec_level2 = torch.cat([inp_dec_level2, out_enc_level2], 1)
+        inp_dec_level2 = self.reduce_chan_level2(inp_dec_level2)
+        out_dec_level2 = self.decoder_level2(inp_dec_level2)
+
+        inp_dec_level1 = self.up2_1(out_dec_level2)
+        inp_dec_level1 = torch.cat([inp_dec_level1, out_enc_level1], 1)
+        out_dec_level1 = self.decoder_level1(inp_dec_level1)
+
+        out_dec_level1 = self.refinement(out_dec_level1)
+
+        #### For Dual-Pixel Defocus Deblurring Task ####
+        if self.dual_pixel_task:
+            out_dec_level1 = out_dec_level1 + self.skip_conv(inp_enc_level1)
+            out_dec_level1 = self.output(out_dec_level1)
+        ###########################
+        else:
+            out_dec_level1 = self.output(out_dec_level1) + inp_img
+
+        return out_dec_level1[:, :, :h_inp, :w_inp]

test_challenge_code/architecture/__init__.py

@@ -0,0 +1,41 @@
+import torch
+from .edsr import EDSR
+from .HDNet import HDNet
+from .hinet import HINet
+from .hrnet import SGN
+from .HSCNN_Plus import HSCNN_Plus
+from .MIRNet import MIRNet
+from .MPRNet import MPRNet
+from .MST import MST
+from .MST_Plus_Plus import MST_Plus_Plus
+from .Restormer import Restormer
+
+def model_generator(method, pretrained_model_path=None):
+    if method == 'mirnet':
+        model = MIRNet(n_RRG=3, n_MSRB=1, height=3, width=1).cuda()
+    elif method == 'mst_plus_plus':
+        model = MST_Plus_Plus().cuda()
+    elif method == 'mst':
+        model = MST(dim=31, stage=2, num_blocks=[4, 7, 5]).cuda()
+    elif method == 'hinet':
+        model = HINet(depth=4).cuda()
+    elif method == 'mprnet':
+        model = MPRNet(num_cab=4).cuda()
+    elif method == 'restormer':
+        model = Restormer().cuda()
+    elif method == 'edsr':
+        model = EDSR().cuda()
+    elif method == 'hdnet':
+        model = HDNet().cuda()
+    elif method == 'hrnet':
+        model = SGN().cuda()
+    elif method == 'hscnn_plus':
+        model = HSCNN_Plus().cuda()
+    else:
+        print(f'Method {method} is not defined !!!!')
+    if pretrained_model_path is not None:
+        print(f'load model from {pretrained_model_path}')
+        checkpoint = torch.load(pretrained_model_path)
+        model.load_state_dict({k.replace('module.', ''): v for k, v in checkpoint['state_dict'].items()},
+                              strict=True)
+    return model

test_challenge_code/architecture/edsr.py

@@ -0,0 +1,87 @@
+import torch.nn as nn
+
+def default_conv(in_channels, out_channels, kernel_size, bias=True):
+    return nn.Conv2d(
+        in_channels, out_channels, kernel_size,
+        padding=(kernel_size//2), bias=bias)
+
+class BasicBlock(nn.Sequential):
+    def __init__(
+        self, conv, in_channels, out_channels, kernel_size, stride=1, bias=False,
+        bn=True, act=nn.ReLU(True)):
+
+        m = [conv(in_channels, out_channels, kernel_size, bias=bias)]
+        if bn:
+            m.append(nn.BatchNorm2d(out_channels))
+        if act is not None:
+            m.append(act)
+
+        super(BasicBlock, self).__init__(*m)
+
+class ResBlock(nn.Module):
+    def __init__(
+        self, conv, n_feats, kernel_size,
+        bias=True, bn=False, act=nn.ReLU(True), res_scale=1):
+
+        super(ResBlock, self).__init__()
+        m = []
+        for i in range(2):
+            m.append(conv(n_feats, n_feats, kernel_size, bias=bias))
+            if bn:
+                m.append(nn.BatchNorm2d(n_feats))
+            if i == 0:
+                m.append(act)
+
+        self.body = nn.Sequential(*m)
+        self.res_scale = res_scale
+
+    def forward(self, x):
+        res = self.body(x).mul(self.res_scale)
+        res += x
+
+        return res
+
+
+
+class EDSR(nn.Module):
+    def __init__(self, conv=default_conv):
+        super(EDSR, self).__init__()
+
+        n_resblocks = 32
+        n_feats = 64
+        kernel_size = 3 
+        n_colors = 3
+        out_channels = 31
+        act = nn.ReLU(True)
+
+
+
+        # define head module
+        m_head = [conv(n_colors, n_feats, kernel_size)]
+
+        # define body module
+        m_body = [
+            ResBlock(
+                conv, n_feats, kernel_size, act=act, res_scale=1
+            ) for _ in range(n_resblocks)
+        ]
+        m_body.append(conv(n_feats, n_feats, kernel_size))
+
+        # define tail module
+        m_tail = [
+            conv(n_feats, out_channels, kernel_size)
+        ]
+
+        self.head = nn.Sequential(*m_head)
+        self.body = nn.Sequential(*m_body)
+        self.tail = nn.Sequential(*m_tail)
+
+    def forward(self, x):
+        x = self.head(x)
+
+        res = self.body(x)
+        res += x
+
+        x = self.tail(res)
+
+        return x

test_challenge_code/architecture/hinet.py

@@ -0,0 +1,212 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+def conv3x3(in_chn, out_chn, bias=True):
+    layer = nn.Conv2d(in_chn, out_chn, kernel_size=3, stride=1, padding=1, bias=bias)
+    return layer
+
+def conv_down(in_chn, out_chn, bias=False):
+    layer = nn.Conv2d(in_chn, out_chn, kernel_size=4, stride=2, padding=1, bias=bias)
+    return layer
+
+def conv(in_channels, out_channels, kernel_size, bias=False, stride = 1):
+    return nn.Conv2d(
+        in_channels, out_channels, kernel_size,
+        padding=(kernel_size//2), bias=bias, stride = stride)
+
+## Supervised Attention Module
+class SAM(nn.Module):
+    def __init__(self, n_feat, kernel_size=3, bias=True):
+        super(SAM, self).__init__()
+        self.conv1 = conv(n_feat, n_feat, kernel_size, bias=bias)
+        self.conv2 = conv(n_feat, n_feat, kernel_size, bias=bias)
+        self.conv3 = conv(n_feat, n_feat, kernel_size, bias=bias)
+
+    def forward(self, x, x_img):
+        x1 = self.conv1(x)
+        img = self.conv2(x) + x_img
+        x2 = torch.sigmoid(self.conv3(img))
+        x1 = x1*x2
+        x1 = x1+x
+        return x1, img
+
+class HINet(nn.Module):
+
+    def __init__(self, in_chn=31, out_chn=31, wf=31, depth=4, relu_slope=0.2, hin_position_left=0, hin_position_right=4):
+        super(HINet, self).__init__()
+
+        self.conv_in = nn.Conv2d(3, in_chn, kernel_size=3, padding=(3 - 1) // 2,
+                                 bias=False)
+        self.depth = depth
+        self.down_path_1 = nn.ModuleList()
+        self.down_path_2 = nn.ModuleList()
+        self.conv_01 = nn.Conv2d(in_chn, wf, 3, 1, 1)
+        self.conv_02 = nn.Conv2d(in_chn, wf, 3, 1, 1)
+
+        prev_channels = self.get_input_chn(wf)
+        for i in range(depth): #0,1,2,3,4
+            use_HIN = True if hin_position_left <= i and i <= hin_position_right else False
+            downsample = True if (i+1) < depth else False
+            self.down_path_1.append(UNetConvBlock(prev_channels, (2**i) * wf, downsample, relu_slope, use_HIN=use_HIN))
+            self.down_path_2.append(UNetConvBlock(prev_channels, (2**i) * wf, downsample, relu_slope, use_csff=downsample, use_HIN=use_HIN))
+            prev_channels = (2**i) * wf
+
+        self.up_path_1 = nn.ModuleList()
+        self.up_path_2 = nn.ModuleList()
+        self.skip_conv_1 = nn.ModuleList()
+        self.skip_conv_2 = nn.ModuleList()
+        for i in reversed(range(depth - 1)):
+            self.up_path_1.append(UNetUpBlock(prev_channels, (2**i)*wf, relu_slope))
+            self.up_path_2.append(UNetUpBlock(prev_channels, (2**i)*wf, relu_slope))
+            self.skip_conv_1.append(nn.Conv2d((2**i)*wf, (2**i)*wf, 3, 1, 1))
+            self.skip_conv_2.append(nn.Conv2d((2**i)*wf, (2**i)*wf, 3, 1, 1))
+            prev_channels = (2**i)*wf
+        self.sam12 = SAM(prev_channels)
+        self.cat12 = nn.Conv2d(prev_channels*2, prev_channels, 1, 1, 0)
+
+        self.last = conv3x3(prev_channels, out_chn, bias=True)
+
+    def forward(self, x):
+
+        b, c, h_inp, w_inp = x.shape
+        hb, wb = 16, 16
+        pad_h = (hb - h_inp % hb) % hb
+        pad_w = (wb - w_inp % wb) % wb
+        x = F.pad(x, [0, pad_w, 0, pad_h], mode='reflect')
+
+        image = self.conv_in(x)
+
+        #stage 1
+        x1 = self.conv_01(image)
+        encs = []
+        decs = []
+        for i, down in enumerate(self.down_path_1):
+            if (i+1) < self.depth:
+                x1, x1_up = down(x1)
+                encs.append(x1_up)
+            else:
+                x1 = down(x1)
+
+        for i, up in enumerate(self.up_path_1):
+            x1 = up(x1, self.skip_conv_1[i](encs[-i-1]))
+            decs.append(x1)
+
+        sam_feature, out_1 = self.sam12(x1, image)
+        #stage 2
+        x2 = self.conv_02(image)
+        x2 = self.cat12(torch.cat([x2, sam_feature], dim=1))
+        blocks = []
+        for i, down in enumerate(self.down_path_2):
+            if (i+1) < self.depth:
+                x2, x2_up = down(x2, encs[i], decs[-i-1])
+                blocks.append(x2_up)
+            else:
+                x2 = down(x2)
+
+        for i, up in enumerate(self.up_path_2):
+            x2 = up(x2, self.skip_conv_2[i](blocks[-i-1]))
+
+        out_2 = self.last(x2)
+        out_2 = out_2 + image
+        return out_2[:, :, :h_inp, :w_inp]
+
+    def get_input_chn(self, in_chn):
+        return in_chn
+
+    def _initialize(self):
+        gain = nn.init.calculate_gain('leaky_relu', 0.20)
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.orthogonal_(m.weight, gain=gain)
+                if not m.bias is None:
+                    nn.init.constant_(m.bias, 0)
+
+
+class UNetConvBlock(nn.Module):
+    def __init__(self, in_size, out_size, downsample, relu_slope, use_csff=False, use_HIN=False):
+        super(UNetConvBlock, self).__init__()
+        self.downsample = downsample
+        self.identity = nn.Conv2d(in_size, out_size, 1, 1, 0)
+        self.use_csff = use_csff
+
+        self.conv_1 = nn.Conv2d(in_size, out_size, kernel_size=3, padding=1, bias=True)
+        self.relu_1 = nn.LeakyReLU(relu_slope, inplace=False)
+        self.conv_2 = nn.Conv2d(out_size, out_size, kernel_size=3, padding=1, bias=True)
+        self.relu_2 = nn.LeakyReLU(relu_slope, inplace=False)
+
+        if downsample and use_csff:
+            self.csff_enc = nn.Conv2d(out_size, out_size, 3, 1, 1)
+            self.csff_dec = nn.Conv2d(out_size, out_size, 3, 1, 1)
+
+        if use_HIN:
+            self.norm = nn.InstanceNorm2d(((out_size+1)//2), affine=True)
+        self.use_HIN = use_HIN
+
+        if downsample:
+            self.downsample = conv_down(out_size, out_size, bias=False)
+
+    def forward(self, x, enc=None, dec=None):
+        out = self.conv_1(x)
+
+        if self.use_HIN:
+            out_1, out_2 = torch.chunk(out, 2, dim=1)
+            out = torch.cat([self.norm(out_1), out_2], dim=1)
+        out = self.relu_1(out)
+        out = self.relu_2(self.conv_2(out))
+
+        out += self.identity(x)
+        if enc is not None and dec is not None:
+            assert self.use_csff
+            out = out + self.csff_enc(enc) + self.csff_dec(dec)
+        if self.downsample:
+            out_down = self.downsample(out)
+            return out_down, out
+        else:
+            return out
+
+
+class UNetUpBlock(nn.Module):
+    def __init__(self, in_size, out_size, relu_slope):
+        super(UNetUpBlock, self).__init__()
+        self.up = nn.ConvTranspose2d(in_size, out_size, kernel_size=2, stride=2, bias=True)
+        self.conv_block = UNetConvBlock(in_size, out_size, False, relu_slope)
+
+    def forward(self, x, bridge):
+        up = self.up(x)
+        out = torch.cat([up, bridge], 1)
+        out = self.conv_block(out)
+        return out
+
+class Subspace(nn.Module):
+
+    def __init__(self, in_size, out_size):
+        super(Subspace, self).__init__()
+        self.blocks = nn.ModuleList()
+        self.blocks.append(UNetConvBlock(in_size, out_size, False, 0.2))
+        self.shortcut = nn.Conv2d(in_size, out_size, kernel_size=1, bias=True)
+
+    def forward(self, x):
+        sc = self.shortcut(x)
+        for i in range(len(self.blocks)):
+            x = self.blocks[i](x)
+        return x + sc
+
+class skip_blocks(nn.Module):
+
+    def __init__(self, in_size, out_size, repeat_num=1):
+        super(skip_blocks, self).__init__()
+        self.blocks = nn.ModuleList()
+        self.re_num = repeat_num
+        mid_c = 128
+        self.blocks.append(UNetConvBlock(in_size, mid_c, False, 0.2))
+        for i in range(self.re_num - 2):
+            self.blocks.append(UNetConvBlock(mid_c, mid_c, False, 0.2))
+        self.blocks.append(UNetConvBlock(mid_c, out_size, False, 0.2))
+        self.shortcut = nn.Conv2d(in_size, out_size, kernel_size=1, bias=True)
+
+    def forward(self, x):
+        sc = self.shortcut(x)
+        for m in self.blocks:
+            x = m(x)
+        return x + sc

test_challenge_code/architecture/hrnet.py

@@ -0,0 +1,484 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn import Parameter
+# ----------------------------------------
+#               Conv2d Block
+# ----------------------------------------
+class Conv2dLayer(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, pad_type='zero',
+                 activation='lrelu', norm='none', sn=False):
+        super(Conv2dLayer, self).__init__()
+        # Initialize the padding scheme
+        if pad_type == 'reflect':
+            self.pad = nn.ReflectionPad2d(padding)
+        elif pad_type == 'replicate':
+            self.pad = nn.ReplicationPad2d(padding)
+        elif pad_type == 'zero':
+            self.pad = nn.ZeroPad2d(padding)
+        else:
+            assert 0, "Unsupported padding type: {}".format(pad_type)
+
+        # Initialize the normalization type
+        if norm == 'bn':
+            self.norm = nn.BatchNorm2d(out_channels)
+        elif norm == 'in':
+            self.norm = nn.InstanceNorm2d(out_channels)
+        elif norm == 'ln':
+            self.norm = LayerNorm(out_channels)
+        elif norm == 'none':
+            self.norm = None
+        else:
+            assert 0, "Unsupported normalization: {}".format(norm)
+
+        # Initialize the activation funtion
+        if activation == 'relu':
+            self.activation = nn.ReLU(inplace=True)
+        elif activation == 'lrelu':
+            self.activation = nn.LeakyReLU(0.2, inplace=True)
+        elif activation == 'prelu':
+            self.activation = nn.PReLU()
+        elif activation == 'selu':
+            self.activation = nn.SELU(inplace=True)
+        elif activation == 'tanh':
+            self.activation = nn.Tanh()
+        elif activation == 'sigmoid':
+            self.activation = nn.Sigmoid()
+        elif activation == 'none':
+            self.activation = None
+        else:
+            assert 0, "Unsupported activation: {}".format(activation)
+
+        # Initialize the convolution layers
+        if sn:
+            self.conv2d = SpectralNorm(
+                nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding=0, dilation=dilation))
+        else:
+            self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding=0, dilation=dilation)
+
+    def forward(self, x):
+        x = self.pad(x)
+        x = self.conv2d(x)
+        if self.norm:
+            x = self.norm(x)
+        if self.activation:
+            x = self.activation(x)
+        return x
+
+
+class TransposeConv2dLayer(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, pad_type='zero',
+                 activation='lrelu', norm='none', sn=False, scale_factor=2):
+        super(TransposeConv2dLayer, self).__init__()
+        # Initialize the conv scheme
+        self.scale_factor = scale_factor
+        self.conv2d = Conv2dLayer(in_channels, out_channels, kernel_size, stride, padding, dilation, pad_type,
+                                  activation, norm, sn)
+
+    def forward(self, x):
+        x = F.interpolate(x, scale_factor=self.scale_factor, mode='nearest')
+        x = self.conv2d(x)
+        return x
+
+
+class ResConv2dLayer(nn.Module):
+    def __init__(self, in_channels, kernel_size=3, stride=1, padding=1, dilation=1, pad_type='zero', activation='lrelu',
+                 norm='none', sn=False, scale_factor=2):
+        super(ResConv2dLayer, self).__init__()
+        # Initialize the conv scheme
+        self.conv2d = nn.Sequential(
+            Conv2dLayer(in_channels, in_channels, kernel_size, stride, padding, dilation, pad_type, activation, norm,
+                        sn),
+            Conv2dLayer(in_channels, in_channels, kernel_size, stride, padding, dilation, pad_type, activation='none',
+                        norm=norm, sn=sn)
+        )
+
+    def forward(self, x):
+        residual = x
+        out = self.conv2d(x)
+        out = 0.1 * out + residual
+        return out
+
+
+class DenseConv2dLayer_5C(nn.Module):
+    def __init__(self, in_channels, latent_channels, kernel_size=3, stride=1, padding=1, dilation=1, pad_type='zero',
+                 activation='lrelu', norm='none', sn=False):
+        super(DenseConv2dLayer_5C, self).__init__()
+        # dense convolutions
+        self.conv1 = Conv2dLayer(in_channels, latent_channels, kernel_size, stride, padding, dilation, pad_type,
+                                 activation, norm, sn)
+        self.conv2 = Conv2dLayer(in_channels + latent_channels, latent_channels, kernel_size, stride, padding, dilation,
+                                 pad_type, activation, norm, sn)
+        self.conv3 = Conv2dLayer(in_channels + latent_channels * 2, latent_channels, kernel_size, stride, padding,
+                                 dilation, pad_type, activation, norm, sn)
+        self.conv4 = Conv2dLayer(in_channels + latent_channels * 3, latent_channels, kernel_size, stride, padding,
+                                 dilation, pad_type, activation, norm, sn)
+        self.conv5 = Conv2dLayer(in_channels + latent_channels * 4, in_channels, kernel_size, stride, padding, dilation,
+                                 pad_type, activation, norm, sn)
+
+    def forward(self, x):
+        x1 = self.conv1(x)
+        x2 = self.conv2(torch.cat((x, x1), 1))
+        x3 = self.conv3(torch.cat((x, x1, x2), 1))
+        x4 = self.conv4(torch.cat((x, x1, x2, x3), 1))
+        x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1))
+        return x5
+
+
+class ResidualDenseBlock_5C(nn.Module):
+    def __init__(self, in_channels, latent_channels, kernel_size=3, stride=1, padding=1, dilation=1, pad_type='zero',
+                 activation='lrelu', norm='none', sn=False):
+        super(ResidualDenseBlock_5C, self).__init__()
+        # dense convolutions
+        self.conv1 = Conv2dLayer(in_channels, latent_channels, kernel_size, stride, padding, dilation, pad_type,
+                                 activation, norm, sn)
+        self.conv2 = Conv2dLayer(in_channels + latent_channels, latent_channels, kernel_size, stride, padding, dilation,
+                                 pad_type, activation, norm, sn)
+        self.conv3 = Conv2dLayer(in_channels + latent_channels * 2, latent_channels, kernel_size, stride, padding,
+                                 dilation, pad_type, activation, norm, sn)
+        self.conv4 = Conv2dLayer(in_channels + latent_channels * 3, latent_channels, kernel_size, stride, padding,
+                                 dilation, pad_type, activation, norm, sn)
+        self.conv5 = Conv2dLayer(in_channels + latent_channels * 4, in_channels, kernel_size, stride, padding, dilation,
+                                 pad_type, activation, norm, sn)
+
+    def forward(self, x):
+        residual = x
+        x1 = self.conv1(x)
+        x2 = self.conv2(torch.cat((x, x1), 1))
+        x3 = self.conv3(torch.cat((x, x1, x2), 1))
+        x4 = self.conv4(torch.cat((x, x1, x2, x3), 1))
+        x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1))
+        x5 = 0.1 * x5 + residual
+        return x5
+
+
+# ----------------------------------------
+#               Layer Norm
+# ----------------------------------------
+class LayerNorm(nn.Module):
+    def __init__(self, num_features, eps=1e-8, affine=True):
+        super(LayerNorm, self).__init__()
+        self.num_features = num_features
+        self.affine = affine
+        self.eps = eps
+
+        if self.affine:
+            self.gamma = Parameter(torch.Tensor(num_features).uniform_())
+            self.beta = Parameter(torch.zeros(num_features))
+
+    def forward(self, x):
+        # layer norm
+        shape = [-1] + [1] * (x.dim() - 1)  # for 4d input: [-1, 1, 1, 1]
+        if x.size(0) == 1:
+            # These two lines run much faster in pytorch 0.4 than the two lines listed below.
+            mean = x.view(-1).mean().view(*shape)
+            std = x.view(-1).std().view(*shape)
+        else:
+            mean = x.view(x.size(0), -1).mean(1).view(*shape)
+            std = x.view(x.size(0), -1).std(1).view(*shape)
+        x = (x - mean) / (std + self.eps)
+        # if it is learnable
+        if self.affine:
+            shape = [1, -1] + [1] * (x.dim() - 2)  # for 4d input: [1, -1, 1, 1]
+            x = x * self.gamma.view(*shape) + self.beta.view(*shape)
+        return x
+
+
+# ----------------------------------------
+#           Spectral Norm Block
+# ----------------------------------------
+def l2normalize(v, eps=1e-12):
+    return v / (v.norm() + eps)
+
+
+class SpectralNorm(nn.Module):
+    def __init__(self, module, name='weight', power_iterations=1):
+        super(SpectralNorm, self).__init__()
+        self.module = module
+        self.name = name
+        self.power_iterations = power_iterations
+        if not self._made_params():
+            self._make_params()
+
+    def _update_u_v(self):
+        u = getattr(self.module, self.name + "_u")
+        v = getattr(self.module, self.name + "_v")
+        w = getattr(self.module, self.name + "_bar")
+
+        height = w.data.shape[0]
+        for _ in range(self.power_iterations):
+            v.data = l2normalize(torch.mv(torch.t(w.view(height, -1).data), u.data))
+            u.data = l2normalize(torch.mv(w.view(height, -1).data, v.data))
+
+        # sigma = torch.dot(u.data, torch.mv(w.view(height,-1).data, v.data))
+        sigma = u.dot(w.view(height, -1).mv(v))
+        setattr(self.module, self.name, w / sigma.expand_as(w))
+
+    def _made_params(self):
+        try:
+            u = getattr(self.module, self.name + "_u")
+            v = getattr(self.module, self.name + "_v")
+            w = getattr(self.module, self.name + "_bar")
+            return True
+        except AttributeError:
+            return False
+
+    def _make_params(self):
+        w = getattr(self.module, self.name)
+
+        height = w.data.shape[0]
+        width = w.view(height, -1).data.shape[1]
+
+        u = Parameter(w.data.new(height).normal_(0, 1), requires_grad=False)
+        v = Parameter(w.data.new(width).normal_(0, 1), requires_grad=False)
+        u.data = l2normalize(u.data)
+        v.data = l2normalize(v.data)
+        w_bar = Parameter(w.data)
+
+        del self.module._parameters[self.name]
+
+        self.module.register_parameter(self.name + "_u", u)
+        self.module.register_parameter(self.name + "_v", v)
+        self.module.register_parameter(self.name + "_bar", w_bar)
+
+    def forward(self, *args):
+        self._update_u_v()
+        return self.module.forward(*args)
+
+
+# ----------------------------------------
+#             Non-local Block
+# ----------------------------------------
+class Self_Attn(nn.Module):
+    """ Self attention Layer for Feature Map dimension"""
+
+    def __init__(self, in_dim, latent_dim=8):
+        super(Self_Attn, self).__init__()
+        self.channel_in = in_dim
+        self.channel_latent = in_dim // latent_dim
+        self.query_conv = nn.Conv2d(in_channels=in_dim, out_channels=in_dim // latent_dim, kernel_size=1)
+        self.key_conv = nn.Conv2d(in_channels=in_dim, out_channels=in_dim // latent_dim, kernel_size=1)
+        self.value_conv = nn.Conv2d(in_channels=in_dim, out_channels=in_dim, kernel_size=1)
+        self.gamma = nn.Parameter(torch.zeros(1))
+        self.softmax = nn.Softmax(dim=-1)
+
+    def forward(self, x):
+        """
+            inputs :
+                x : input feature maps(B X C X H X W)
+            returns :
+                out : self attention value + input feature
+                attention: B X N X N (N is Height * Width)
+        """
+        batchsize, C, height, width = x.size()
+        # proj_query: reshape to B x N x c, N = H x W
+        proj_query = self.query_conv(x).view(batchsize, -1, height * width).permute(0, 2, 1)
+        # proj_query: reshape to B x c x N, N = H x W
+        proj_key = self.key_conv(x).view(batchsize, -1, height * width)
+        # transpose check, energy: B x N x N, N = H x W
+        energy = torch.bmm(proj_query, proj_key)
+        # attention: B x N x N, N = H x W
+        attention = self.softmax(energy)
+        # proj_value is normal convolution, B x C x N
+        proj_value = self.value_conv(x).view(batchsize, -1, height * width)
+        # out: B x C x N
+        out = torch.bmm(proj_value, attention.permute(0, 2, 1))
+        out = out.view(batchsize, C, height, width)
+
+        out = self.gamma * out + x
+        return out
+
+
+# ----------------------------------------
+#              Global Block
+# ----------------------------------------
+class SELayer(nn.Module):
+    def __init__(self, channel, reduction=16):
+        super(SELayer, self).__init__()
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.fc = nn.Sequential(
+            nn.Linear(channel, channel // reduction, bias=False),
+            nn.ReLU(inplace=True),
+            nn.Linear(channel // reduction, channel // reduction, bias=False),
+            nn.ReLU(inplace=True),
+            nn.Linear(channel // reduction, channel, bias=False),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        b, c, _, _ = x.size()
+        y = self.avg_pool(x).view(b, c)
+        y = self.fc(y).view(b, c, 1, 1)
+        return x * y.expand_as(x)
+
+
+class GlobalBlock(nn.Module):
+    def __init__(self, in_channels, kernel_size, stride=1, padding=0, dilation=1, pad_type='zero', activation='lrelu',
+                 norm='none', sn=False, reduction=8):
+        super(GlobalBlock, self).__init__()
+        self.conv1 = Conv2dLayer(in_channels, in_channels, kernel_size, stride, padding, dilation, pad_type, activation,
+                                 norm, sn)
+        self.conv2 = Conv2dLayer(in_channels, in_channels, kernel_size, stride, padding, dilation, pad_type, activation,
+                                 norm, sn)
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.fc = nn.Sequential(
+            nn.Linear(in_channels, in_channels // reduction, bias=False),
+            nn.ReLU(inplace=True),
+            nn.Linear(in_channels // reduction, in_channels // reduction, bias=False),
+            nn.ReLU(inplace=True),
+            nn.Linear(in_channels // reduction, in_channels, bias=False),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        # residual
+        residual = x
+        # Sequeeze-and-Excitation(SE)
+        b, c, _, _ = x.size()
+        x = self.conv1(x)
+        y = self.avg_pool(x).view(b, c)
+        y = self.fc(y).view(b, c, 1, 1)
+        y = x * y.expand_as(x)
+        y = self.conv2(x)
+        # addition
+        out = 0.1 * y + residual
+        return out
+def pixel_unshuffle(input, downscale_factor):
+    '''
+    input: batchSize * c * k*w * k*h
+    downscale_factor: k
+    batchSize * c * k*w * k*h -> batchSize * k*k*c * w * h
+    '''
+    c = input.shape[1]
+    kernel = torch.zeros(size = [downscale_factor * downscale_factor * c, 1, downscale_factor, downscale_factor],
+                        device = input.device)
+    for y in range(downscale_factor):
+        for x in range(downscale_factor):
+            kernel[x + y * downscale_factor::downscale_factor * downscale_factor, 0, y, x] = 1
+    return F.conv2d(input, kernel, stride = downscale_factor, groups = c)
+
+class PixelUnShuffle(nn.Module):
+    def __init__(self, downscale_factor):
+        super(PixelUnShuffle, self).__init__()
+        self.downscale_factor = downscale_factor
+
+    def forward(self, input):
+        '''
+        input: batchSize * c * k*w * k*h
+        downscale_factor: k
+        batchSize * c * k*w * k*h -> batchSize * k*k*c * w * h
+        '''
+        return pixel_unshuffle(input, self.downscale_factor)
+
+# ----------------------------------------
+#         Initialize the networks
+# ----------------------------------------
+def weights_init(net, init_type = 'normal', init_gain = 0.02):
+    """Initialize network weights.
+    Parameters:
+        net (network)   -- network to be initialized
+        init_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonal
+        init_gain (float)    -- scaling factor for normal, xavier and orthogonal
+    In our paper, we choose the default setting: zero mean Gaussian distribution with a standard deviation of 0.02
+    """
+    def init_func(m):
+        classname = m.__class__.__name__
+        if hasattr(m, 'weight') and classname.find('Conv') != -1:
+            if init_type == 'normal':
+                torch.nn.init.normal_(m.weight.data, 0.0, init_gain)
+            elif init_type == 'xavier':
+                torch.nn.init.xavier_normal_(m.weight.data, gain = init_gain)
+            elif init_type == 'kaiming':
+                torch.nn.init.kaiming_normal_(m.weight.data, a = 0, mode = 'fan_in')
+            elif init_type == 'orthogonal':
+                torch.nn.init.orthogonal_(m.weight.data, gain = init_gain)
+            else:
+                raise NotImplementedError('initialization method [%s] is not implemented' % init_type)
+        elif classname.find('BatchNorm2d') != -1:
+            torch.nn.init.normal_(m.weight.data, 1.0, 0.02)
+            torch.nn.init.constant_(m.bias.data, 0.0)
+
+    # apply the initialization function <init_func>
+    print('initialize network with %s type' % init_type)
+    net.apply(init_func)
+
+# ----------------------------------------
+#                 Generator
+# ----------------------------------------
+class SGN(nn.Module):
+    def __init__(self, in_channels=3, out_channels=31, start_channels=64, pad='zero', activ='lrelu', norm='none', ):
+        super(SGN, self).__init__()
+        # Top subnetwork, K = 3
+        self.top1 = Conv2dLayer(in_channels * (4 ** 3), start_channels * (2 ** 3), 3, 1, 1, pad_type = pad, activation = activ, norm = norm)
+        self.top21 = ResidualDenseBlock_5C(start_channels * (2 ** 3), start_channels * (2 ** 2), 3, 1, 1, pad_type = pad, activation = activ, norm = norm)
+        self.top22 = GlobalBlock(start_channels * (2 ** 3), 3, 1, 1, pad_type = pad, activation = activ, norm = norm, sn = False, reduction = 4)
+        self.top3 = Conv2dLayer(start_channels * (2 ** 3), start_channels * (2 ** 3), 1, 1, 0, pad_type = pad, activation = activ, norm = norm)
+        # Middle subnetwork, K = 2
+        self.mid1 = Conv2dLayer(in_channels * (4 ** 2), start_channels * (2 ** 2), 3, 1, 1, pad_type = pad, activation = activ, norm = norm)
+        self.mid2 = Conv2dLayer(int(start_channels * (2 ** 2 + 2 ** 3 / 4)), start_channels * (2 ** 2), 3, 1, 1, pad_type = pad, activation = activ, norm = norm)
+        self.mid31 = ResidualDenseBlock_5C(start_channels * (2 ** 2), start_channels * (2 ** 1), 3, 1, 1, pad_type = pad, activation = activ, norm = norm)
+        self.mid32 = GlobalBlock(start_channels * (2 ** 2), 3, 1, 1, pad_type = pad, activation = activ, norm = norm, sn = False, reduction = 4)
+        self.mid4 = Conv2dLayer(start_channels * (2 ** 2), start_channels * (2 ** 2), 3, 1, 1, pad_type = pad, activation = activ, norm = norm)
+        # Bottom subnetwork, K = 1
+        self.bot1 = Conv2dLayer(in_channels * (4 ** 1), start_channels * (2 ** 1), 3, 1, 1, pad_type = pad, activation = activ, norm = norm)
+        self.bot2 = Conv2dLayer(int(start_channels * (2 ** 1 + 2 ** 2 / 4)), start_channels * (2 ** 1), 3, 1, 1, pad_type = pad, activation = activ, norm = norm)
+        self.bot31 = ResidualDenseBlock_5C(start_channels * (2 ** 1), start_channels * (2 ** 0), 3, 1, 1, pad_type = pad, activation = activ, norm = norm)
+        self.bot32 = ResidualDenseBlock_5C(start_channels * (2 ** 1), start_channels * (2 ** 0), 3, 1, 1, pad_type = pad, activation = activ, norm = norm)
+        self.bot33 = GlobalBlock(start_channels * (2 ** 1), 3, 1, 1, pad_type = pad, activation = activ, norm = norm, sn = False, reduction = 4)
+        self.bot4 = Conv2dLayer(start_channels * (2 ** 1), start_channels * (2 ** 1), 3, 1, 1, pad_type = pad, activation = activ, norm = norm)
+        # Mainstream
+        self.main1 = Conv2dLayer(in_channels, start_channels, 3, 1, 1, pad_type = pad, activation = activ, norm = norm)
+        self.main2 = Conv2dLayer(int(start_channels * (2 ** 0 + 2 ** 1 / 4)), start_channels, 3, 1, 1, pad_type = pad, activation = activ, norm = norm)
+        self.main31 = ResidualDenseBlock_5C(start_channels, start_channels // 2, 3, 1, 1, pad_type = pad, activation = activ, norm = norm)
+        self.main32 = ResidualDenseBlock_5C(start_channels, start_channels // 2, 3, 1, 1, pad_type = pad, activation = activ, norm = norm)
+        self.main33 = ResidualDenseBlock_5C(start_channels, start_channels // 2, 3, 1, 1, pad_type = pad, activation = activ, norm = norm)
+        self.main34 = ResidualDenseBlock_5C(start_channels, start_channels // 2, 3, 1, 1, pad_type = pad, activation = activ, norm = norm)
+        self.main35 = GlobalBlock(start_channels, 3, 1, 1, pad_type = pad, activation = activ, norm = norm, sn = False, reduction = 4)
+        self.main4 = Conv2dLayer(start_channels, out_channels, 3, 1, 1, pad_type = pad, activation = activ, norm = norm)
+
+    def forward(self, x):
+        b, c, h_inp, w_inp = x.shape
+        hb, wb = 8, 8
+        pad_h = (hb - h_inp % hb) % hb
+        pad_w = (wb - w_inp % wb) % wb
+        x = F.pad(x, [0, pad_w, 0, pad_h], mode='reflect')
+        # PixelUnShuffle                                        input: batch * 3 * 256 * 256
+        x1 = pixel_unshuffle(x, 2)               # out: batch * 12 * 128 * 128
+        x2 = pixel_unshuffle(x, 4)               # out: batch * 48 * 64 * 64
+        x3 = pixel_unshuffle(x, 8)               # out: batch * 192 * 32 * 32
+        # Top subnetwork                                        suppose the start_channels = 32
+        x3 = self.top1(x3)                                      # out: batch * 256 * 32 * 32
+        x3 = self.top21(x3)                                     # out: batch * 256 * 32 * 32
+        x3 = self.top22(x3)                                     # out: batch * 256 * 32 * 32
+        x3 = self.top3(x3)                                      # out: batch * 256 * 32 * 32
+        x3 = F.pixel_shuffle(x3, 2)                             # out: batch * 64 * 64 * 64, ready to be concatenated
+        # Middle subnetwork
+        x2 = self.mid1(x2)                                      # out: batch * 128 * 64 * 64
+        x2 = torch.cat((x2, x3), 1)                             # out: batch * (128 + 64) * 64 * 64
+        x2 = self.mid2(x2)                                      # out: batch * 128 * 64 * 64
+        x2 = self.mid31(x2)                                     # out: batch * 128 * 64 * 64
+        x2 = self.mid32(x2)                                     # out: batch * 128 * 64 * 64
+        x2 = self.mid4(x2)                                      # out: batch * 128 * 64 * 64
+        x2 = F.pixel_shuffle(x2, 2)                             # out: batch * 32 * 128 * 128, ready to be concatenated
+        # Bottom subnetwork
+        x1 = self.bot1(x1)                                      # out: batch * 64 * 128 * 128
+        x1 = torch.cat((x1, x2), 1)                             # out: batch * (64 + 32) * 128 * 128
+        x1 = self.bot2(x1)                                      # out: batch * 64 * 128 * 128
+        x1 = self.bot31(x1)                                     # out: batch * 64 * 128 * 128
+        x1 = self.bot32(x1)                                     # out: batch * 64 * 128 * 128
+        x1 = self.bot33(x1)                                     # out: batch * 64 * 128 * 128
+        x1 = self.bot4(x1)                                      # out: batch * 64 * 128 * 128
+        x1 = F.pixel_shuffle(x1, 2)                             # out: batch * 16 * 256 * 256, ready to be concatenated
+        # U-Net generator with skip connections from encoder to decoder
+        x = self.main1(x)                                       # out: batch * 32 * 256 * 256
+        x = torch.cat((x, x1), 1)                               # out: batch * (32 + 16) * 256 * 256
+        x = self.main2(x)                                       # out: batch * 32 * 256 * 256
+        x = self.main31(x)                                      # out: batch * 32 * 256 * 256
+        x = self.main32(x)                                      # out: batch * 32 * 256 * 256
+        x = self.main33(x)                                      # out: batch * 32 * 256 * 256
+        x = self.main34(x)                                      # out: batch * 32 * 256 * 256
+        x = self.main35(x)                                      # out: batch * 32 * 256 * 256
+        x = self.main4(x)                                       # out: batch * 3 * 256 * 256
+
+        return x[:, :, :h_inp, :w_inp]
+