spikezoo 0.1.2__py3-none-any.whl → 0.2.1__py3-none-any.whl

spikezoo/archs/ssir/utils.py

@@ -0,0 +1,154 @@
+import numpy as np
+import torch
+import torch.nn.functional as F
+import os
+import os.path as osp
+import random
+import cv2
+
+def set_seeds(_seed_):
+    random.seed(_seed_)
+    np.random.seed(_seed_)
+    torch.manual_seed(_seed_)  # use torch.manual_seed() to seed the RNG for all devices (both CPU and CUDA)
+    torch.cuda.manual_seed_all(_seed_)
+
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+    # os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8'
+    # set a debug environment variable CUBLAS_WORKSPACE_CONFIG to ":16:8" (may limit overall performance) or ":4096:8" (will increase library footprint in GPU memory by approximately 24MiB).
+    # torch.use_deterministic_algorithms(True)
+
+
+def make_dir(path):
+    if not osp.exists(path):
+        os.makedirs(path)
+    return
+
+
+def add_args_to_cfg(cfg, args, args_list):
+    for aa in args_list:
+        cfg['train'][aa] = eval('args.{:s}'.format(aa))
+    return cfg
+
+
+# class AverageMeter(object):
+#     """Computes and stores the average and current value"""
+#     def __init__(self, precision=3):
+#         self.precision = precision
+#         self.reset()
+
+#     def reset(self):
+#         self.val = 0
+#         self.avg = 0
+#         self.sum = 0
+#         self.count = 0
+
+#     def update(self, val, n=1):
+#         self.val = val
+#         self.sum += val * n
+#         self.count += n
+#         self.avg = self.sum / self.count
+
+#     def __repr__(self):
+#         return '{:.{}f} ({:.{}f})'.format(self.val, self.precision, self.avg, self.precision)
+
+
+class AverageMeter(object):
+    """Computes and stores the average and current value"""
+
+    def __init__(self, i=1, precision=3, names=None):
+        self.meters = i
+        self.precision = precision
+        self.reset(self.meters)
+        self.names = names
+        if names is not None:
+            assert self.meters == len(self.names)
+        else:
+            self.names = [''] * self.meters
+
+    def reset(self, i):
+        self.val = [0] * i
+        self.avg = [0] * i
+        self.sum = [0] * i
+        self.count = [0] * i
+
+    def update(self, val, n=1):
+        if not isinstance(val, list):
+            val = [val]
+        if not isinstance(n, list):
+            n = [n] * self.meters
+        assert (len(val) == self.meters and len(n) == self.meters)
+        for i in range(self.meters):
+            self.count[i] += n[i]
+        for i, v in enumerate(val):
+            self.val[i] = v
+            self.sum[i] += v * n[i]
+            self.avg[i] = self.sum[i] / self.count[i]
+
+    def __repr__(self):
+        # val = ' '.join(['{} {:.{}f}'.format(n, v, self.precision) for n, v in
+        #                 zip(self.names, self.val)])
+        # avg = ' '.join(['{} {:.{}f}'.format(n, a, self.precision) for n, a in
+        #                 zip(self.names, self.avg)])
+        out = ' '.join(['{} {:.{}f} ({:.{}f})'.format(n, v, self.precision, a, self.precision) for n, v, a in
+                        zip(self.names, self.val, self.avg)])
+        # return '{} ({})'.format(val, avg)
+        return '{}'.format(out)
+
+
+def normalize_image_torch(image, percentile_lower=1, percentile_upper=99):
+    b, c, h, w = image.shape
+    image_reshape = image.reshape([b, c, h*w])
+    mini = torch.quantile(image_reshape, 0.01, dim=2, keepdim=True).unsqueeze_(dim=3)
+    maxi = torch.quantile(image_reshape, 0.99, dim=2, keepdim=True).unsqueeze_(dim=3)
+    # if mini == maxi:
+    #     return 0 * image + 0.5  # gray image
+    return torch.clip((image - mini) / (maxi - mini + 1e-5), 0, 1)
+
+def normalize_image_torch2(image):
+    return torch.clip(image, 0, 1)
+
+# --------------------------------------------
+# Torch to Numpy 0~255
+# --------------------------------------------
+def torch2numpy255(im):
+    im = im[0, 0].detach().cpu().numpy()
+    im = (im * 255).astype(np.float64)
+    return im
+
+def torch2torch255(im):
+    return im * 255.0
+
+class InputPadder:
+    """ Pads images such that dimensions are divisible by padsize """
+    def __init__(self, dims, padsize=16):
+        self.ht, self.wd = dims[-2:]
+        pad_ht = (((self.ht // padsize) + 1) * padsize - self.ht) % padsize
+        pad_wd = (((self.wd // padsize) + 1) * padsize - self.wd) %padsize
+        self._pad = [pad_wd//2, pad_wd - pad_wd//2, 0, pad_ht]
+
+    def pad(self, *inputs):
+        return [F.pad(x, self._pad, mode='replicate') for x in inputs]
+
+    def unpad(self,x):
+        ht, wd = x.shape[-2:]
+        c = [self._pad[2], ht-self._pad[3], self._pad[0], wd-self._pad[1]]
+        return x[..., c[0]:c[1], c[2]:c[3]]
+    
+
+
+def vis_img(vis_path: str, img: torch.Tensor, vis_name: str = 'vis'):
+    ww = 0
+    rows = []
+    for ii in range(4):
+        cur_row = []
+        for jj in range(img.shape[0]//4):
+            cur_img = img[ww, 0].detach().cpu().numpy() * 255
+            cur_img = cur_img.astype(np.uint8)
+            cur_row.append(cur_img)
+            ww += 1
+        cur_row_cat = np.concatenate(cur_row, axis=1)
+        rows.append(cur_row_cat)
+    out_img = np.concatenate(rows, axis=0)
+    cv2.imwrite(osp.join(vis_path, vis_name+'.png'), out_img)
+    return

spikezoo/archs/ssml/cbam.py

@@ -0,0 +1,224 @@
+import torch
+import math
+import torch.nn as nn
+import torch.nn.functional as F
+
+class crop(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        N, C, H, W = x.shape
+        x = x[0:N, 0:C, 0:H-1, 0:W]
+        return x
+
+class shift(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.shift_down = nn.ZeroPad2d((0,0,1,0))
+        self.crop = crop()
+
+    def forward(self, x):
+        x = self.shift_down(x)
+        x = self.crop(x)
+        return x
+
+class Conv(nn.Module):
+    def __init__(self, in_channels, out_channels, bias=False, blind=True,stride=1,padding=0,kernel_size=3):
+        super().__init__()
+        self.blind = blind
+        if blind:
+            self.shift_down = nn.ZeroPad2d((0,0,1,0))
+            self.crop = crop()
+        self.replicate = nn.ReplicationPad2d(1)
+#         self.conv = nn.Conv2d(in_channels, out_channels, 3, bias=bias)
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, stride=stride,padding=padding,bias=bias) 
+        self.relu = nn.LeakyReLU(0.1, inplace=True)
+#         self.ln = nn.GroupNorm(1,out_channels)
+
+    def forward(self, x):
+        if self.blind:
+            x = self.shift_down(x)
+        x = self.replicate(x)
+        x = self.conv(x)
+        x = self.relu(x)
+        
+        if self.blind:
+            x = self.crop(x)
+        return x
+
+class BasicConv(nn.Module):
+    def __init__(self, in_planes, out_planes, kernel_size, stride=1, padding=0, dilation=1, groups=1, relu=False, bn=False, bias=True,blind=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 Flatten(nn.Module):
+    def forward(self, x):
+        return x.view(x.size(0), -1)
+
+class ChannelGate(nn.Module):
+    def __init__(self, gate_channels, reduction_ratio=16, pool_types=['avg', 'max']):
+        super(ChannelGate, self).__init__()
+        self.gate_channels = gate_channels
+        self.mlp = nn.Sequential(
+            Flatten(),
+            nn.Linear(gate_channels, gate_channels // reduction_ratio),
+            nn.ReLU(),
+            nn.Linear(gate_channels // reduction_ratio, gate_channels)
+            )
+        self.pool_types = pool_types
+    def forward(self, x):
+        channel_att_sum = None
+        for pool_type in self.pool_types:
+            if pool_type=='avg':
+                avg_pool = F.avg_pool2d( x, (x.size(2), x.size(3)), stride=(x.size(2), x.size(3)))
+                channel_att_raw = self.mlp( avg_pool )
+            elif pool_type=='max':
+                max_pool = F.max_pool2d( x, (x.size(2), x.size(3)), stride=(x.size(2), x.size(3)))
+                channel_att_raw = self.mlp( max_pool )
+            elif pool_type=='lp':
+                lp_pool = F.lp_pool2d( x, 2, (x.size(2), x.size(3)), stride=(x.size(2), x.size(3)))
+                channel_att_raw = self.mlp( lp_pool )
+            elif pool_type=='lse':
+                # LSE pool only
+                lse_pool = logsumexp_2d(x)
+                channel_att_raw = self.mlp( lse_pool )
+
+            if channel_att_sum is None:
+                channel_att_sum = channel_att_raw
+            else:
+                channel_att_sum = channel_att_sum + channel_att_raw
+
+        scale = F.sigmoid( channel_att_sum ).unsqueeze(2).unsqueeze(3).expand_as(x)
+        return x * scale
+
+def logsumexp_2d(tensor):
+    tensor_flatten = tensor.view(tensor.size(0), tensor.size(1), -1)
+    s, _ = torch.max(tensor_flatten, dim=2, keepdim=True)
+    outputs = s + (tensor_flatten - s).exp().sum(dim=2, keepdim=True).log()
+    return outputs
+
+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 SpatialGate(nn.Module):
+    def __init__(self,bias=False,blind=False):
+        super(SpatialGate, self).__init__()
+        kernel_size = 7
+        self.compress = ChannelPool()
+        self.spatial = BasicConv(2, 1, kernel_size, stride=1, padding=(kernel_size-1) // 2, relu=False,bias=bias,blind=False)
+    def forward(self, x):
+        x_compress = self.compress(x)
+        x_out = self.spatial(x_compress)
+        scale = F.sigmoid(x_out) # broadcasting
+        return x * scale
+
+class CBAM(nn.Module):
+    def __init__(self, gate_channels, reduction_ratio=16, pool_types=['avg', 'max'], no_spatial=False):
+        super(CBAM, self).__init__()
+        self.ChannelGate = ChannelGate(gate_channels, reduction_ratio, pool_types)
+        self.no_spatial=no_spatial
+        if not no_spatial:
+            self.SpatialGate = SpatialGate()
+    def forward(self, x):
+        x_out = self.ChannelGate(x)
+        if not self.no_spatial:
+            x_out = self.SpatialGate(x_out)
+        return x_out
+
+def weights_init_rcan(m):
+    """
+    custom weights initialization called on netG and netD
+    https://github.com/pytorch/examples/blob/master/dcgan/main.py
+    """
+    classname = m.__class__.__name__
+    if classname.find('Conv') != -1:
+        if classname.find('BasicConv') != -1:
+            m.conv.weight.data.normal_(0.0, 0.02)
+            if m.bn != None:
+                m.bn.bias.data.fill_(0)
+        else:
+            m.weight.data.normal_(0.0, 0.02)
+    elif classname.find('BatchNorm') != -1:
+        m.weight.data.normal_(1.0, 0.02)
+        m.bias.data.fill_(0)
+        
+class Temporal_Fusion(nn.Module):
+
+    def __init__(self, nf=64, nframes=3, center=1,bias=False):
+        super(Temporal_Fusion, self).__init__()
+        self.center = center
+
+        self.tAtt_1 = nn.Conv2d(nf, nf, 3, 1, 1, bias=bias)
+        self.tAtt_2 = nn.Conv2d(nf, nf, 3, 1, 1, bias=bias)
+
+        self.fea_fusion = nn.Conv2d(nframes * nf, nf, 1, 1, bias=bias)
+
+        self.sAtt_1 = nn.Conv2d(nframes * nf, nf, 1, 1, bias=bias)
+        self.maxpool = nn.MaxPool2d(3, stride=2, padding=1)
+        self.avgpool = nn.AvgPool2d(3, stride=2, padding=1)
+        self.sAtt_2 = nn.Conv2d(nf * 2, nf, 1, 1, bias=bias)
+        self.sAtt_3 = nn.Conv2d(nf, nf, 3, 1, 1, bias=bias)
+        self.sAtt_4 = nn.Conv2d(nf, nf, 1, 1, bias=bias)
+        self.sAtt_5 = nn.Conv2d(nf, nf, 3, 1, 1, bias=bias)
+        self.sAtt_L1 = nn.Conv2d(nf, nf, 1, 1, bias=bias)
+        self.sAtt_L2 = nn.Conv2d(nf * 2, nf, 3, 1, 1, bias=bias)
+        self.sAtt_L3 = nn.Conv2d(nf, nf, 3, 1, 1, bias=bias)
+        self.sAtt_add_1 = nn.Conv2d(nf, nf, 1, 1, bias=bias)
+        self.sAtt_add_2 = nn.Conv2d(nf, nf, 1, 1, bias=bias)
+
+        self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=True)
+
+    def forward(self, nonlocal_fea):
+        B, N, C, H, W = nonlocal_fea.size()  
+
+        emb_ref = self.tAtt_2(nonlocal_fea[:, self.center, :, :, :].clone())
+        emb = self.tAtt_1(nonlocal_fea.view(-1, C, H, W)).view(B, N, -1, H, W)  
+
+        cor_l = []
+        for i in range(N):
+            emb_nbr = emb[:, i, :, :, :]
+            cor_tmp = torch.sum(emb_nbr * emb_ref, 1).unsqueeze(1) 
+            cor_l.append(cor_tmp)
+        cor_prob = torch.sigmoid(torch.cat(cor_l, dim=1))  
+        cor_prob = cor_prob.unsqueeze(2).repeat(1, 1, C, 1, 1)
+        cor_prob = cor_prob.view(B, -1, H, W)
+        nonlocal_fea = nonlocal_fea.view(B, -1, H, W) * cor_prob
+
+        fea = self.lrelu(self.fea_fusion(nonlocal_fea))
+
+        att = self.lrelu(self.sAtt_1(nonlocal_fea))
+        att_max = self.maxpool(att)
+        att_avg = self.avgpool(att)
+        att = self.lrelu(self.sAtt_2(torch.cat([att_max, att_avg], dim=1)))
+
+        att_L = self.lrelu(self.sAtt_L1(att))
+        att_max = self.maxpool(att_L)
+        att_avg = self.avgpool(att_L)
+        att_L = self.lrelu(self.sAtt_L2(torch.cat([att_max, att_avg], dim=1)))
+        att_L = self.lrelu(self.sAtt_L3(att_L))
+        att_L = F.interpolate(att_L, scale_factor=2, mode='bilinear', align_corners=False)
+
+        att = self.lrelu(self.sAtt_3(att))
+        att = att + att_L
+        att = self.lrelu(self.sAtt_4(att))
+        att = F.interpolate(att, scale_factor=2, mode='bilinear', align_corners=False)
+        att = self.sAtt_5(att)
+        att_add = self.sAtt_add_2(self.lrelu(self.sAtt_add_1(att)))
+        att = torch.sigmoid(att)
+
+        fea = fea * att * 2 + att_add
+
+        return fea

spikezoo/archs/ssml/model.py

@@ -0,0 +1,290 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+# from utils import *
+import numpy as np
+import os
+import sys
+current_dir = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(current_dir)
+
+from cbam import SpatialGate,ChannelGate,Temporal_Fusion
+
+class crop(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        N, C, H, W = x.shape
+        x = x[0:N, 0:C, 0:H-1, 0:W]
+        return x
+
+class shift(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.shift_down = nn.ZeroPad2d((0,0,1,0))
+        self.crop = crop()
+
+    def forward(self, x):
+        x = self.shift_down(x)
+        x = self.crop(x)
+        return x
+
+class Conv(nn.Module):
+    def __init__(self, in_channels, out_channels, bias=False, blind=True,stride=1,padding=0,kernel_size=3):
+        super().__init__()
+        self.blind = blind
+        if blind:
+            self.shift_down = nn.ZeroPad2d((0,0,1,0))
+            self.crop = crop()
+        self.replicate = nn.ReplicationPad2d(1)
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, stride=stride,padding=padding,bias=bias) 
+        self.relu = nn.LeakyReLU(0.1, inplace=True)
+
+
+    def forward(self, x):
+        if self.blind:
+            x = self.shift_down(x)
+        x = self.replicate(x)
+        x = self.conv(x)
+        x = self.relu(x)        
+        if self.blind:
+            x = self.crop(x)
+        return x
+
+class Pool(nn.Module):
+    def __init__(self, blind=True):
+        super().__init__()
+        self.blind = blind
+        if blind:
+            self.shift = shift()
+        self.pool = nn.MaxPool2d(2)
+
+    def forward(self, x):
+        if self.blind:
+            x = self.shift(x)
+        x = self.pool(x)
+        return x
+
+class rotate(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        x90 = x.transpose(2,3).flip(3)
+        x180 = x.flip(2).flip(3)
+        x270 = x.transpose(2,3).flip(2)
+        x = torch.cat((x,x90,x180,x270), dim=0)
+        return x
+
+class unrotate(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        x0, x90, x180, x270 = torch.chunk(x, 4, dim=0)
+        x90 = x90.transpose(2,3).flip(2)
+        x180 = x180.flip(2).flip(3)
+        x270 = x270.transpose(2,3).flip(3)
+        x = torch.cat((x0,x90,x180,x270), dim=1)
+        return x
+
+class ENC_Conv(nn.Module):
+    def __init__(self, in_channels, mid_channels, out_channels, bias=False, reduce=True, blind=True):
+        super().__init__()
+        self.reduce = reduce
+        self.conv1 = Conv(in_channels, mid_channels, bias=bias, blind=blind)
+        self.conv2 = Conv(mid_channels, mid_channels, bias=bias, blind=blind)
+        self.conv3 = Conv(mid_channels, out_channels, bias=bias, blind=blind)
+        if reduce:
+            self.pool = Pool(blind=blind)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.conv2(x)
+        x = self.conv3(x)
+        if self.reduce:
+            x = self.pool(x)
+        return x
+
+class DEC_Conv(nn.Module):
+    def __init__(self, in_channels, mid_channels, out_channels, bias=False, blind=True):
+        super().__init__()
+        self.upsample = nn.Upsample(scale_factor=2, mode='nearest')
+        self.conv1 = Conv(in_channels, mid_channels, bias=bias, blind=blind)
+        self.conv2 = Conv(mid_channels, mid_channels, bias=bias, blind=blind)
+        self.conv3 = Conv(mid_channels, mid_channels, bias=bias, blind=blind)
+        self.conv4 = Conv(mid_channels, out_channels, bias=bias, blind=blind)
+
+    def forward(self, x, x_in):
+        x = self.upsample(x)
+
+        # Smart Padding
+        diffY = x_in.size()[2] - x.size()[2]
+        diffX = x_in.size()[3] - x.size()[3]
+        x = F.pad(x, [diffX // 2, diffX - diffX // 2,
+                      diffY // 2, diffY - diffY // 2])
+
+        x = torch.cat((x, x_in), dim=1)
+        x = self.conv1(x)
+        x = self.conv2(x)
+        x = self.conv3(x)
+        x = self.conv4(x)
+        return x
+
+class Blind_UNet(nn.Module):
+    def __init__(self, n_channels=3, n_output=96, bias=False, blind=True):
+        super().__init__()
+        self.n_channels = n_channels
+        self.bias = bias
+        self.enc1 = ENC_Conv(n_channels, 48, 48, bias=bias, blind=blind)
+        self.enc2 = ENC_Conv(48, 48, 48, bias=bias, blind=blind)
+        self.enc3 = ENC_Conv(48, 96, 48, bias=bias, reduce=False, blind=blind)
+        self.dec2 = DEC_Conv(96, 96, 96, bias=bias, blind=blind)
+        self.dec1 = DEC_Conv(96+n_channels, 96, n_output, bias=bias, blind=blind)
+
+    def forward(self, input):
+        x1 = self.enc1(input)
+        x2 = self.enc2(x1)
+        x = self.enc3(x2)
+        x = self.dec2(x, x1)
+        x = self.dec1(x, input)
+        return x
+
+def middleTFI(spike, middle, window=50):
+    #左右找1
+    spike = spike.squeeze(1).numpy() 
+    C, H, W = spike.shape
+    lindex, rindex = np.zeros([H, W]), np.zeros([H, W])
+    l, r = middle+1, middle+1
+    for r in range(middle+1, middle + window+1): #往左包括自己50个,往右不包括自己也是50个
+        l = l - 1
+        if l>=0:
+            newpos = spike[l, :, :]*(1 - np.sign(lindex)) 
+            distance = l*newpos
+            lindex += distance
+        if r<C:
+            newpos = spike[r, :, :]*(1 - np.sign(rindex))
+            distance = r*newpos
+            rindex += distance
+
+    rindex[rindex==0] = window+middle
+    lindex[lindex==0] = middle-window
+    interval = rindex - lindex
+    tfi = 1.0 / interval
+    
+    return tfi
+
+class MotionInference(nn.Module):
+    def __init__(self,n_frame=41,bias=False,blind=False):
+        super().__init__()
+        self.middle = n_frame//2
+        self.conv0 = nn.Conv2d(5*2+1,1,1,bias=bias)
+        self.conv1 = nn.Conv2d(9*2+1,1,1,bias=bias)
+        self.conv2 = nn.Conv2d(13*2+1,1,1,bias=bias)
+        self.tfpconv = Conv(in_channels=3, out_channels=16, bias=bias,blind=blind)
+        self.tficonv = Conv(in_channels=1, out_channels=16, bias=bias,blind=blind)
+        self.ChannelGate = ChannelGate(gate_channels=16, reduction_ratio=4)
+        self.SpatialGate = SpatialGate(bias=bias,blind=blind)
+        self.blind = blind
+    def forward(self, x):  
+        N, C, H, W = x.shape
+        tmp=[]
+        ttt=[]
+        for j in range(N):
+            tmp2 = middleTFI(x[j].cpu(), self.middle, window=12)    
+            tmp2 = torch.tensor(tmp2,dtype=torch.float32).unsqueeze_(dim=0)
+            tmp.append(tmp2) #1 40 40
+            ttt5=torch.mean(x[j,self.middle-3:self.middle+3+1,:,:].cpu(),dim=0).unsqueeze_(0) 
+            ttt.append(ttt5)
+        tfi_label = torch.stack(tmp,0).cuda()
+        tfp_label = torch.stack(ttt,0).cuda()
+
+#         tfi_label = (torch.clamp(tfi_label,0,1))**(1/2.2)
+#         tfp_label = (torch.clamp(tfp_label,0,1))**(1/2.2)
+        
+        tfp0 = self.conv0(x[:,self.middle-5:self.middle+5+1,:,:]) #b 1 h w，
+        tfp1 = self.conv1(x[:,self.middle-9:self.middle+9+1,:,:])
+        tfp2 = self.conv2(x[:,self.middle-13:self.middle+13+1,:,:])
+        tfps = torch.cat([tfp0,tfp1,tfp2],dim=1) #b 3 h w
+        
+        tfp_fea = self.tfpconv(tfps)        
+        tfi_fea = self.tficonv(tfi_label)
+        
+        if not self.blind:
+            tfp_fea = self.SpatialGate(tfp_fea) #b 16 h w
+            tfi_fea = self.SpatialGate(tfi_fea) 
+            fusion_fea = self.ChannelGate(tfp_fea+tfi_fea) #b 16 h w
+        else:
+            fusion_fea = tfp_fea+tfi_fea
+#         tfi_label = (torch.clamp(tfi_label,0,1))**(1/2.2)
+#         tfp_label = (torch.clamp(tfp_label,0,1))**(1/2.2)        
+        return fusion_fea,tfi_label,tfp_label
+    
+
+class BSN(nn.Module): 
+    def __init__(self, n_channels=3, n_output=3, bias=False, blind=True, sigma_known=True):
+        super().__init__()
+        self.n_channels = n_channels
+        self.c = n_channels
+        self.n_output = n_output
+        self.bias = bias
+        self.blind = blind
+        self.sigma_known = sigma_known
+        self.rotate = rotate()
+        self.unet = Blind_UNet(n_channels=n_channels+16, bias=bias, blind=blind) 
+        self.shift = shift()
+        self.unrotate = unrotate()
+        self.nin_A = nn.Conv2d(384, 384, 1, bias=bias)
+        self.nin_B = nn.Conv2d(384, 96, 1, bias=bias)
+        self.nin_C = nn.Conv2d(96, n_output, 1, bias=bias)
+        self.MotionInference = MotionInference(n_frame=41,bias=bias,blind=blind)
+    
+    def forward(self, x):  
+        N, C, H, W = x.shape   
+        _,tfi_label,tfp_label = self.MotionInference(x) 
+        if(H > W):
+            diff = H - W
+            x = F.pad(x, [diff // 2, diff - diff // 2, 0, 0], mode = 'reflect')
+        elif(W > H):
+            diff = W - H
+            x = F.pad(x, [0, 0, diff // 2, diff - diff // 2], mode = 'reflect')
+            
+        x = self.rotate(x)
+        
+        fea1,tfi,tfp = self.MotionInference(x)
+        x = torch.cat([x,fea1],1)
+        
+        x = self.unet(x)   #4 3 100 100 -> 4 96 100 100
+        if self.blind:
+            x = self.shift(x)
+        x = self.unrotate(x) #4 96 100 100 -> 1 384 100 100
+        
+        x0 = F.leaky_relu_(self.nin_A(x), negative_slope=0.1)
+        x0 = F.leaky_relu_(self.nin_B(x0), negative_slope=0.1)
+        x0 = self.nin_C(x0)
+
+        # Unsquare
+        if(H > W):
+            diff = H - W
+            x0 = x0[:, :, 0:H, (diff // 2):(diff // 2 + W)]
+        elif(W > H):
+            diff = W - H
+            x0 = x0[:, :, (diff // 2):(diff // 2 + H), 0:W]
+            
+        return x0,tfi_label,tfp_label   
+
+class DoubleNet(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.nbsn = BSN(n_channels=41, n_output=1,blind=False)
+        # self.bsn = BSN(n_channels=41, n_output=1,blind=True)
+    
+    def forward(self, x):  
+        out1,_,_ = self.nbsn(x)
+        
+        return out1
+
+if __name__ == '__main__':
+    a=DoubleNet().cuda()
+    print(a(torch.ones(2,41,40,40).cuda()))