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

spikezoo/archs/wgse/README.md

@@ -0,0 +1,64 @@
+# WGSE-SpikeCamera
+
+Codes and Datasets of **"Learning Temporal-ordered Representation for Spike Streams Based on Discrete Wavelet Transforms"**.
+
+Jiyuan Zhang*, Shanshan Jia*, Zhaofei Yu $\dagger$ and Tiejun Huang
+
+ Thirty-Seventh AAAI Conference on Artificial Intelligence (**AAAI 2023**) .
+ __________________________________________________
+ ## Code for WGSE
+
+* **Key Requirements of the Code Environment**
+
+  * Pytorch > 1.7.0
+  * pywavelets
+  * pytorch-wavelets
+  * scikit-image
+  * einops
+
+* **Train**
+
+To train the network for WGSE, use the file `train.py`
+
+the parameter list is:
+
+```python
+parser = argparse.ArgumentParser(description='AAAI - WGSE - REDS')
+parser.add_argument('-c', '--cuda', type=str, default='1', help='select gpu card')
+parser.add_argument('-b', '--batch_size', type=int, default=16)
+parser.add_argument('-e', '--epoch', type=int, default=600)
+parser.add_argument('-w', '--wvl', type=str, default='db8', help='select wavelet base function')
+parser.add_argument('-j', '--jlevels', type=int, default=5)
+parser.add_argument('-k', '--kernel_size', type=int, default=3)
+parser.add_argument('-l', '--logpath', type=str, default='WGSE-Dwt1dNet')
+parser.add_argument('-r', '--resume_from', type=str, default=None)
+parser.add_argument('--dataroot', type=str, default=None)
+```
+
+`-c` is the CUDA device index on your computer, `-b` is the batchsize, `-e` is the number of epoch of training, `-w` is the wavelet function, `-j` is the decompostion level, `-k` is the kernel size in the WGSE, `-r` is the folder where saving the weights that you want to load and resume training, `--dataroot` is the root path of the dataset and `-l` is the folder path where you want to save the log files and weight file of the model.
+
+In our implementation, the training script should be:
+
+``` bash
+python train.py -c 0 -b 16 -e 600 --dataroot "rootpath/of/dataset"  -l "folder/for/save/logfiles"
+```
+
+* **Test**
+
+To test the network for WGSE, use the file `demo.py`
+
+the parameter list is:
+
+```python
+parser = argparse.ArgumentParser(description='AAAI - WGSE - REDS')
+parser.add_argument('-w', '--wvl', type=str, default='db8', help='select wavelet base function')
+parser.add_argument('-j', '--jlevels', type=int, default=5)
+parser.add_argument('-k', '--kernel_size', type=int, default=3)
+parser.add_argument('-l', '--logpath', type=str, default='WGSE-Dwt1dNet')
+parser.add_argument('-f', '--datfile', type=str, default=None, help='path of the spike data to be tested')
+```
+
+`-w` is the wavelet function, `-j` is the decompostion level, `-k` is the kernel size in the WGSE, `-l` is the folder path where you save the weight file, `-f` is the `.dat` data path of spikes that you want to test.
+
+## The Dataset for WGSE
+In this work, we propose a synthetic dataset "Spike-Cityscapes" for semantic segmentation based on the spike streams generated from [Cityscapes](https://www.cityscapes-dataset.com/). The spike data is available and you can download them at [https://pan.baidu.com/s/1lB4qpfZwaVN6WDFo5MRR-w](https://pan.baidu.com/s/1lB4qpfZwaVN6WDFo5MRR-w) with the password **svpg**.

spikezoo/archs/wgse/dataset.py

@@ -0,0 +1,59 @@
+import glob
+import os
+
+import numpy as np
+import torch
+import torch.utils.data as data
+from skimage import io
+
+import utils
+
+
+class DatasetREDS(data.Dataset):
+    def __init__(self, cfg, transform=None):
+        super(DatasetREDS, self).__init__()
+        self.cfg = cfg
+        self.rootfolder = cfg['rootfolder']
+        self.spikefolder = os.path.join(self.rootfolder, cfg['spikefolder'])
+        self.imagefolder = os.path.join(self.rootfolder, cfg['imagefolder'])
+        
+        self.spike_list = os.listdir(self.spikefolder)
+        self.image_list = os.listdir(self.imagefolder)
+        
+        self.H = int(cfg['H'])
+        self.W = int(cfg['W'])
+        self.C_SPIKE = int(cfg['C'])
+
+        self.transform = transform
+
+    def __getitem__(self, index: int):
+        item_name = self.spike_list[index][:-4]
+        spike_path = os.path.join(self.spikefolder, item_name+'.dat')
+        image_path = os.path.join(self.imagefolder, item_name+'.png')
+
+        path_grayframe = glob.glob(image_path)
+        gray_frame = io.imread(path_grayframe[0], as_gray=False).astype(np.float32)
+
+        gray_frame /= 255.0  # normalize
+        gray_frame = np.expand_dims(gray_frame, axis=0)  # expand to [1 x H x W]
+        gray_frame = torch.from_numpy(gray_frame)
+
+        f = open(spike_path, 'rb')
+        spike_seq = f.read()
+        spike_seq = np.frombuffer(spike_seq, 'b')
+        spikes = utils.RawToSpike(spike_seq, self.H, self.W)
+        spikes = spikes.astype(np.float32)
+        spikes = torch.from_numpy(spikes)
+        f.close()
+
+        if self.transform:
+            gray_frame, spikes = self.transform(gray_frame, spikes)
+
+        item = {}
+        item['spikes'] = spikes
+        item['image'] = gray_frame
+        
+        return item
+
+    def __len__(self) -> int:
+        return len(self.spike_list)

spikezoo/archs/wgse/demo.py

@@ -0,0 +1,83 @@
+import os
+
+import numpy as np
+import cv2
+import torch
+import argparse
+
+from dwtnets import Dwt1dResnetX_TCN
+from utils import RawToSpike
+from einops import rearrange
+from pytorch_wavelets import DWT1DForward
+
+
+parser = argparse.ArgumentParser(description='AAAI - WGSE - REDS')
+parser.add_argument('-w', '--wvl', type=str, default='db8', help='select wavelet base function')
+parser.add_argument('-j', '--jlevels', type=int, default=5)
+parser.add_argument('-k', '--kernel_size', type=int, default=3)
+parser.add_argument('-l', '--logpath', type=str, default='WGSE-Dwt1dNet')
+parser.add_argument('-f', '--datfile', type=str, default="/home/chenkang455/chenk/myproject/SpikeCLIP/selfcode/data/recVidarReal2019/classB/train-350kmh.dat", help='path of the spike data to be tested')
+
+args = parser.parse_args()
+dataroot = args.datfile
+wvlname = args.wvl
+j = args.jlevels
+logfolder = args.logpath
+ks = args.kernel_size
+
+
+def nor(x):
+    return (x-np.min(x))/(np.max(x)-np.min(x))
+
+def ensure_dir(path):
+    if not os.path.exists(path):
+        os.makedirs(path)
+
+def progress_bar_time(total_time):
+    hour = int(total_time) // 3600
+    minu = (int(total_time) % 3600) // 60
+    sec = int(total_time) % 60
+    return '%d:%02d:%02d' % (hour, minu, sec)
+
+
+def main():
+
+    f = open(dataroot, 'rb')
+    spike_seq = f.read()
+    spike_seq = np.frombuffer(spike_seq, 'b')
+    spikes = RawToSpike(spike_seq, 250, 400)
+    spikes = spikes.astype(np.float32)
+    spikes = torch.from_numpy(spikes)
+    f.close()
+
+    spikes = spikes[None, 130:171, :, :]
+
+    s = spikes[:, :, 0:1, 0:1]
+    dwt = DWT1DForward(wave=wvlname, J=j)
+    s_r = rearrange(s, 'b t h w -> b h w t')
+    s_r = rearrange(s_r, 'b h w t -> (b h w) 1 t')
+    yl, yh = dwt(s_r)
+    yl_size = yl.shape[-1]
+    yh_size = [yhi.shape[-1] for yhi in yh]
+    print(yl_size,yh_size)
+    model = Dwt1dResnetX_TCN(
+        wvlname=wvlname, J=j, yl_size=yl_size, yh_size=yh_size, num_residual_blocks=3, norm=None, ks=ks, store_features=True
+    )
+    print(model)
+
+    saved_state_dict = torch.load('model_best.pt')
+    model.load_state_dict(saved_state_dict.module.state_dict())
+    
+    model = model.cuda()
+    model.eval()
+
+    pred = model(spikes.cuda())        
+    prediction = pred[0].permute(1,2,0).cpu().detach().numpy()
+    print(os.path.join(logfolder, 'demo.png'))
+    cv2.imwrite( 'demo.png', prediction * 255.0)
+    
+
+if __name__ == '__main__':
+    main()
+
+    

spikezoo/archs/wgse/dwtnets.py

@@ -0,0 +1,145 @@
+import torch.nn as nn
+from einops import rearrange
+import os
+import sys
+current_dir = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(current_dir)
+
+from pytorch_wavelets import DWT1DForward, DWT1DInverse
+
+from submodules import ResidualBlock
+
+
+class TcnResidualLayer(nn.Module):
+    def __init__(self, in_c, out_c, dilated=1, k=3, s=1, p=1, store_features=False):
+        super().__init__()
+        self.tcn0 = nn.Sequential(
+            nn.Conv1d(in_c, out_c, kernel_size=k, stride=s, padding=p, dilation=dilated),
+            nn.ReLU(),
+        )
+        self.tcn1 = nn.Sequential(
+            nn.Conv1d(out_c, out_c, kernel_size=k, stride=s, padding=p, dilation=dilated),
+        )
+        self.relu = nn.ReLU(inplace=False)
+        self.store_features = store_features
+        self.features = {}
+
+    def forward(self, x):
+        residual = x
+        out = self.tcn0(x)
+        if self.store_features:
+            self.features['after_tcn0'] = out
+        out = self.tcn1(out)
+        out = out + residual
+        out = self.relu(out)
+        return out
+
+
+class Dwt1dModule_Tcn(nn.Module):
+    def __init__(
+        self, 
+        wvlname='db1', 
+        J=3, 
+        yl_size=14, 
+        yh_size=[26, 18, 14], 
+        ks = 3,
+        store_features=False
+    ):
+        super().__init__()
+        self.wvlname = wvlname
+        self.J = J
+        self.yl_num = yl_size
+        self.yh_num = yh_size
+        self.yh_blocks = nn.ModuleList()
+
+        self.store_features = store_features
+        self.features = {}
+
+        for i in self.yh_num:
+            self.yh_blocks.append(
+                nn.Sequential(
+                    TcnResidualLayer(1, 32, store_features=store_features, k=ks, p=ks//2),
+                    nn.Conv1d(32, 1, kernel_size=ks, padding=ks//2, dilation=1),
+                    nn.ReLU(),
+                )
+            )
+        self.yl_block = nn.Sequential(
+            TcnResidualLayer(1, 32, store_features=store_features, k=ks, p=ks//2),
+            nn.Conv1d(32, 1, kernel_size=ks, padding=ks//2, dilation=1),
+            nn.ReLU(),
+        )
+        self.dwt = DWT1DForward(wave=self.wvlname, J=self.J)
+        self.idwt = DWT1DInverse(wave=self.wvlname)
+
+    def forward(self, x):
+        B, T, H, W = x.shape
+        x_r = rearrange(x, 'b t h w -> b h w t')
+        x_r = rearrange(x_r, 'b h w t -> (b h w) 1 t')
+
+        yl, yh = self.dwt(x_r)
+        yl_out = self.yl_block(yl)
+        yh_out = []
+        for i, yhi in enumerate(yh):
+            yhi_out = self.yh_blocks[i](yhi)
+            yh_out.append(yhi_out)
+
+        out = self.idwt((yl_out, yh_out)) 
+        out = rearrange(out, '(b h w) 1 t -> b h w t', b=B, h=H, w=W)
+        out = rearrange(out, 'b h w t -> b t h w')
+
+        return out
+
+
+
+class Dwt1dResnetX_TCN(nn.Module):
+    def __init__(
+        self, 
+        wvlname='db1', 
+        J=3, 
+        yl_size=14, 
+        yh_size=[26, 18, 14], 
+        num_residual_blocks=2, 
+        norm=None, 
+        inc=41, 
+        ks=3,
+        store_features=False
+    ):
+        super().__init__()
+
+        self.wvl = Dwt1dModule_Tcn(wvlname, J, yl_size, yh_size, store_features=store_features, ks=ks)
+
+        self.norm = norm
+        self.num_residual_blocks = num_residual_blocks
+        self.resblocks = nn.ModuleList()
+        for _ in range(self.num_residual_blocks):
+            self.resblocks.append(ResidualBlock(256, 256, norm=self.norm))
+
+
+        self.conv = nn.Sequential(
+            nn.Conv2d(inc if inc%2==0 else inc+1, 256, kernel_size=3, padding=1, bias=True),
+            nn.ReLU(),
+        )
+
+
+        self.tail = nn.Sequential(
+            nn.Conv2d(256, 64, kernel_size=3, padding=1, bias=True),
+            nn.ReLU(),
+
+            nn.Conv2d(64, 1, kernel_size=3, padding=1, bias=True),
+            nn.ReLU(),
+        )
+
+        self.store_features = store_features
+        self.features = {}
+
+    def forward(self, x):
+        y = self.wvl(x)
+
+        y = self.conv(y)
+
+        for resi, resblock in enumerate(self.resblocks):
+            y = resblock(y)
+
+        out = self.tail(y)
+
+        return out

spikezoo/archs/wgse/eval.py

@@ -0,0 +1,133 @@
+import os
+import time
+
+import numpy as np
+import cv2
+import torch
+import argparse
+
+from dataset import DatasetREDS
+from dwtnets import Dwt1dResnetX_TCN
+from utils import calculate_psnr, calculate_ssim
+from einops import rearrange
+from pytorch_wavelets import DWT1DForward
+
+
+parser = argparse.ArgumentParser(description='AAAI - WGSE - REDS')
+parser.add_argument('-w', '--wvl', type=str, default='db8', help='select wavelet base function')
+parser.add_argument('-j', '--jlevels', type=int, default=5)
+parser.add_argument('-k', '--kernel_size', type=int, default=3)
+parser.add_argument('-l', '--logpath', type=str, default='WGSE-Dwt1dNet')
+parser.add_argument('--dataroot', type=str, default=None)
+
+args = parser.parse_args()
+dataroot = args.dataroot
+wvlname = args.wvl
+j = args.jlevels
+logfolder = args.logpath
+ks = args.kernel_size
+
+
+def nor(x):
+    return (x-np.min(x))/(np.max(x)-np.min(x))
+
+def save_wvl(wvl, savefolder, saveprefix):
+    wvl = wvl.squeeze()
+    t, h, w = wvl.shape
+    wvl = nor(wvl)
+    
+    for i in range(t):
+        wvl_t = wvl[i] * 255
+        cv2.imwrite(os.path.join(savefolder, saveprefix+'_{:03d}.png'.format(i)), wvl_t)
+
+def ensure_dir(path):
+    if not os.path.exists(path):
+        os.makedirs(path)
+
+def progress_bar_time(total_time):
+    hour = int(total_time) // 3600
+    minu = (int(total_time) % 3600) // 60
+    sec = int(total_time) % 60
+    return '%d:%02d:%02d' % (hour, minu, sec)
+
+
+def main():
+    cfg = {}
+    cfg['rootfolder'] = os.path.join(dataroot, 'val')
+    cfg['spikefolder'] = 'input'
+    cfg['imagefolder'] = 'gt'
+    cfg['H'] = 250
+    cfg['W'] = 400
+    cfg['C'] = 41
+    test_set = DatasetREDS(cfg)
+    
+    test_data_loader = torch.utils.data.DataLoader(
+        dataset=test_set,
+        batch_size=1,
+        shuffle=False,
+        num_workers=1,
+        drop_last=False)
+
+    item0 = test_set[0]
+    s = item0['spikes']
+    s = s[None, :, 0:1, 0:1]
+    dwt = DWT1DForward(wave=wvlname, J=j)
+    B, T, H, W = s.shape
+    s_r = rearrange(s, 'b t h w -> b h w t')
+    s_r = rearrange(s_r, 'b h w t -> (b h w) 1 t')
+    yl, yh = dwt(s_r)
+    yl_size = yl.shape[-1]
+    yh_size = [yhi.shape[-1] for yhi in yh]
+    
+    model = Dwt1dResnetX_TCN(
+        wvlname=wvlname, J=j, yl_size=yl_size, yh_size=yh_size, num_residual_blocks=3, norm=None, ks=ks, store_features=True
+    )
+    print(model)
+
+    saved_state_dict = torch.load(logfolder + '/model_best.pt')
+    model.load_state_dict(saved_state_dict.module.state_dict())
+    
+    model = model.cuda()
+    # model = torch.nn.DataParallel(model).cuda()
+
+    model.eval()
+    
+    with torch.no_grad():
+        sum_ssim = 0.0
+        sum_psnr = 0.0
+        sum_num = 0
+        total_time = 0
+        for i, item in enumerate(test_data_loader):
+            start_time = time.time()
+
+            spikes = item['spikes'][:, 130:171, :, :].cuda()
+            image = item['image'].cuda()
+
+            pred = model(spikes)
+            
+            prediction = pred[0].permute(1,2,0).cpu().numpy()
+            gt = image[0].permute(1,2,0).cpu().numpy()
+
+            sum_ssim += calculate_ssim(gt * 255.0, prediction * 255.0)
+            sum_psnr += calculate_psnr(gt * 255.0, prediction * 255.0)
+            sum_num += 1
+            elapse_time = time.time() - start_time
+            total_time += elapse_time
+
+            print('\r[evaluating] %3.2f%% | %6d/%6d [%s<%s, %.2fs/it]' % (
+                float(i + 1) / int(len(test_data_loader)) * 100, i + 1, int(len(test_data_loader)),
+                progress_bar_time(total_time),
+                progress_bar_time(total_time / (i + 1) * int(len(test_data_loader))),
+                total_time / (i + 1)), end='')
+                            
+        sum_psnr /= sum_num
+        sum_ssim /= sum_num
+
+    print('')
+    print('\r[Evaluation Result] PSNR: %.4f | SSIM: %.4f' % (sum_psnr, sum_ssim))
+
+
+if __name__ == '__main__':
+    main()
+
+    

spikezoo/archs/wgse/logs/WGSE-Dwt1dNet-db8-5-ks3/log.txt

@@ -0,0 +1,11 @@
+[training] 0.50% |      1/   200 [0:00:03<0:11:10, 3.35s/it] | LOSS: 0.4502 | LR: 0.0001 
+[training] 1.00% |      2/   200 [0:00:03<0:05:38, 1.69s/it] | LOSS: 0.3112 | LR: 0.0001 
+[training] 1.50% |      3/   200 [0:00:03<0:03:48, 1.14s/it] | LOSS: 0.2487 | LR: 0.0001 
+[training] 2.00% |      4/   200 [0:00:03<0:02:53, 0.87s/it] | LOSS: 0.1424 | LR: 0.0001 
+[training] 2.50% |      5/   200 [0:00:03<0:02:20, 0.70s/it] | LOSS: 0.2075 | LR: 0.0001 
+[training] 3.00% |      6/   200 [0:00:03<0:01:58, 0.59s/it] | LOSS: 0.2425 | LR: 0.0001 
+[training] 3.50% |      7/   200 [0:00:03<0:01:42, 0.51s/it] | LOSS: 0.2178 | LR: 0.0001 
+[training] 4.00% |      8/   200 [0:00:03<0:01:30, 0.45s/it] | LOSS: 0.1428 | LR: 0.0001 
+[training] 4.50% |      9/   200 [0:00:03<0:01:21, 0.41s/it] | LOSS: 0.1512 | LR: 0.0001 
+[training] 5.00% |     10/   200 [0:00:03<0:01:14, 0.37s/it] | LOSS: 0.1258 | LR: 0.0001 
+[training] 5.50% |     11/   200 [0:00:03<0:01:08, 0.34s/it] | LOSS: 0.1700 | LR: 0.0001 

spikezoo/archs/wgse/submodules.py

@@ -0,0 +1,68 @@
+import torch
+import torch.nn as nn
+
+
+class ConvLayer(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, activation='relu', norm=None,
+                 BN_momentum=0.1):
+        super(ConvLayer, self).__init__()
+
+        bias = False if norm == 'BN' else True
+        self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=bias)
+        if activation is not None:
+            self.activation = getattr(torch, activation)
+        else:
+            self.activation = None
+
+        self.norm = norm
+        if norm == 'BN':
+            self.norm_layer = nn.BatchNorm2d(out_channels, momentum=BN_momentum)
+        elif norm == 'IN':
+            self.norm_layer = nn.InstanceNorm2d(out_channels, track_running_stats=True)
+
+    def forward(self, x):
+        out = self.conv2d(x)
+
+        if self.norm in ['BN', 'IN']:
+            out = self.norm_layer(out)
+
+        if self.activation is not None:
+            out = self.activation(out)
+
+        return out
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, stride=1, downsample=None, norm=None,
+                 BN_momentum=0.1):
+        super(ResidualBlock, self).__init__()
+        bias = False if norm == 'BN' else True
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=bias)
+        self.norm = norm
+        if norm == 'BN':
+            self.bn1 = nn.BatchNorm2d(out_channels, momentum=BN_momentum)
+            self.bn2 = nn.BatchNorm2d(out_channels, momentum=BN_momentum)
+        elif norm == 'IN':
+            self.bn1 = nn.InstanceNorm2d(out_channels)
+            self.bn2 = nn.InstanceNorm2d(out_channels)
+
+        self.relu = nn.ReLU(inplace=False)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=bias)
+        self.downsample = downsample
+
+    def forward(self, x):
+        residual = x
+        out = self.conv1(x)
+        if self.norm in ['BN', 'IN']:
+            out = self.bn1(out)
+        out = self.relu(out)
+        out = self.conv2(out)
+        if self.norm in ['BN', 'IN']:
+            out = self.bn2(out)
+
+        if self.downsample:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+        return out