scbc

.gitattributes

@@ -33,3 +33,8 @@ 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
+SCBC/Input/IMG_20240215_213330.png filter=lfs diff=lfs merge=lfs -text
+SCBC/Input/IMG_20240215_213619.png filter=lfs diff=lfs merge=lfs -text
+SCBC/Input/IMG_20240215_214449.png filter=lfs diff=lfs merge=lfs -text
+SCBC/Output/IMG_20240215_213330.png filter=lfs diff=lfs merge=lfs -text
+SCBC/Output/IMG_20240215_214449.png filter=lfs diff=lfs merge=lfs -text

SCBC/CPNet_model.py

@@ -0,0 +1,629 @@
+from __future__ import division
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.nn.init as init
+import torch.utils.model_zoo as model_zoo
+from torchvision import models
+from torchvision import transforms
+import cv2
+import matplotlib.pyplot as plt
+from PIL import Image
+import numpy as np
+import math
+import time
+import tqdm
+import os
+import argparse
+import copy
+import sys
+import networks as N
+from model_module import *
+sys.path.insert(0, '.')
+# from .common import *
+sys.path.insert(0, '../utils/')
+
+
+class LiteISPNet(nn.Module):
+    def __init__(self,):
+        super(LiteISPNet, self).__init__()
+
+        ch_1 = 64
+        ch_2 = 128
+        ch_3 = 128
+        n_blocks = 4
+
+
+        self.head = N.seq(
+            N.conv(3, ch_1, mode='C')
+        )  # shape: (N, ch_1, H/2, W/2)
+
+        self.down1 = N.seq(
+            N.conv(ch_1, ch_1, mode='C'),
+            N.RCAGroup(in_channels=ch_1, out_channels=ch_1, nb=n_blocks),
+            N.conv(ch_1, ch_1, mode='C'),
+            N.DWTForward(ch_1)
+        )  # shape: (N, ch_1*4, H/4, W/4)
+
+        self.down2 = N.seq(
+            N.conv(ch_1*4, ch_1, mode='C'),
+            N.RCAGroup(in_channels=ch_1, out_channels=ch_1, nb=n_blocks),
+            N.DWTForward(ch_1)
+        )  # shape: (N, ch_1*4, H/8, W/8)
+
+        self.down3 = N.seq(
+            N.conv(ch_1*4, ch_2, mode='C'),
+            N.RCAGroup(in_channels=ch_2, out_channels=ch_2, nb=n_blocks),
+            N.DWTForward(ch_2)
+        )  # shape: (N, ch_2*4, H/16, W/16)
+
+        self.middle = N.seq(
+            N.conv(ch_2*4, ch_3, mode='C'),
+            N.RCAGroup(in_channels=ch_3, out_channels=ch_3, nb=n_blocks),
+            N.RCAGroup(in_channels=ch_3, out_channels=ch_3, nb=n_blocks),
+            N.conv(ch_3, ch_2*4, mode='C')
+        )  # shape: (N, ch_2*4, H/16, W/16)
+
+        self.up3 = N.seq(
+            N.DWTInverse(ch_2*4),
+            N.RCAGroup(in_channels=ch_2, out_channels=ch_2, nb=n_blocks),
+            N.conv(ch_2, ch_1*4, mode='C')
+        )  # shape: (N, ch_1*4, H/8, W/8)
+
+        self.up2 = N.seq(
+            N.DWTInverse(ch_1*4),
+            N.RCAGroup(in_channels=ch_1, out_channels=ch_1, nb=n_blocks),
+            N.conv(ch_1, ch_1*4, mode='C')
+        )  # shape: (N, ch_1*4, H/4, W/4)
+
+        self.up1 = N.seq(
+            N.DWTInverse(ch_1*4),
+            N.RCAGroup(in_channels=ch_1, out_channels=ch_1, nb=n_blocks),
+            N.conv(ch_1, ch_1, mode='C')
+        )  # shape: (N, ch_1, H/2, W/2)
+
+        self.tail = N.seq(
+            #N.conv(ch_1, ch_1*4, mode='C'),
+            #nn.PixelShuffle(upscale_factor=2),
+            N.conv(ch_1, 3, mode='C')
+        )  # shape: (N, 3, H, W)
+
+    def forward(self, raw):
+        # input = raw
+        input = torch.pow(raw, 1/2.2)
+
+        h = self.head(input)
+        h_coord = h
+
+        d1 = self.down1(h_coord)
+        d2 = self.down2(d1)
+        d3 = self.down3(d2)
+        m = self.middle(d3) + d3
+        u3 = self.up3(m) + d2
+        u2 = self.up2(u3) + d1
+        u1 = self.up1(u2) + h
+        out = self.tail(u1)
+
+        return out
+
+
+class LiteAWBISPNet(nn.Module):
+    def __init__(self,):
+        super(LiteAWBISPNet, self).__init__()
+
+        ch_1 = 64
+        ch_2 = 128
+        ch_3 = 128
+        n_blocks = 4
+
+
+        self.head = N.seq(
+            N.conv(3, ch_1, mode='C')
+        )  # shape: (N, ch_1, H/2, W/2)
+
+        self.down1 = N.seq(
+            N.conv(ch_1, ch_1, mode='C'),
+            N.RCAGroup(in_channels=ch_1, out_channels=ch_1, nb=n_blocks),
+            N.conv(ch_1, ch_1, mode='C'),
+            N.DWTForward(ch_1)
+        )  # shape: (N, ch_1*4, H/4, W/4)
+
+        self.down2 = N.seq(
+            N.conv(ch_1*4, ch_1, mode='C'),
+            N.RCAGroup(in_channels=ch_1, out_channels=ch_1, nb=n_blocks),
+            N.DWTForward(ch_1)
+        )  # shape: (N, ch_1*4, H/8, W/8)
+
+        self.down3 = N.seq(
+            N.conv(ch_1*4, ch_2, mode='C'),
+            N.RCAGroup(in_channels=ch_2, out_channels=ch_2, nb=n_blocks),
+            N.DWTForward(ch_2)
+        )  # shape: (N, ch_2*4, H/16, W/16)
+
+        self.middle = N.seq(
+            N.conv(ch_2*4, ch_3, mode='C'),
+            N.RCAGroup(in_channels=ch_3, out_channels=ch_3, nb=n_blocks),
+            N.RCAGroup(in_channels=ch_3, out_channels=ch_3, nb=n_blocks),
+            N.conv(ch_3, ch_2*4, mode='C')
+        )  # shape: (N, ch_2*4, H/16, W/16)
+
+        self.up3 = N.seq(
+            N.DWTInverse(ch_2*4),
+            N.RCAGroup(in_channels=ch_2, out_channels=ch_2, nb=n_blocks),
+            N.conv(ch_2, ch_1*4, mode='C')
+        )  # shape: (N, ch_1*4, H/8, W/8)
+
+        self.up2 = N.seq(
+            N.DWTInverse(ch_1*4),
+            N.RCAGroup(in_channels=ch_1, out_channels=ch_1, nb=n_blocks),
+            N.conv(ch_1, ch_1*4, mode='C')
+        )  # shape: (N, ch_1*4, H/4, W/4)
+
+        self.up1 = N.seq(
+            N.DWTInverse(ch_1*4),
+            N.RCAGroup(in_channels=ch_1, out_channels=ch_1, nb=n_blocks),
+            N.conv(ch_1, ch_1, mode='C')
+        )  # shape: (N, ch_1, H/2, W/2)
+
+        self.tail = N.seq(
+            #N.conv(ch_1, ch_1*4, mode='C'),
+            #nn.PixelShuffle(upscale_factor=2),
+            N.conv(ch_1, 3, mode='C')
+        )  # shape: (N, 3, H, W)
+
+    def forward(self, raw):
+        # input = raw
+
+        input = raw
+        h = self.head(input)
+        h_coord = h
+
+        d1 = self.down1(h_coord)
+        d2 = self.down2(d1)
+        d3 = self.down3(d2)
+        m = self.middle(d3) + d3
+        u3 = self.up3(m) + d2
+        u2 = self.up2(u3) + d1
+        u1 = self.up1(u2) + h
+        out = self.tail(u1)
+
+        return out
+        
+        
+# Alignment Encoder
+class A_Encoder(nn.Module):
+    def __init__(self):
+        super(A_Encoder, self).__init__()
+        self.conv12 = Conv2d(3, 64, kernel_size=5, stride=2, padding=2, activation=nn.ReLU()) # 2
+        self.conv2 = Conv2d(64, 64, kernel_size=3, stride=1, padding=1, activation=nn.ReLU()) # 2
+        self.conv23 = Conv2d(64, 128, kernel_size=3, stride=2, padding=1, activation=nn.ReLU()) # 4
+        self.conv3 = Conv2d(128, 128, kernel_size=3, stride=1, padding=1, activation=nn.ReLU()) # 4
+        self.conv34 = Conv2d(128, 256, kernel_size=3, stride=2, padding=1, activation=nn.ReLU()) # 8
+        self.conv4a = Conv2d(256, 256, kernel_size=3, stride=1, padding=1, activation=nn.ReLU()) # 8
+        self.conv4b = Conv2d(256, 256, kernel_size=3, stride=1, padding=1, activation=nn.ReLU()) # 8
+        init_He(self)
+        self.register_buffer('mean', torch.FloatTensor([0.485, 0.456, 0.406]).view(1,3,1,1))
+        self.register_buffer('std', torch.FloatTensor([0.229, 0.224, 0.225]).view(1,3,1,1))
+        
+    def forward(self, in_f):
+        f = (in_f - self.mean) / self.std
+        x = f
+        x = F.upsample(x, size=(224, 224), mode='bilinear', align_corners=False)
+        x = self.conv12(x)
+        x = self.conv2(x)
+        x = self.conv23(x)
+        x = self.conv3(x)
+        x = self.conv34(x)
+        x = self.conv4a(x)
+        x = self.conv4b(x)
+        return x
+
+# Alignment Regressor
+class A_Regressor(nn.Module):
+    def __init__(self):
+        super(A_Regressor, self).__init__()
+        self.conv45 = Conv2d(512, 512, kernel_size=3, stride=2, padding=1, activation=nn.ReLU()) # 16
+        self.conv5a = Conv2d(512, 512, kernel_size=3, stride=1, padding=1, activation=nn.ReLU()) # 16
+        self.conv5b = Conv2d(512, 512, kernel_size=3, stride=1, padding=1, activation=nn.ReLU()) # 16
+        self.conv56 = Conv2d(512, 512, kernel_size=3, stride=2, padding=1, activation=nn.ReLU()) # 32
+        self.conv6a = Conv2d(512, 512, kernel_size=3, stride=1, padding=1, activation=nn.ReLU()) # 32
+        self.conv6b = Conv2d(512, 512, kernel_size=3, stride=1, padding=1, activation=nn.ReLU()) # 32
+        init_He(self)
+        
+        self.fc = nn.Linear(512, 6)
+        self.fc.weight.data.zero_()
+        self.fc.bias.data.copy_(torch.tensor([1, 0, 0, 0, 1, 0], dtype=torch.float32))
+
+    def forward(self, feat1, feat2):
+        x = torch.cat([feat1, feat2], dim=1)
+        x = self.conv45(x)
+        x = self.conv5a(x)
+        x = self.conv5b(x)
+        x = self.conv56(x)
+        x = self.conv5a(x)
+        x = self.conv5b(x)
+
+        x = F.avg_pool2d(x, x.shape[2])
+        x = x.view(-1, x.shape[1])
+
+        theta = self.fc(x)
+        theta = theta.view(-1, 2, 3)
+
+        return theta
+
+# Encoder (Copy network)
+class Encoder(nn.Module):
+    def __init__(self):
+        super(Encoder, self).__init__()
+        self.conv12 = Conv2d(4, 64, kernel_size=5, stride=2, padding=2, activation=nn.ReLU()) # 2
+        self.conv2 = Conv2d(64, 64, kernel_size=3, stride=1, padding=1, activation=nn.ReLU()) # 2
+        self.conv23 = Conv2d(64, 128, kernel_size=3, stride=2, padding=1, activation=nn.ReLU()) # 4
+        self.conv3 = Conv2d(128, 128, kernel_size=3, stride=1, padding=1, activation=nn.ReLU()) # 4
+        self.value3 = Conv2d(128, 128, kernel_size=3, stride=1, padding=1, activation=None) # 4
+        init_He(self)
+        self.register_buffer('mean', torch.FloatTensor([0.485, 0.456, 0.406]).view(1,3,1,1))
+        self.register_buffer('std', torch.FloatTensor([0.229, 0.224, 0.225]).view(1,3,1,1))
+        
+    def forward(self, in_f, in_v):
+        f = (in_f - self.mean) / self.std
+        x = torch.cat([f, in_v], dim=1)
+        x = self.conv12(x)
+        x = self.conv2(x)
+        x = self.conv23(x)
+        x = self.conv3(x)
+        v = self.value3(x)
+        return v
+
+# Decoder (Paste network)
+class Decoder(nn.Module):
+    def __init__(self):
+        super(Decoder, self).__init__()
+        self.conv4 = Conv2d(257, 257, kernel_size=3, stride=1, padding=1, activation=nn.ReLU())
+        self.conv5_1 = Conv2d(257, 257, kernel_size=3, stride=1, padding=1, activation=nn.ReLU())
+        self.conv5_2 = Conv2d(257, 257, kernel_size=3, stride=1, padding=1, activation=nn.ReLU())
+        
+        # dilated convolution blocks
+        self.convA4_1 = Conv2d(257, 257, kernel_size=3, stride=1, padding=2, D=2, activation=nn.ReLU())
+        self.convA4_2 = Conv2d(257, 257, kernel_size=3, stride=1, padding=4, D=4, activation=nn.ReLU())
+        self.convA4_3 = Conv2d(257, 257, kernel_size=3, stride=1, padding=8, D=8, activation=nn.ReLU())
+        self.convA4_4 = Conv2d(257, 257, kernel_size=3, stride=1, padding=16, D=16,activation=nn.ReLU())
+
+        self.conv3c = Conv2d(257, 257, kernel_size=3, stride=1, padding=1, activation=nn.ReLU()) # 4
+        self.conv3b = Conv2d(257, 128, kernel_size=3, stride=1, padding=1, activation=nn.ReLU()) # 4
+        self.conv3a = Conv2d(128, 128, kernel_size=3, stride=1, padding=1, activation=nn.ReLU()) # 4
+        self.conv32 = Conv2d(128, 64, kernel_size=3, stride=1, padding=1, activation=nn.ReLU()) # 2
+        self.conv2 = Conv2d(64, 64, kernel_size=3, stride=1, padding=1, activation=nn.ReLU()) # 2
+        self.conv21 = Conv2d(64, 3, kernel_size=5, stride=1, padding=2, activation=None) # 1
+
+        self.register_buffer('mean', torch.FloatTensor([0.485, 0.456, 0.406]).view(1,3,1,1))
+        self.register_buffer('std', torch.FloatTensor([0.229, 0.224, 0.225]).view(1,3,1,1))
+    
+    def forward(self, x):
+        x = self.conv4(x)
+        x = self.conv5_1(x)
+        x = self.conv5_2(x)
+
+        x = self.convA4_1(x)
+        x = self.convA4_2(x)
+        x = self.convA4_3(x)
+        x = self.convA4_4(x)
+
+        x = self.conv3c(x)
+        x = self.conv3b(x)
+        x = self.conv3a(x)
+        x = F.upsample(x, scale_factor=2, mode='nearest') # 2
+        x = self.conv32(x)
+        x = self.conv2(x)
+        x = F.upsample(x, scale_factor=2, mode='nearest') # 2
+        x = self.conv21(x)
+
+        p = (x *self.std) + self.mean
+        return p
+
+
+# Context Matching Module
+class CM_Module(nn.Module):
+    def __init__(self):
+        super(CM_Module, self).__init__()
+    
+    def masked_softmax(self, vec, mask, dim):
+        masked_vec = vec * mask.float()
+        max_vec = torch.max(masked_vec, dim=dim, keepdim=True)[0]
+        exps = torch.exp(masked_vec-max_vec)
+        masked_exps = exps * mask.float()
+        masked_sums = masked_exps.sum(dim, keepdim=True)
+        zeros = (masked_sums <1e-4)        
+        masked_sums += zeros.float()
+        return masked_exps/masked_sums
+
+    def forward(self, values, tvmap, rvmaps):
+
+        B, C, T, H, W = values.size()
+        # t_feat: target feature
+        t_feat = values[:, :, 0]
+        # r_feats: refetence features
+        r_feats = values[:, :, 1:]
+        
+        B, Cv, T, H, W = r_feats.size()
+        # vmap: visibility map
+        # tvmap: target visibility map
+        # rvmap: reference visibility map
+        # gs: cosine similarity
+        # c_m: c_match
+        gs_,vmap_ = [], []
+        tvmap_t = (F.upsample(tvmap, size=(H, W), mode='bilinear', align_corners=False)>0.5).float()
+        for r in range(T):
+            rvmap_t = (F.upsample(rvmaps[:,:,r], size=(H, W), mode='bilinear', align_corners=False)>0.5).float()
+            # vmap: visibility map
+            vmap = tvmap_t * rvmap_t
+            gs = (vmap * t_feat * r_feats[:,:,r]).sum(-1).sum(-1).sum(-1)
+            #valid sum
+            v_sum = vmap[:,0].sum(-1).sum(-1)
+            zeros = (v_sum <1e-4)
+            gs[zeros] = 0
+            v_sum += zeros.float()
+            gs = gs / v_sum / C
+            gs = torch.ones(t_feat.shape).float().cuda() * gs.view(B,1,1,1)
+            gs_.append(gs)
+            vmap_.append(rvmap_t)
+
+        gss = torch.stack(gs_, dim=2)
+        vmaps = torch.stack(vmap_, dim=2)
+        
+        #weighted pixelwise masked softmax
+        c_match = self.masked_softmax(gss, vmaps, dim=2)
+        c_out = torch.sum(r_feats * c_match, dim=2)
+
+        # c_mask
+        c_mask = (c_match * vmaps)
+        c_mask = torch.sum(c_mask,2)
+        c_mask = 1. - (torch.mean(c_mask, 1, keepdim=True))
+
+        return torch.cat([t_feat, c_out, c_mask], dim=1), c_mask
+
+
+class GCMModel(nn.Module):
+    def __init__(self):
+        super(GCMModel, self).__init__()
+        self.ch_1 = 16
+        self.ch_2 = 32
+        guide_input_channels = 3
+        align_input_channels = 3
+        self.gcm_coord = None
+
+        if not self.gcm_coord:
+            guide_input_channels = 3
+            align_input_channels = 3
+
+        self.guide_net = N.seq(
+            N.conv(guide_input_channels, self.ch_1, 7, stride=2, padding=0, mode='CR'),
+            N.conv(self.ch_1, self.ch_1, kernel_size=3, stride=1, padding=1, mode='CRC'),
+            nn.AdaptiveAvgPool2d(1),
+            N.conv(self.ch_1, self.ch_2, 1, stride=1, padding=0, mode='C')
+        )
+
+        self.align_head = N.conv(align_input_channels, self.ch_2, 1, padding=0, mode='CR')
+
+        self.align_base = N.seq(
+            N.conv(self.ch_2, self.ch_2, kernel_size=1, stride=1, padding=0, mode='CRCRCRCRCR')
+        )
+        self.align_tail = N.seq(
+            N.conv(self.ch_2, 3, 1, padding=0, mode='C')
+        )
+
+    def forward(self, demosaic_raw):
+        demosaic_raw = torch.pow(demosaic_raw, 1 / 2.2)
+        guide_input = demosaic_raw
+        base_input =demosaic_raw
+        guide = self.guide_net(guide_input)
+        out = self.align_head(base_input)
+        out = guide * out + out
+        out = self.align_base(out)
+        out = self.align_tail(out)+demosaic_raw
+
+        return out
+
+class Fusion(nn.Module):
+    def __init__(self):
+        super(Fusion, self).__init__()
+        self.ch_1 = 16
+        self.ch_2 = 32
+        guide_input_channels = 9
+        align_input_channels = 9
+        self.gcm_coord = None
+
+        if not self.gcm_coord:
+            guide_input_channels = 9
+            align_input_channels = 9
+
+        self.guide_net = N.seq(
+            N.conv(guide_input_channels, self.ch_1, 7, stride=2, padding=0, mode='CR'),
+            N.conv(self.ch_1, self.ch_1, kernel_size=3, stride=1, padding=1, mode='CRC'),
+            nn.AdaptiveAvgPool2d(1),
+            N.conv(self.ch_1, self.ch_2, 1, stride=1, padding=0, mode='C')
+        )
+
+        self.align_head = N.conv(align_input_channels, self.ch_2, 1, padding=0, mode='CR')
+
+        self.align_base = N.seq(
+            N.conv(self.ch_2, self.ch_2, kernel_size=1, stride=1, padding=0, mode='CRCRCR')
+        )
+        self.align_tail = N.seq(
+            N.conv(self.ch_2, 3, 1, padding=0, mode='C')
+        )
+
+    def forward(self, demosaic_raw):
+        #demosaic_raw = torch.pow(demosaic_raw, 1 / 2.2)
+        guide_input = demosaic_raw
+        base_input =demosaic_raw
+        guide = self.guide_net(guide_input)
+        out = self.align_head(base_input)
+        out = guide * out + out
+        out = self.align_base(out)
+        out = self.align_tail(out)
+
+        return out
+
+
+
+
+class CPNet(nn.Module):
+    def __init__(self, mode='Train'):
+        super(CPNet, self).__init__()
+        self.A_Encoder = A_Encoder()  # Align 
+        self.A_Regressor = A_Regressor() # output: alignment network
+        self.GCMModel = GCMModel()
+        self.Encoder = Encoder()  # Merge
+        self.CM_Module = CM_Module()
+
+        self.Decoder = Decoder() 
+        
+        self.register_buffer('mean', torch.FloatTensor([0.485, 0.456, 0.406]).view(1,3,1,1))
+        self.register_buffer('mean3d', torch.FloatTensor([0.485, 0.456, 0.406]).view(1,3,1,1,1))
+
+    
+    def encoding(self, frames, holes):
+
+        batch_size, _, num_frames, height, width = frames.size()
+        # padding
+        (frames, holes), pad = pad_divide_by([frames, holes], 8, (frames.size()[3], frames.size()[4]))
+        
+        feat_ = []
+        for t in range(num_frames):
+            feat = self.A_Encoder(frames[:,:,t], holes[:,:,t])
+            feat_.append(feat)
+        feats = torch.stack(feat_, dim=2)
+        return feats
+
+    def inpainting(self, rfeats, rframes, rholes, frame, hole, gt): 
+
+        batch_size, _, height, width = frame.size() # B C H W
+        num_r = rfeats.size()[2] # # of reference frames
+
+        # padding
+        (rframes, rholes, frame, hole, gt), pad = pad_divide_by([rframes, rholes, frame, hole, gt], 8, (height, width))
+        
+        # Target embedding
+        tfeat = self.A_Encoder(frame, hole)
+        
+        # c_feat: Encoder(Copy Network) features
+        c_feat_ = [self.Encoder(frame, hole)]
+        L_align = torch.zeros_like(frame)
+        
+        # aligned_r: aligned reference frames
+        aligned_r_ = []
+
+        # rvmap: aligned reference frames valid maps
+        rvmap_ = []
+        
+        for r in range(num_r):
+            theta_rt = self.A_Regressor(tfeat, rfeats[:,:,r])
+            grid_rt = F.affine_grid(theta_rt, frame.size())
+
+            # aligned_r: aligned reference frame
+            # reference frame affine transformation
+            aligned_r = F.grid_sample(rframes[:,:,r], grid_rt)
+            
+            # aligned_v: aligned reference visiblity map
+            # reference mask affine transformation
+            aligned_v = F.grid_sample(1-rholes[:,:,r], grid_rt)
+            aligned_v = (aligned_v>0.5).float()
+
+            aligned_r_.append(aligned_r)
+
+            #intersection of target and reference valid map
+            trvmap = (1-hole) * aligned_v
+            # compare the aligned frame - target frame 
+            
+            c_feat_.append(self.Encoder(aligned_r, aligned_v))
+            
+            rvmap_.append(aligned_v)
+
+        aligned_rs = torch.stack(aligned_r_, 2)
+        
+        c_feats =torch.stack(c_feat_, dim=2)
+        rvmaps = torch.stack(rvmap_, dim=2)
+
+        # p_in: paste network input(target features + c_out + c_mask)
+        p_in, c_mask = self.CM_Module(c_feats, 1-hole, rvmaps)
+        
+        pred = self.Decoder(p_in)
+        
+        _, _, _, H, W = aligned_rs.shape
+        c_mask = (F.upsample(c_mask, size=(H, W), mode='bilinear', align_corners=False)).detach()
+
+        comp = pred * (hole) + gt * (1.-hole)
+
+
+        if pad[2]+pad[3] > 0:
+            comp = comp[:,:,pad[2]:-pad[3],:]
+
+        if pad[0]+pad[1] > 0:
+            comp = comp[:,:,:,pad[0]:-pad[1]]
+            
+        comp = torch.clamp(comp, 0, 1)
+
+        return comp
+
+    def forward(self, Source, Target):
+
+        feat_target =self.A_Encoder(Target)
+        feat_source = self.A_Encoder(Source)
+
+        theta = self.A_Regressor(feat_target,feat_source)
+        grid_rt = F.affine_grid(theta, Target.size())
+        aligned = F.grid_sample(Source, grid_rt)
+        mask = torch.ones_like(Source)
+        mask = F.grid_sample(mask,grid_rt)
+
+        return aligned,mask
+
+
+class AC(nn.Module):
+    def __init__(self):
+        super(AC, self).__init__()
+        self.ch_1 = 32
+        self.ch_2 = 64
+        guide_input_channels = 8
+        align_input_channels = 5
+        self.gcm_coord = None
+
+        if not self.gcm_coord:
+            guide_input_channels = 6
+            align_input_channels = 3
+
+        self.guide_net = N.seq(
+            N.conv(guide_input_channels, self.ch_1, 7, stride=2, padding=0, mode='CR'),
+            N.conv(self.ch_1, self.ch_1, kernel_size=3, stride=1, padding=1, mode='CRC'),
+            nn.AdaptiveAvgPool2d(1),
+            N.conv(self.ch_1, self.ch_2, 1, stride=1, padding=0, mode='C')
+        )
+
+        self.align_head = N.conv(align_input_channels, self.ch_2, 1, padding=0, mode='CR')
+
+        self.align_base = N.seq(
+            N.conv(self.ch_2, self.ch_2, kernel_size=1, stride=1, padding=0, mode='CRCRCR')
+        )
+        self.align_tail = N.seq(
+            N.conv(self.ch_2, 3, 1, padding=0, mode='C')
+        )
+
+    def forward(self, demosaic_raw, dslr, coord=None):
+        demosaic_raw = demosaic_raw+0.01*torch.ones_like(demosaic_raw   )
+        demosaic_raw = torch.pow(demosaic_raw, 1 / 2.2)
+        demosaic_raw = demosaic_raw/2
+        if self.gcm_coord:
+            guide_input = torch.cat((demosaic_raw, dslr, coord), 1)
+            base_input = torch.cat((demosaic_raw, coord), 1)
+        else:
+            guide_input = torch.cat((demosaic_raw, dslr), 1)
+            base_input = demosaic_raw
+
+        guide = self.guide_net(guide_input)
+
+        out = self.align_head(base_input)
+        out = guide * out + out
+        out = self.align_base(out)
+        out = self.align_tail(out) +demosaic_raw
+
+        return out

SCBC/Dockerfile

@@ -0,0 +1,16 @@
+From python:3.8
+
+COPY . /SCBC
+WORKDIR /SCBC
+
+ARG DEBIAN_FRONTEND=noninteractive
+ENV TZ=Asia/Shanghai
+
+RUN apt-get update && apt-get install -y \
+    libpng-dev libjpeg-dev \
+    libopencv-dev ffmpeg \
+    libgl1-mesa-glx
+
+RUN python -m pip install --no-cache -r requirements.txt
+
+CMD ["./run.sh"]

SCBC/Input/IMG_20240215_213330.json

@@ -0,0 +1,25 @@
+{
+    "black_level": [
+        256,
+        256,
+        256,
+        256
+    ],
+    "white_level": 4095,
+    "noise_profile": [
+        0.001180699005,
+        6.3947934705e-06
+    ],
+    "cfa_pattern": [
+        0,
+        1,
+        1,
+        2
+    ],
+    "orientation": "Horizontal (normal)",
+    "as_shot_neutral": [
+        0.4234199302,
+        1.0,
+        0.2275
+    ]
+}

SCBC/Input/IMG_20240215_213619.json

@@ -0,0 +1,25 @@
+{
+    "black_level": [
+        256,
+        256,
+        256,
+        256
+    ],
+    "white_level": 4095,
+    "noise_profile": [
+        0.000575730186,
+        3.09754693248e-06
+    ],
+    "cfa_pattern": [
+        0,
+        1,
+        1,
+        2
+    ],
+    "orientation": "Horizontal (normal)",
+    "as_shot_neutral": [
+        0.4354066986,
+        1.0,
+        0.2288348701
+    ]
+}

SCBC/Input/IMG_20240215_214449.json

@@ -0,0 +1,25 @@
+{
+    "black_level": [
+        256,
+        256,
+        256,
+        256
+    ],
+    "white_level": 4095,
+    "noise_profile": [
+        0.002300534904,
+        2.25042231834722e-05
+    ],
+    "cfa_pattern": [
+        0,
+        1,
+        1,
+        2
+    ],
+    "orientation": "Horizontal (normal)",
+    "as_shot_neutral": [
+        0.4204851752,
+        1.0,
+        0.224368194
+    ]
+}

SCBC/Readme.txt

@@ -0,0 +1,2 @@
+> docker build -t scbc .
+> docker run --gpus all -it --rm -v $PWD/:/SCBC scbc sh run.sh

SCBC/SCBC_Solution.py

@@ -0,0 +1,130 @@
+import os
+import cv2
+import json
+import torch
+import torchvision.transforms as transforms
+from CPNet_model import LiteAWBISPNet
+import torchvision
+import numpy as np
+from Utiles import white_balance,apply_color_space_transform, transform_xyz_to_srgb, apply_gamma,fix_orientation,binning,Four2One,One2Four
+import time
+from net.mwrcanet import Net
+import torch.nn as nn
+from PIL import Image
+import torch.nn.functional as F
+
+#######Set Raw path###########
+Rpath = './Input'
+image_files = []
+
+####### Temp ###############################
+
+
+infer_times = []
+
+
+#######Color Matrix from Baseline#############
+color_matrix = [1.06835938, -0.29882812, -0.14257812,
+                -0.43164062, 1.35546875, 0.05078125,
+                -0.1015625, 0.24414062, 0.5859375]
+
+
+#######Data Transfer###########################
+transforms_ = [ transforms.ToTensor(),
+                transforms.Resize([768,1024])]
+transform = transforms.Compose(transforms_)
+
+transforms_ = [ transforms.ToTensor()]
+transformo = transforms.Compose(transforms_)
+
+########Load  the pretrained refinement model####
+model = LiteAWBISPNet()
+model.cuda()
+model.load_state_dict(torch.load('./model_zoo/CC2.pth') )
+
+######load pretrianed Denoised model##############
+last_ckpt = './model_zoo/dn_mwrcanet_raw_c1.pth'
+dn_net = Net()
+dn_model = nn.DataParallel(dn_net).cuda()
+tmp_ckpt = torch.load(last_ckpt)
+pretrained_dict = tmp_ckpt['state_dict']
+model_dict=dn_model.state_dict()
+pretrained_dict_update = {k: v for k, v in pretrained_dict.items() if k in model_dict}
+assert(len(pretrained_dict)==len(pretrained_dict_update))
+assert(len(pretrained_dict_update)==len(model_dict))
+model_dict.update(pretrained_dict_update)
+dn_model.load_state_dict(model_dict)
+
+############################Start Processing!#########
+
+for filename in os.listdir(Rpath):
+
+    if os.path.splitext(filename)[-1].lower() == ".png":
+        image_files.append(filename)
+
+with torch.no_grad():
+  for fp in image_files:
+  
+      fp = os.path.join(Rpath, fp)
+      mn = os.path.splitext(fp)[-2]
+      mf = str(mn) + '.json'
+  
+      raw_image = cv2.imread(fp, -1)
+      with open(mf, 'r') as file:
+          data = json.load(file)
+  
+  ############Bleack & Whilte##########################
+      time_BL_S = time.time()
+      
+      raw_image = (raw_image.astype(np.float32) - 256.)
+      raw_image = raw_image / (4095. - 256.)
+      raw_image = np.clip(raw_image, 0.0, 1.0)
+      
+  
+  
+  #############  Binning   ############################
+  
+      raw_image = binning(raw_image,data)
+  
+          
+  ############# Down sample ###########################
+  
+      
+      raw_image = cv2.resize(raw_image, [1024,768])
+      
+  
+  ############   Raw Denoise  ##########################
+        
+      Temp_I = Four2One(raw_image)
+      Temp_I = transformo(Temp_I).unsqueeze(0).cuda()
+      Temp_I = dn_model(Temp_I)
+      Temp_I = np.asarray(Temp_I.squeeze(0).squeeze(0).cpu())
+      raw_image = One2Four(Temp_I)
+      #raw_image = cv2.resize(raw_image, [1024,768])                   
+   
+  #############White Balance, Color M, Vignet #########
+      
+      raw_image = white_balance(raw_image, data['as_shot_neutral'])
+      raw_image = apply_color_space_transform(raw_image, color_matrix)
+      raw_image = transform_xyz_to_srgb(raw_image)
+      raw_image = apply_gamma(raw_image)
+
+          
+  #############Refinement#############################
+  
+      Source = transform(raw_image).unsqueeze(0).float().cuda()
+      Out = model(Source)
+
+  #################Saving#############################
+
+      Out = Out.clip(0,1)
+      OA = np.asarray(Out.squeeze(0).cpu()).transpose(1,2,0).astype(np.float32)
+      OA = OA*255.
+      OA = OA.astype(np.uint8)
+      OA = fix_orientation(OA,data["orientation"])
+      time_Save_F = time.time()
+      OA = cv2.cvtColor(OA, cv2.COLOR_RGB2BGR)
+      OA = cv2.imwrite('./Output/' + str(os.path.basename(fp)),OA)
+
+      infer_times.append(time_Save_F-time_BL_S)
+print(f"Average inference time: {np.mean(infer_times)} seconds")

SCBC/Utiles.py

@@ -0,0 +1,143 @@
+import  numpy as  np
+from fractions import Fraction
+import cv2
+import numpy as np
+import exifread
+from exifread.utils import Ratio
+import struct
+import json
+import torch
+import time
+
+Temp = np.ones([1536,2048]).astype(np.float32)
+Timg = np.ones([768,1024,3]).astype(np.float32)
+
+def apply_gamma(x):
+    # return x ** (1.0 / 2.2)
+    x = x.copy()
+    idx = x <= 0.0031308
+    x[idx] *= 12.92
+    x[idx == False] = (x[idx == False] ** (1.0 / 2.4)) * 1.055 - 0.055
+    return x
+
+def binning(img,data):
+
+    if data['cfa_pattern'] == [0,1,1,2]: 
+
+      ch_R  = img[0::2, 0::2]
+      ch_G  = (img[1::2, 0::2]+img[0::2,1::2])/2
+      ch_B  = img[1::2, 1::2]
+      out = np.dstack((ch_R, ch_G, ch_B))
+
+    if data['cfa_pattern'] == [2,1,1,0]:      
+    
+      ch_R  = img[1::2, 1::2]
+      ch_G  = (img[1::2, 0::2]+img[0::2,1::2])/2
+      ch_B  = img[0::2, 0::2]
+      out = np.dstack((ch_R, ch_G, ch_B))
+    
+    return out
+
+def Four2One(img):
+  Temp[0::2,0::2] = img[:,:,0]
+  Temp[1::2,0::2] = img[:,:,1]
+  Temp[0::2,1::2] = img[:,:,1]
+  Temp[1::2,1::2] = img[:,:,2]
+  
+  return Temp
+  
+def One2Four(Temp):
+  Timg[:,:,0] = Temp[0::2,0::2]
+  Timg[:,:,1] = (Temp[1::2,0::2]+Temp[0::2,1::2])/2
+  Timg[:,:,2] = Temp[1::2,1::2]
+  
+  return Timg
+  
+  
+def white_balance(demosaic_img, as_shot_neutral):
+    if type(as_shot_neutral[0]) is Ratio:
+        as_shot_neutral = ratios2floats(as_shot_neutral)
+
+    as_shot_neutral = np.asarray(as_shot_neutral)
+    # transform vector into matrix
+    if as_shot_neutral.shape == (3,):
+        as_shot_neutral = np.diag(1. / as_shot_neutral)
+
+    assert as_shot_neutral.shape == (3, 3)
+
+    white_balanced_image = np.dot(demosaic_img, as_shot_neutral.T)
+    white_balanced_image = np.clip(white_balanced_image, 0.0, 1.0)
+
+    return white_balanced_image
+    
+
+
+    
+
+
+def apply_color_space_transform(demosaiced_image, color_matrix):
+    xyz2cam = np.reshape(np.asarray(color_matrix), (3, 3))
+    # normalize rows (needed?)
+    xyz2cam = xyz2cam / np.sum(xyz2cam, axis=1, keepdims=True)
+    # inverse
+    cam2xyz = np.linalg.inv(xyz2cam)
+    # simplified matrix multiplication
+    xyz_image = cam2xyz[np.newaxis, np.newaxis, :, :] * \
+                demosaiced_image[:, :, np.newaxis, :]
+    xyz_image = np.sum(xyz_image, axis=-1)
+    xyz_image = np.clip(xyz_image, 0.0, 1.0)
+    return xyz_image
+
+
+    
+
+
+def transform_xyz_to_srgb(xyz_image):
+    xyz2srgb = np.array([[3.2404542, -1.5371385, -0.4985314],
+                         [-0.9692660, 1.8760108, 0.0415560],
+                         [0.0556434, -0.2040259, 1.0572252]])
+
+    # normalize rows (needed?)
+    xyz2srgb = xyz2srgb / np.sum(xyz2srgb, axis=-1, keepdims=True)
+
+    srgb_image = xyz2srgb[np.newaxis, np.newaxis, :, :] * xyz_image[:, :, np.newaxis, :]
+    srgb_image = np.sum(srgb_image, axis=-1)
+    srgb_image = np.clip(srgb_image, 0.0, 1.0)
+    return srgb_image
+
+
+
+
+def fix_orientation(image, orientation):
+    # 1 = Horizontal(normal)
+    # 2 = Mirror horizontal
+    # 3 = Rotate 180
+    # 4 = Mirror vertical
+    # 5 = Mirror horizontal and rotate 270 CW
+    # 6 = Rotate 90 CW
+    # 7 = Mirror horizontal and rotate 90 CW
+    # 8 = Rotate 270 CW
+
+    if type(orientation) is list:
+        orientation = orientation[0]
+
+    if orientation == "Horizontal(normal)":
+        pass
+    elif orientation == "Mirror horizonta":
+        image = cv2.flip(image, 0)
+    elif orientation == "Rotate 180":
+        image = cv2.rotate(image, cv2.ROTATE_180)
+    elif orientation == "Mirror vertical":
+        image = cv2.flip(image, 1)
+    elif orientation == "Mirror horizontal and rotate 270 CW":
+        image = cv2.flip(image, 0)
+        image = cv2.rotate(image, cv2.ROTATE_90_COUNTERCLOCKWISE)
+    elif orientation == "Rotate 90 CW":
+        image = cv2.rotate(image, cv2.ROTATE_90_CLOCKWISE)
+    elif orientation == "Mirror horizontal and rotate 90 CW":
+        image = cv2.flip(image, 0)
+        image = cv2.rotate(image, cv2.ROTATE_90_CLOCKWISE)
+    elif orientation == "Rotate 270 CW":
+        image = cv2.rotate(image, cv2.ROTATE_90_COUNTERCLOCKWISE)
+
+    return image

SCBC/model_module.py

@@ -0,0 +1,49 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.autograd import Variable
+import sys
+
+
+class Conv2d(nn.Module):
+    def __init__(self, in_ch, out_ch, kernel_size=3, stride=1, padding=1, D=1, activation=nn.ReLU()):
+        super(Conv2d, self).__init__()
+        if activation:
+            self.conv = nn.Sequential(
+                nn.Conv2d(in_ch, out_ch, kernel_size=kernel_size, stride=stride, padding=padding, dilation=D),
+                activation
+            )
+        else:
+            self.conv = nn.Sequential(
+                nn.Conv2d(in_ch, out_ch, kernel_size=kernel_size, stride=stride, padding=padding, dilation=D)
+            )
+
+    def forward(self, x):
+        x = self.conv(x)
+        return x
+        
+def init_He(module):
+    for m in module.modules():
+        if isinstance(m, nn.Conv2d):
+            nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+        elif isinstance(m, nn.BatchNorm2d):
+            nn.init.constant_(m.weight, 1)
+            nn.init.constant_(m.bias, 0)
+
+def pad_divide_by(in_list, d, in_size):
+    out_list = []
+    h, w = in_size
+    if h % d > 0:
+        new_h = h + d - h % d
+    else:
+        new_h = h
+    if w % d > 0:
+        new_w = w + d - w % d
+    else:
+        new_w = w
+    lh, uh = int((new_h-h) / 2), int(new_h-h) - int((new_h-h) / 2)
+    lw, uw = int((new_w-w) / 2), int(new_w-w) - int((new_w-w) / 2)
+    pad_array = (int(lw), int(uw), int(lh), int(uh))
+    for inp in in_list:
+        out_list.append(F.pad(inp, pad_array))
+    return out_list, pad_array

SCBC/model_zoo/CC2.pth

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

SCBC/model_zoo/dn_mwrcanet_raw_c1.pth

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

SCBC/models.py

@@ -0,0 +1,92 @@
+import torch.nn as nn
+import torch.nn.functional as F
+
+class ResidualBlock(nn.Module):
+    def __init__(self, in_features):
+        super(ResidualBlock, self).__init__()
+
+        conv_block = [  nn.ReflectionPad2d(1),
+                        nn.Conv2d(in_features, in_features, 3),
+                        nn.InstanceNorm2d(in_features),
+                        nn.ReLU(inplace=True),
+                        nn.ReflectionPad2d(1),
+                        nn.Conv2d(in_features, in_features, 3),
+                        nn.InstanceNorm2d(in_features)  ]
+
+        self.conv_block = nn.Sequential(*conv_block)
+
+    def forward(self, x):
+        return x + self.conv_block(x)
+
+class Generator(nn.Module):
+    def __init__(self, input_nc, output_nc, n_residual_blocks=9):
+        super(Generator, self).__init__()
+
+        # Initial convolution block       
+        model = [   nn.ReflectionPad2d(3),
+                    nn.Conv2d(input_nc, 64, 7),
+                    nn.InstanceNorm2d(64),
+                    nn.ReLU(inplace=True) ]
+
+        # Downsampling
+        in_features = 64
+        out_features = in_features*2
+        for _ in range(2):
+            model += [  nn.Conv2d(in_features, out_features, 3, stride=2, padding=1),
+                        nn.InstanceNorm2d(out_features),
+                        nn.ReLU(inplace=True) ]
+            in_features = out_features
+            out_features = in_features*2
+
+        # Residual blocks
+        for _ in range(n_residual_blocks):
+            model += [ResidualBlock(in_features)]
+
+        # Upsampling
+        out_features = in_features//2
+        for _ in range(2):
+            model += [  nn.ConvTranspose2d(in_features, out_features, 3, stride=2, padding=1, output_padding=1),
+                        nn.InstanceNorm2d(out_features),
+                        nn.ReLU(inplace=True) ]
+            in_features = out_features
+            out_features = in_features//2
+
+        # Output layer
+        model += [  nn.ReflectionPad2d(3),
+                    nn.Conv2d(64, output_nc, 7),
+                    nn.Tanh() ]
+
+        self.model = nn.Sequential(*model)
+
+    def forward(self, x):
+        return self.model(x)
+
+class Discriminator(nn.Module):
+    def __init__(self, input_nc):
+        super(Discriminator, self).__init__()
+
+        # A bunch of convolutions one after another
+        model = [   nn.Conv2d(input_nc, 64, 4, stride=2, padding=1),
+                    nn.LeakyReLU(0.2, inplace=True) ]
+
+        model += [  nn.Conv2d(64, 128, 4, stride=2, padding=1),
+                    nn.InstanceNorm2d(128), 
+                    nn.LeakyReLU(0.2, inplace=True) ]
+
+        model += [  nn.Conv2d(128, 256, 4, stride=2, padding=1),
+                    nn.InstanceNorm2d(256), 
+                    nn.LeakyReLU(0.2, inplace=True) ]
+
+        model += [  nn.Conv2d(256, 512, 4, padding=1),
+                    nn.InstanceNorm2d(512), 
+                    nn.LeakyReLU(0.2, inplace=True) ]
+
+        # FCN classification layer
+        model += [nn.Conv2d(512, 1, 4, padding=1)]
+
+        self.model = nn.Sequential(*model)
+
+    def forward(self, x):
+        x =  self.model(x)
+        # Average pooling and flatten
+        return F.avg_pool2d(x, x.size()[2:]).view(x.size()[0], -1)

SCBC/net/mwrcanet.py

@@ -0,0 +1,167 @@
+# -*- coding: utf-8 -*-
+# Yue Cao (cscaoyue@gmail.com) (cscaoyue@hit.edu.cn)
+# supervisor : Wangmeng Zuo (cswmzuo@gmail.com)
+# github: https://github.com/happycaoyue
+# personal link:   happycaoyue.com
+import torch
+import torch.nn as nn
+import numpy as np
+import torch.nn.init as init
+import torch.nn.functional as F
+class HITVPCTeam:
+    r"""
+        DWT and IDWT block written by: Yue Cao
+        """
+    class CALayer(nn.Module):
+        def __init__(self, channel=64, reduction=16):
+            super(HITVPCTeam.CALayer, 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
+
+    # conv - prelu - conv - sum
+    class RB(nn.Module):
+        def __init__(self, filters):
+            super(HITVPCTeam.RB, self).__init__()
+            self.conv1 = nn.Conv2d(filters, filters, 3, 1, 1)
+            self.act = nn.PReLU()
+            self.conv2 = nn.Conv2d(filters, filters, 3, 1, 1)
+            self.cuca = HITVPCTeam.CALayer(channel=filters)
+
+        def forward(self, x):
+            c0 = x
+            x = self.conv1(x)
+            x = self.act(x)
+            x = self.conv2(x)
+            out = self.cuca(x)
+            return out + c0
+
+    class NRB(nn.Module):
+        def __init__(self, n, f):
+            super(HITVPCTeam.NRB, self).__init__()
+            nets = []
+            for i in range(n):
+                nets.append(HITVPCTeam.RB(f))
+            self.body = nn.Sequential(*nets)
+            self.tail = nn.Conv2d(f, f, 3, 1, 1)
+
+        def forward(self, x):
+            return x + self.tail(self.body(x))
+
+    class DWTForward(nn.Module):
+        def __init__(self):
+            super(HITVPCTeam.DWTForward, self).__init__()
+            ll = np.array([[0.5, 0.5], [0.5, 0.5]])
+            lh = np.array([[-0.5, -0.5], [0.5, 0.5]])
+            hl = np.array([[-0.5, 0.5], [-0.5, 0.5]])
+            hh = np.array([[0.5, -0.5], [-0.5, 0.5]])
+            filts = np.stack([ll[None,::-1,::-1], lh[None,::-1,::-1],
+                              hl[None,::-1,::-1], hh[None,::-1,::-1]],
+                             axis=0)
+            self.weight = nn.Parameter(
+                torch.tensor(filts).to(torch.get_default_dtype()),
+                requires_grad=False)
+        def forward(self, x):
+            C = x.shape[1]
+            filters = torch.cat([self.weight,] * C, dim=0)
+            y = F.conv2d(x, filters, groups=C, stride=2)
+            return y
+
+    class DWTInverse(nn.Module):
+        def __init__(self):
+            super(HITVPCTeam.DWTInverse, self).__init__()
+            ll = np.array([[0.5, 0.5], [0.5, 0.5]])
+            lh = np.array([[-0.5, -0.5], [0.5, 0.5]])
+            hl = np.array([[-0.5, 0.5], [-0.5, 0.5]])
+            hh = np.array([[0.5, -0.5], [-0.5, 0.5]])
+            filts = np.stack([ll[None, ::-1, ::-1], lh[None, ::-1, ::-1],
+                              hl[None, ::-1, ::-1], hh[None, ::-1, ::-1]],
+                             axis=0)
+            self.weight = nn.Parameter(
+                torch.tensor(filts).to(torch.get_default_dtype()),
+                requires_grad=False)
+
+        def forward(self, x):
+            C = int(x.shape[1] / 4)
+            filters = torch.cat([self.weight, ] * C, dim=0)
+            y = F.conv_transpose2d(x, filters, groups=C, stride=2)
+            return y
+
+
+class Net(nn.Module):
+    def __init__(self, channels=1, filters_level1=96, filters_level2=256//2, filters_level3=256//2, n_rb=4*5):
+        super(Net, self).__init__()
+
+        self.head = HITVPCTeam.DWTForward()
+
+        self.down1 = nn.Sequential(
+            nn.Conv2d(channels * 4, filters_level1, 3, 1, 1),
+            nn.PReLU(),
+            HITVPCTeam.NRB(n_rb, filters_level1))
+
+        # sum 1
+        # self.down1 = HITVPCTeam.NRB(n_rb, filters_level1),
+
+        # sum 2
+        self.down2 = nn.Sequential(
+            HITVPCTeam.DWTForward(),
+            nn.Conv2d(filters_level1 * 4, filters_level2, 3, 1, 1),
+            nn.PReLU(),
+            HITVPCTeam.NRB(n_rb, filters_level2))
+
+        self.down3 = nn.Sequential(
+            HITVPCTeam.DWTForward(),
+            nn.Conv2d(filters_level2 * 4, filters_level3, 3, 1, 1),
+            nn.PReLU())
+
+        self.middle = HITVPCTeam.NRB(n_rb, filters_level3)
+
+        self.up1 = nn.Sequential(
+            nn.Conv2d(filters_level3, filters_level2 * 4, 3, 1, 1),
+            nn.PReLU(),
+            HITVPCTeam.DWTInverse())
+
+        self.up2 = nn.Sequential(
+            HITVPCTeam.NRB(n_rb, filters_level2),
+            nn.Conv2d(filters_level2, filters_level1 * 4, 3, 1, 1),
+            nn.PReLU(),
+            HITVPCTeam.DWTInverse())
+
+        self.up3 = nn.Sequential(
+            HITVPCTeam.NRB(n_rb, filters_level1),
+            nn.Conv2d(filters_level1, channels * 4, 3, 1, 1))
+
+        self.tail = HITVPCTeam.DWTInverse()
+
+    def forward(self, inputs):
+        c0 = inputs
+        c1 = self.head(c0)
+        c2 = self.down1(c1)
+        c3 = self.down2(c2)
+        c4 = self.down3(c3)
+        m = self.middle(c4)
+        c5 = self.up1(m) + c3
+        c6 = self.up2(c5) + c2
+        c7 = self.up3(c6) + c1
+        return self.tail(c7)
+
+    def _initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                init.orthogonal_(m.weight)
+                print('init weight')
+                if m.bias is not None:
+                    init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm2d):
+                init.constant_(m.weight, 1)
+                init.constant_(m.bias, 0)

SCBC/networks.py

@@ -0,0 +1,294 @@
+import torch
+import torch.nn as nn
+from torch.nn import init
+from torch.optim import lr_scheduler
+from collections import OrderedDict
+
+
+def get_scheduler(optimizer, opt):
+	if opt.lr_policy == 'linear':
+		def lambda_rule(epoch):
+			return 1 - max(0, epoch-opt.niter) / max(1, float(opt.niter_decay))
+		scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule)
+	elif opt.lr_policy == 'step':
+		scheduler = lr_scheduler.StepLR(optimizer,
+										step_size=opt.lr_decay_iters,
+										gamma=0.5)
+	elif opt.lr_policy == 'plateau':
+		scheduler = lr_scheduler.ReduceLROnPlateau(optimizer,
+												   mode='min',
+												   factor=0.2,
+												   threshold=0.01,
+												   patience=5)
+	elif opt.lr_policy == 'cosine':
+		scheduler = lr_scheduler.CosineAnnealingLR(optimizer,
+												   T_max=opt.niter,
+												   eta_min=0)
+	else:
+		return NotImplementedError('lr [%s] is not implemented', opt.lr_policy)
+	return scheduler
+
+def init_weights(net, init_type='normal', init_gain=0.02):
+	def init_func(m):  # define the initialization function
+		classname = m.__class__.__name__
+		if hasattr(m, 'weight') and (classname.find('Conv') != -1 \
+				or classname.find('Linear') != -1):
+			if init_type == 'normal':
+				init.normal_(m.weight.data, 0.0, init_gain)
+			elif init_type == 'xavier':
+				init.xavier_normal_(m.weight.data, gain=init_gain)
+			elif init_type == 'kaiming':
+				init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
+			elif init_type == 'orthogonal':
+				init.orthogonal_(m.weight.data, gain=init_gain)
+			elif init_type == 'uniform':
+				init.uniform_(m.weight.data, b=init_gain)
+			else:
+				raise NotImplementedError('[%s] is not implemented' % init_type)
+			if hasattr(m, 'bias') and m.bias is not None:
+				init.constant_(m.bias.data, 0.0)
+		elif classname.find('BatchNorm2d') != -1:
+			init.normal_(m.weight.data, 1.0, init_gain)
+			init.constant_(m.bias.data, 0.0)
+
+	print('initialize network with %s' % init_type)
+	net.apply(init_func)  # apply the initialization function <init_func>
+
+def init_net(net, init_type='default', init_gain=0.02, gpu_ids=[]):
+	if len(gpu_ids) > 0:
+		assert(torch.cuda.is_available())
+		net.to(gpu_ids[0])
+		net = torch.nn.DataParallel(net, gpu_ids)  # multi-GPUs
+	if init_type != 'default' and init_type is not None:
+		init_weights(net, init_type, init_gain=init_gain)
+	return net
+
+
+'''
+# ===================================
+# Advanced nn.Sequential
+# reform nn.Sequentials and nn.Modules
+# to a single nn.Sequential
+# ===================================
+'''
+
+def seq(*args):
+	if len(args) == 1:
+		args = args[0]
+	if isinstance(args, nn.Module):
+		return args
+	modules = OrderedDict()
+	if isinstance(args, OrderedDict):
+		for k, v in args.items():
+			modules[k] = seq(v)
+		return nn.Sequential(modules)
+	assert isinstance(args, (list, tuple))
+	return nn.Sequential(*[seq(i) for i in args])
+
+'''
+# ===================================
+# Useful blocks
+# --------------------------------
+# conv (+ normaliation + relu)
+# concat
+# sum
+# resblock (ResBlock)
+# resdenseblock (ResidualDenseBlock_5C)
+# resinresdenseblock (RRDB)
+# ===================================
+'''
+
+# -------------------------------------------------------
+# return nn.Sequantial of (Conv + BN + ReLU)
+# -------------------------------------------------------
+def conv(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1,
+		 output_padding=0, dilation=1, groups=1, bias=True,
+		 padding_mode='zeros', mode='CBR'):
+	L = []
+	for t in mode:
+		if t == 'C':
+			L.append(nn.Conv2d(in_channels=in_channels,
+							   out_channels=out_channels,
+							   kernel_size=kernel_size,
+							   stride=stride,
+							   padding=padding,
+							   dilation=dilation,
+							   groups=groups,
+							   bias=bias,
+							   padding_mode=padding_mode))
+		elif t == 'X':
+			assert in_channels == out_channels
+			L.append(nn.Conv2d(in_channels=in_channels,
+							   out_channels=out_channels,
+							   kernel_size=kernel_size,
+							   stride=stride,
+							   padding=padding,
+							   dilation=dilation,
+							   groups=in_channels,
+							   bias=bias,
+							   padding_mode=padding_mode))
+		elif t == 'T':
+			L.append(nn.ConvTranspose2d(in_channels=in_channels,
+										out_channels=out_channels,
+										kernel_size=kernel_size,
+										stride=stride,
+										padding=padding,
+										output_padding=output_padding,
+										groups=groups,
+										bias=bias,
+										dilation=dilation,
+										padding_mode=padding_mode))
+		elif t == 'B':
+			L.append(nn.BatchNorm2d(out_channels))
+		elif t == 'I':
+			L.append(nn.InstanceNorm2d(out_channels, affine=True))
+		elif t == 'i':
+			L.append(nn.InstanceNorm2d(out_channels))
+		elif t == 'R':
+			L.append(nn.ReLU(inplace=True))
+		elif t == 'r':
+			L.append(nn.ReLU(inplace=False))
+		elif t == 'S':
+			L.append(nn.Sigmoid())
+		elif t == 'P':
+			L.append(nn.PReLU())
+		elif t == 'L':
+			L.append(nn.LeakyReLU(negative_slope=1e-1, inplace=True))
+		elif t == 'l':
+			L.append(nn.LeakyReLU(negative_slope=1e-1, inplace=False))
+		elif t == '2':
+			L.append(nn.PixelShuffle(upscale_factor=2))
+		elif t == '3':
+			L.append(nn.PixelShuffle(upscale_factor=3))
+		elif t == '4':
+			L.append(nn.PixelShuffle(upscale_factor=4))
+		elif t == 'U':
+			L.append(nn.Upsample(scale_factor=2, mode='nearest'))
+		elif t == 'u':
+			L.append(nn.Upsample(scale_factor=3, mode='nearest'))
+		elif t == 'M':
+			L.append(nn.MaxPool2d(kernel_size=kernel_size,
+								  stride=stride,
+								  padding=0))
+		elif t == 'A':
+			L.append(nn.AvgPool2d(kernel_size=kernel_size,
+								  stride=stride,
+								  padding=0))
+		else:
+			raise NotImplementedError('Undefined type: '.format(t))
+	return seq(*L)
+
+
+class DWTForward(nn.Conv2d):
+    def __init__(self, in_channels=64):
+        super(DWTForward, self).__init__(in_channels, in_channels*4, 2, 2,
+                                  groups=in_channels, bias=False)
+        weight = torch.tensor([[[[0.5,  0.5], [ 0.5,  0.5]]],
+                               [[[0.5,  0.5], [-0.5, -0.5]]],
+                               [[[0.5, -0.5], [ 0.5, -0.5]]],
+                               [[[0.5, -0.5], [-0.5,  0.5]]]],
+                              dtype=torch.get_default_dtype()
+                             ).repeat(in_channels, 1, 1, 1)# / 2
+        self.weight.data.copy_(weight)
+        self.requires_grad_(False)
+
+
+class DWTInverse(nn.ConvTranspose2d):
+    def __init__(self, in_channels=64):
+        super(DWTInverse, self).__init__(in_channels, in_channels//4, 2, 2,
+                                  groups=in_channels//4, bias=False)
+        weight = torch.tensor([[[[0.5,  0.5], [ 0.5,  0.5]]],
+                               [[[0.5,  0.5], [-0.5, -0.5]]],
+                               [[[0.5, -0.5], [ 0.5, -0.5]]],
+                               [[[0.5, -0.5], [-0.5,  0.5]]]],
+                              dtype=torch.get_default_dtype()
+                             ).repeat(in_channels//4, 1, 1, 1)# * 2
+        self.weight.data.copy_(weight)
+        self.requires_grad_(False)
+
+
+# -------------------------------------------------------
+# Channel Attention (CA) Layer
+# -------------------------------------------------------
+class CALayer(nn.Module):
+	def __init__(self, channel=64, reduction=16):
+		super(CALayer, 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 
+
+
+# -------------------------------------------------------
+# Res Block: x + conv(relu(conv(x)))
+# -------------------------------------------------------
+class ResBlock(nn.Module):
+	def __init__(self, in_channels=64, out_channels=64, kernel_size=3, stride=1,
+				 padding=1, bias=True, mode='CRC'):
+		super(ResBlock, self).__init__()
+
+		assert in_channels == out_channels
+		if mode[0] in ['R','L']:
+			mode = mode[0].lower() + mode[1:]
+
+		self.res = conv(in_channels, out_channels, kernel_size,
+						stride, padding=padding, bias=bias, mode=mode)
+
+	def forward(self, x):
+		res = self.res(x)
+		return x + res
+
+
+# -------------------------------------------------------
+# Residual Channel Attention Block (RCAB)
+# -------------------------------------------------------
+class RCABlock(nn.Module):
+	def __init__(self, in_channels=64, out_channels=64, kernel_size=3, stride=1,
+				 padding=1, bias=True, mode='CRC', reduction=16):
+		super(RCABlock, self).__init__()
+		assert in_channels == out_channels
+		if mode[0] in ['R','L']:
+			mode = mode[0].lower() + mode[1:]
+
+		self.res = conv(in_channels, out_channels, kernel_size,
+						stride, padding, bias=bias, mode=mode)
+		self.ca = CALayer(out_channels, reduction)
+
+	def forward(self, x):
+		res = self.res(x)
+		res = self.ca(res) 
+		return res + x
+
+
+# -------------------------------------------------------
+# Residual Channel Attention Group (RG)
+# -------------------------------------------------------
+class RCAGroup(nn.Module):
+	def __init__(self, in_channels=64, out_channels=64, kernel_size=3, stride=1,
+				 padding=1, bias=True, mode='CRC', reduction=16, nb=12):
+		super(RCAGroup, self).__init__()
+		assert in_channels == out_channels
+		if mode[0] in ['R','L']:
+			mode = mode[0].lower() + mode[1:]
+
+		RG = [RCABlock(in_channels, out_channels, kernel_size, stride, padding,
+					   bias, mode, reduction) for _ in range(nb)]
+		# RG = [ResBlock(in_channels, out_channels, kernel_size, stride, padding,
+		#                bias, mode) for _ in range(nb)]
+		RG.append(conv(out_channels, out_channels, mode='C'))
+
+		self.rg = nn.Sequential(*RG)
+
+	def forward(self, x):
+		res = self.rg(x)
+		return res + x
+

SCBC/requirements.txt

@@ -0,0 +1,13 @@
+opencv-python
+scipy
+numpy
+torch
+pandas
+torchvision
+Pillow
+matplotlib
+tqdm
+imageio
+seaborn
+hdf5storage
+exifread

SCBC/run.sh

@@ -0,0 +1 @@
+python SCBC_Solution.py