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

spikezoo/archs/bsf/test.py

@@ -0,0 +1,16 @@
+import numpy as np
+import torch
+# import cupy as cp
+
+def compute_dsft_core(spike):
+    H, W, T = spike.shape
+    time = spike * torch.arange(T, device='cuda').reshape(1, 1, T)
+    l_idx, _ = time.cummax(dim=2)
+    time[time==0] = T
+    r_idx, _ = torch.flip(time, [2]).cummin(dim=2)
+    r_idx = torch.flip(r_idx, [2])
+    r_idx = torch.cat([r_idx[:, :, 1:], torch.ones([H, W, 1], device='cuda') * T], dim=2)
+    res = r_idx - l_idx
+    
+    res = torch.clip(res, 0)
+    return res

spikezoo/archs/bsf/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.uint8)
+    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/spikeclip/nets.py

@@ -0,0 +1,40 @@
+import torch.nn as nn
+import torch
+
+def conv_layer(inDim, outDim, ks, s, p, norm_layer='none'):
+    ## convolutional layer
+    conv = nn.Conv2d(inDim, outDim, kernel_size=ks, stride=s, padding=p)
+    relu = nn.ReLU(True)
+    assert norm_layer in ('batch', 'instance', 'none')
+    if norm_layer == 'none':
+        seq = nn.Sequential(*[conv, relu])
+    else:
+        if (norm_layer == 'instance'):
+            norm = nn.InstanceNorm2d(outDim, affine=False, track_running_stats=False) # instance norm
+        else:
+            momentum = 0.1
+            norm = nn.BatchNorm2d(outDim, momentum = momentum, affine=True, track_running_stats=True)
+        seq = nn.Sequential(*[conv, norm, relu])
+    return seq
+
+def LRN(inDim=50, outDim=1, norm='none'):  
+    convBlock1 = conv_layer(inDim,64,3,1,1)
+    convBlock2 = conv_layer(64,128,3,1,1,norm)
+    convBlock3 = conv_layer(128,64,3,1,1,norm)
+    convBlock4 = conv_layer(64,16,3,1,1,norm)
+    conv = nn.Conv2d(16, outDim, 3, 1, 1) 
+    seq = nn.Sequential(*[convBlock1, convBlock2, convBlock3, convBlock4, conv])
+    return seq
+
+    
+from thop import profile
+if __name__ == "__main__":
+    net = LRN()
+    total = sum(p.numel() for p in net.parameters())
+    spike = torch.zeros((1,50,250,400))
+    flops, _ = profile((net), inputs=(spike,))
+    re_msg = (
+        "Total params: %.4fM" % (total / 1e6),
+        "FLOPs=" + str(flops / 1e9) + '{}'.format("G"),
+    )    
+    print(re_msg)

spikezoo/archs/spikeformer/CheckPoints/readme

@@ -0,0 +1 @@
+This is a folder for saving the trained model !

spikezoo/archs/spikeformer/DataProcess/DataExtactor.py

@@ -0,0 +1,60 @@
+import os
+import numpy as np
+
+
+class DataExtractor():
+
+    def __init__(self, dataPath='', type='train'):
+
+        self.type = type
+        self.rootPath = dataPath
+
+    def GetData(self):
+
+        if self.type == "train":
+            return self.__GetTrainData()
+        if self.type == "valid":
+            return self.__GetValidData()
+        if self.type == "test":
+            return self.__GetTestData()
+
+
+    def __GetTrainData(self):
+
+        pathList = []
+
+        root = os.path.join(self.rootPath, 'train')
+        fileNames = os.listdir(root)
+        fileNames.sort()
+        for name in fileNames:
+            path = os.path.join(root, name)
+            pathList.append(path)
+
+        return pathList
+
+    def __GetValidData(self):
+
+        pathList = []
+
+        root = os.path.join(self.rootPath, 'valid')
+        fileNames = os.listdir(root)
+        fileNames.sort()
+        for name in fileNames:
+            path = os.path.join(root, name)
+            pathList.append(path)
+
+        return pathList
+
+    def __GetTestData(self):
+
+        pathList = []
+
+        root = os.path.join(self.rootPath, 'test')
+        fileNames = os.listdir(root)
+        fileNames.sort()
+        for name in fileNames:
+            path = os.path.join(root, name)
+            pathList.append(path)
+
+        return pathList
+

spikezoo/archs/spikeformer/DataProcess/DataLoader.py

@@ -0,0 +1,115 @@
+import os
+import torch
+# from torchvision import transforms
+from torch.utils import data
+import numpy as np
+from PIL import Image
+import cv2
+import random
+
+
+from DataProcess.DataExtactor import DataExtractor
+from DataProcess.LoadSpike import LoadSpike, load_spike_raw
+
+class Dataset(data.Dataset):
+
+    def __init__(self, pathList, dataType, spikeRadius):
+
+        self.pathList = pathList
+        self.dataType = dataType
+        self.spikeRadius = spikeRadius
+
+        #Random Rotation
+        if self.dataType == "train":
+            self.choice = [0, 1, 2, 3]
+        else:
+            self.choice = [0]
+
+    def __getitem__(self, index):
+
+        spSeq, gtFrames = self.GetItem(index)
+
+        return spSeq, gtFrames
+
+    def __len__(self):
+
+        return len(self.pathList)
+
+    def GetItem(self, index):
+
+        path = self.pathList[index]
+        spSeq, gtFrames = LoadSpike(path)
+
+        spLen, _, _ = spSeq.shape
+        gtLen, _, _ = gtFrames.shape
+        spCenter = spLen // 2
+        gtCenter = gtLen // 2
+
+        spLeft, spRight = (spCenter - self.spikeRadius,
+                           spCenter + self.spikeRadius)
+        spRight = spRight + 1
+        spSeq = spSeq[spLeft:spRight]
+
+        gtFrame = gtFrames[gtCenter]
+
+        spSeq = np.pad(spSeq, ((0, 0), (3, 3), (0, 0)), mode='constant')
+        spSeq = spSeq.astype(float) * 2 - 1
+
+        gtFrame = gtFrame.astype(float) / 255. * 2.0 - 1.
+
+
+        spSeq = torch.FloatTensor(spSeq)
+        gtFrame = torch.FloatTensor(gtFrame)
+
+        '''
+        Rotate the spike frame and Gt frame by ramdom degree, 
+        depending on the values of 'self.choice'
+        '''
+        # choice = random.choice(self.choice)
+        # spSeq = torch.rot90(spSeq, choice, dims=(1,2))
+        # gtFrame =torch.rot90(gtFrame, choice, dims=(1,2))
+        return spSeq, gtFrame
+
+
+
+
+
+class DataContainer():
+
+    def __init__(self, dataPath='', dataType='train',
+                 spikeRadius=16, batchSize=128, numWorks=0):
+
+        self.dataPath = dataPath
+        self.dataType = dataType
+        self.spikeRadius = spikeRadius
+        self.batchSize = batchSize
+        self.numWorks = numWorks
+
+        self.__GetData()
+
+    def __GetData(self):
+
+        dataset = None
+
+        dataset = DataExtractor(dataPath=self.dataPath, type=self.dataType)
+        self.pathList = dataset.GetData()
+
+    def GetLoader(self):
+
+        dataset = Dataset(self.pathList, self.dataType, self.spikeRadius)
+        dataLoader = None
+        if self.dataType == "train":
+            dataLoader = data.DataLoader(dataset, batch_size=self.batchSize, shuffle=True,
+                                         num_workers=self.numWorks, pin_memory=False)
+        else:
+            dataLoader = data.DataLoader(dataset, batch_size=self.batchSize, shuffle=False,
+                                         num_workers=self.numWorks, pin_memory=False)
+
+        return dataLoader
+
+if __name__ == "__main__":
+
+    pass
+
+
+

spikezoo/archs/spikeformer/DataProcess/LoadSpike.py

@@ -0,0 +1,39 @@
+import numpy as np
+
+def load_spike_numpy(path: str) -> (np.ndarray, np.ndarray):
+    '''
+    Load a spike sequence with it's tag from prepacked `.npz` file.\n
+    The sequence is of shape (`length`, `height`, `width`) and tag of
+        shape (`height`, `width`).
+    '''
+    data = np.load(path)
+    seq, tag, length = data['seq'], data['tag'], int(data['length'])
+    seq = np.array([(seq[i // 8] >> (i & 7)) & 1 for i in range(length)])
+    return seq, tag
+
+def LoadSpike(path: str) -> (np.ndarray, np.ndarray):
+    '''
+    Load a spike sequence,  the corresponding ground-truth frame sequence,
+    and sequence length.
+    spSeq: an ndarray of shape('sequence number', 'height', 'width')
+    gtFrames: an ndarray of shape('sequence length', 'height', 'width')
+    '''
+    data = np.load(path)
+    spSeq, gtFrames, length = data['spSeq'], data['gt'], int(data['length'])
+    spSeq = np.array([(spSeq[i // 8] >> (i & 7)) & 1 for i in range(length)])
+    return spSeq, gtFrames
+
+def load_spike_raw(path: str, width=400, height=250) -> np.ndarray:
+    '''
+    Load bit-compact raw spike data into an ndarray of shape
+        (`sequence length`, `height`, `width`).
+    '''
+    with open(path, 'rb') as f:
+        fbytes = f.read()
+    fnum = (len(fbytes) * 8) // (width * height)  # number of frames
+    frames = np.frombuffer(fbytes, dtype=np.uint8)
+    frames = np.array([frames & (1 << i) for i in range(8)])
+    frames = frames.astype(np.bool).astype(np.uint8)
+    frames = frames.transpose(1, 0).reshape(fnum, height, width)
+    frames = np.flip(frames, 1)
+    return frames

spikezoo/archs/spikeformer/EvalResults/readme

@@ -0,0 +1 @@
+This is a folder for saving the images reconstructed from validation/testing set !

spikezoo/archs/spikeformer/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2022 YangChenUcas
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

spikezoo/archs/spikeformer/Metrics/Metrics.py

@@ -0,0 +1,50 @@
+import numpy as np
+from skimage import metrics
+
+class Metrics():
+
+    def __init__(self):
+        self.best_psnr = 0.
+        self.best_ssim = 0.
+        self.best_niqe = 0.
+
+    def Update(self, psnr=0., ssim=0., niqe=0.):
+        self.best_psnr = psnr
+        self.best_ssim = ssim
+        self.best_niqe = niqe
+
+    def GetBestMetrics(self):
+
+        return self.best_psnr, self.best_ssim, self.best_niqe
+
+    def Cal_PSNR(self, preImgs, gtImgs): #shape:[B, H, W]
+
+        B, _, _ = preImgs.shape
+        total_psnr = 0.
+        for i, (pre, gt) in enumerate(zip(preImgs, gtImgs)):
+            print(i+1, metrics.peak_signal_noise_ratio(gt, pre))
+            total_psnr += metrics.peak_signal_noise_ratio(gt, pre)
+
+        avg_psnr = total_psnr / B
+
+        return avg_psnr
+
+    def Cal_SSIM(self, preImgs, gtImgs): #shape:[B, H, W]
+
+        B, _, _ = preImgs.shape
+        total_ssim = 0.
+        for i, (pre, gt) in enumerate(zip(preImgs, gtImgs)):
+            total_ssim += metrics.structural_similarity(pre, gt)
+
+        avg_ssim = total_ssim / B
+
+        return avg_ssim
+
+
+if __name__ == "__main__":
+
+    a = np.random.random((2,256,256))
+    b = np.random.random((2,256,256))
+    metrics = Metrics()
+
+    print(metrics.Cal_NIQE(a))

spikezoo/archs/spikeformer/Model/Loss.py

@@ -0,0 +1,89 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class CharbonnierLoss(nn.Module):
+    """Charbonnier Loss (L1)"""
+
+    def __init__(self, eps=1e-3):
+        super(CharbonnierLoss, self).__init__()
+        self.eps = eps
+
+    def forward(self, x, y):
+        diff = x - y
+        # loss = torch.sum(torch.sqrt(diff * diff + self.eps))
+        loss = torch.mean(torch.sqrt((diff * diff) + (self.eps*self.eps)))
+        return loss
+
+class EdgeLoss(nn.Module):
+    def __init__(self):
+        super(EdgeLoss, self).__init__()
+        k = torch.Tensor([[.05, .25, .4, .25, .05]])
+        # self.kernel = torch.matmul(k.t(),k).unsqueeze(0).repeat(3,1,1,1)
+        self.kernel = torch.matmul(k.t(),k).unsqueeze(0).unsqueeze(0).repeat(1,3,1,1) #这个的repeat也是后加的
+        # print(self.kernel.shape)
+        if torch.cuda.is_available():
+            self.kernel = self.kernel.cuda()
+        self.loss = CharbonnierLoss()
+
+    def conv_gauss(self, img):
+        # print('aaaa')
+        # print(img.shape)
+        n_channels, _, kw, kh = self.kernel.shape
+        img = F.pad(img, (kw//2, kh//2, kw//2, kh//2), mode='replicate')
+        # return F.conv2d(img, self.kernel, groups=n_channels)
+        return F.conv2d(img, self.kernel)
+
+    def laplacian_kernel(self, current):
+        filtered    = self.conv_gauss(current)    # filter
+        down        = filtered[:,:,::2,::2]               # downsample
+        new_filter  = torch.zeros_like(filtered)
+        new_filter[:,:,::2,::2] = down*4                  # upsample
+        filtered    = self.conv_gauss(new_filter.repeat(1,3,1,1)) # filter  #这里为什么需要repeat一下？原文的目的是什么？否则不能正常运行
+        diff = current - filtered
+        return diff
+
+    def forward(self, x, y):
+        y = y.repeat(1,3,1,1)
+        x = x.repeat(1,3,1,1)
+        # print('bbbbbb')
+        # print(x.shape)
+        # print(y.shape)
+        loss = self.loss(self.laplacian_kernel(x), self.laplacian_kernel(y))
+        return loss
+
+
+class VGGLoss4(nn.Module):
+    def __init__(self, path: str):
+        super().__init__()
+        self.features = nn.Sequential(
+            nn.Conv2d(3, 64, 3, 1, 1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(64, 64, 3, 1, 1),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(2),
+            nn.Conv2d(64, 128, 3, 1, 1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(128, 128, 3, 1, 1),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(2),
+            nn.Conv2d(128, 256, 3, 1, 1),
+            nn.ReLU(inplace=True),
+            # nn.Conv2d(256, 256, 3, 1, 1),
+            # nn.ReLU(inplace=True),
+            # nn.Conv2d(256, 256, 3, 1, 1),
+            # nn.ReLU(inplace=True),
+        )
+        self.load_state_dict(torch.load(path))
+        for param in self.parameters():
+            param.requires_grad = False
+
+    def forward(self, real_y, fake_y):
+        real_y = real_y.repeat((1, 3, 1, 1))
+        fake_y = fake_y.repeat((1, 3, 1, 1))
+        with torch.no_grad():
+            real_f = self.features(real_y)
+        fake_f = self.features(fake_y)
+        return F.mse_loss(real_f, fake_f)
+
+