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

spikezoo/archs/spikeformer/Model/SpikeFormer.py

@@ -0,0 +1,230 @@
+'''
+Ref: To implenment SpikeFormer, we referred to the code of ”segefomer-pytorch” published on github
+     (link: https://github.com/lucidrains/segformer-pytorch.git)
+'''
+from math import sqrt
+from functools import partial
+import torch
+from torch import nn, einsum
+from einops import rearrange, reduce
+
+# helpers
+
+def exists(val):
+    return val is not None
+
+def cast_tuple(val, depth):
+    return val if isinstance(val, tuple) else (val,) * depth
+
+LayerNorm = partial(nn.InstanceNorm2d, affine = True)
+
+# classes
+
+class DsConv2d(nn.Module):
+    def __init__(self, dim_in, dim_out, kernel_size, padding, stride = 1, bias = True):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Conv2d(dim_in, dim_in, kernel_size = kernel_size, padding = padding, groups = dim_in, stride = stride, bias = bias),
+            nn.GELU(),
+            nn.Conv2d(dim_in, dim_out, kernel_size = 1, bias = bias),
+            nn.GELU(),
+        )
+    def forward(self, x):
+        return self.net(x)
+
+class PreNorm(nn.Module):
+    def __init__(self, dim, fn):
+        super().__init__()
+        self.fn = fn
+        self.norm = LayerNorm(dim)
+
+    def forward(self, x):
+        # return self.fn(x)
+        return self.fn(self.norm(x))
+
+class EfficientSelfAttention(nn.Module):
+    def __init__(
+        self,
+        *,
+        dim,
+        heads,
+        reduction_ratio
+    ):
+        super().__init__()
+        self.scale = (dim // heads) ** -0.5
+        self.heads = heads
+
+        self.to_q = nn.Conv2d(dim, dim, 1, bias = False)
+        self.to_kv = nn.Conv2d(dim, dim * 2, reduction_ratio, stride = reduction_ratio, bias = False)
+        self.to_out = nn.Sequential(
+            nn.Conv2d(dim, dim, 1, bias=False),
+            nn.GELU()
+        )
+
+    def forward(self, x):
+        h, w = x.shape[-2:]
+        heads = self.heads
+
+        q, k, v = (self.to_q(x), *self.to_kv(x).chunk(2, dim = 1))
+        q, k, v = map(lambda t: rearrange(t, 'b (h c) x y -> (b h) (x y) c', h = heads), (q, k, v))
+
+        sim = einsum('b i d, b j d -> b i j', q, k) * self.scale
+        attn = sim.softmax(dim = -1)
+
+        out = einsum('b i j, b j d -> b i d', attn, v)
+        out = rearrange(out, '(b h) (x y) c -> b (h c) x y', h = heads, x = h, y = w)
+        return self.to_out(out)
+
+class MixFeedForward(nn.Module):
+    def __init__(
+        self,
+        *,
+        dim,
+        expansion_factor
+    ):
+        super().__init__()
+        hidden_dim = dim * expansion_factor
+        self.net = nn.Sequential(
+            nn.Conv2d(dim, hidden_dim, 1),
+            nn.GELU(),
+            DsConv2d(hidden_dim, hidden_dim, 3, padding = 1),
+            # nn.GELU(),
+            nn.Conv2d(hidden_dim, dim, 1),
+            nn.GELU(),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+class MiT(nn.Module):
+    def __init__(
+        self,
+        *,
+        channels,
+        dims,
+        heads,
+        ff_expansion,
+        reduction_ratio,
+        num_layers
+    ):
+        super().__init__()
+        stage_kernel_stride_pad = ((7, 4, 3), (3, 2, 1), (3, 2, 1), (3, 2, 1))
+        channels_mod = channels/16
+        dims = (channels_mod, *dims)
+        dim_pairs = list(zip(dims[:-1], dims[1:]))
+
+        self.stages = nn.ModuleList([])
+
+        for (dim_in, dim_out), (kernel, stride, padding), num_layers, ff_expansion, heads, reduction_ratio in zip(dim_pairs, stage_kernel_stride_pad, num_layers, ff_expansion, heads, reduction_ratio):
+            get_overlap_patches = nn.Unfold(kernel, stride = stride, padding = padding)
+            # if count == 0:
+                # overlap_patch_embed = nn.Conv2d(int((dim_in * kernel ** 2)/16), dim_out, 1)
+            # else:
+            overlap_patch_embed = nn.Conv2d(int(dim_in * kernel ** 2), dim_out, 1)
+            # count+=1
+
+            layers = nn.ModuleList([])
+
+            for _ in range(num_layers):
+                layers.append(nn.ModuleList([
+                    PreNorm(dim_out, EfficientSelfAttention(dim = dim_out, heads = heads, reduction_ratio = reduction_ratio)),
+                    PreNorm(dim_out, MixFeedForward(dim = dim_out, expansion_factor = ff_expansion)),
+                ]))
+
+            self.stages.append(nn.ModuleList([
+                get_overlap_patches,
+                overlap_patch_embed,
+                layers
+            ]))
+
+    def forward(
+        self,
+        x,
+        return_layer_outputs = False
+    ):
+        h, w = x.shape[-2:]
+        # h = int(h/4)
+        # w = int(w/4)
+        # print(h)
+
+        layer_outputs = []
+        for (get_overlap_patches, overlap_embed, layers) in self.stages:
+            x = get_overlap_patches(x)
+            # print('aaa')
+            # print(x.shape)
+            num_patches = int(x.shape[-1])
+
+
+            # num_patches = int(x.shape[-1]/16)
+            # print(num_patches)
+            ratio = int(sqrt((h * w) / num_patches))
+            # print(ratio)
+            x = rearrange(x, 'b c (h w) -> b c h w', h = h // ratio)
+            # print(x.shape)
+            x = x.type(torch.cuda.FloatTensor)
+
+            x = overlap_embed(x)
+            for (attn, ff) in layers:
+                x = attn(x) + x
+                x = ff(x) + x
+
+            layer_outputs.append(x)
+
+        ret = x if not return_layer_outputs else layer_outputs
+        return ret
+
+class SpikeFormer(nn.Module):
+    def __init__(
+        self,
+        inputDim=64,
+        dims = (32, 64, 160, 256),
+        heads = (1, 2, 5, 8),
+        ff_expansion = (8, 8, 4, 4),
+        reduction_ratio = (8, 4, 2, 1),
+        num_layers = 2,
+        channels =64,
+        decoder_dim = 256,
+        out_channel = 1
+    ):
+        super().__init__()
+        dims, heads, ff_expansion, reduction_ratio, num_layers = map(partial(cast_tuple, depth = 4), (dims, heads, ff_expansion, reduction_ratio, num_layers))
+        assert all([*map(lambda t: len(t) == 4, (dims, heads, ff_expansion, reduction_ratio, num_layers))]), 'only four stages are allowed, all keyword arguments must be either a single value or a tuple of 4 values'
+
+        self.mit = MiT(
+            channels = channels,
+            dims = dims,
+            heads = heads,
+            ff_expansion = ff_expansion,
+            reduction_ratio = reduction_ratio,
+            num_layers = num_layers
+        )
+        self.channel_transform = nn.Sequential(
+            nn.Conv2d(inputDim, 64, 3, 1, 1),
+            nn.GELU()
+        )
+
+        self.to_fused = nn.ModuleList([nn.Sequential(
+            nn.Conv2d(dim, decoder_dim, 1),
+            nn.PixelShuffle(2 ** i),
+            nn.GELU(),
+        ) for i, dim in enumerate(dims)])
+
+        self.to_restore = nn.Sequential(
+            nn.Conv2d(256+64+16+4, decoder_dim, 1),
+            nn.GELU(),
+            nn.Conv2d(decoder_dim, out_channel, 1),
+        )
+        self.fournew = nn.PixelShuffle(4)
+
+
+
+    def forward(self, x):
+        x = self.channel_transform(x)
+        x = self.fournew(x)
+        layer_outputs = self.mit(x, return_layer_outputs = True)
+
+        fused = [to_fused(output) for output, to_fused in zip(layer_outputs, self.to_fused)]
+
+        fused = torch.cat(fused, dim = 1)
+
+        return self.to_restore(fused)

spikezoo/archs/spikeformer/README.md

@@ -0,0 +1,30 @@
+# SpikeFormer [![MIT Licence](https://badges.frapsoft.com/os/mit/mit.svg?v=103)](https://opensource.org/licenses/mit-license.php)
+Pytorch Implementation of "SpikeFormer: Image Reconstruction from the Sequence of Spike Camera Based on Transformer"[[Paper]](https://dl.acm.org/doi/abs/10.1145/3512388.3512399)
+
+## Prerequisites
+* Create a conda environment by running `conda env create -f requirements.yml`
+
+## Dataset Structure
+* To train the SpikeFormer, please organize file structure of the dataset as follows:
+```
+Dataset  
+├── test  
+│   └── c.npz  
+├── train  
+│   └── a.npz  
+└── valid  
+    └── b.npz  
+```
+
+## Pretrained Model
+* Download the pretrained model [here](https://pan.baidu.com/s/1aeW15vQh0GXgRJtfStBHDg) using password: nwh5.
+* Put the model to the path ./CheckPoints/
+
+## Training
+* Run `python train.py` to train SpikeFormer on training set.
+
+## Validation
+* Run `python evaluate.py` to evaluate the performance of trained model on testing set.
+
+## Reconstruct Images from Real Spike Data
+* Run `python recon_real_data.py`.

spikezoo/archs/spikeformer/evaluate.py

@@ -0,0 +1,87 @@
+import os
+os.environ['CUDA_VISIBLE_DEVICES'] = "0"
+import torch
+import numpy as np
+from DataProcess import DataLoader as dl
+from Model import Loss
+from PIL import Image
+from Metrics.Metrics import Metrics
+from Model.SpikeFormer import SpikeFormer
+from utils import LoadModel
+
+if __name__ == "__main__":
+
+    dataPath = "/home/storage2/shechen/Spike_Sample_250x400"
+    spikeRadius = 32  # half length of input spike sequence expcept for the middle frame
+    spikeLen = 2 * spikeRadius + 1  # length of input spike sequence
+    batchSize = 4
+
+    reuse = True
+    checkPath = "CheckPoints/best.pth"
+
+    validContainer = dl.DataContainer(dataPath=dataPath, dataType='valid',
+                                      spikeRadius=spikeRadius,batchSize=batchSize)
+    validData = validContainer.GetLoader()
+
+    metrics = Metrics()
+    # model = Spk2Img(spikeRadius, frameRadius, frameStride).cuda()
+
+    model = SpikeFormer(
+        inputDim=spikeLen,
+        dims=(32, 64, 160, 256),  # dimensions of each stage
+        heads=(1, 2, 5, 8),  # heads of each stage
+        ff_expansion=(8, 8, 4, 4),  # feedforward expansion factor of each stage
+        reduction_ratio=(8, 4, 2, 1),  # reduction ratio of each stage for efficient attention
+        num_layers=2,  # num layers of each stage
+        decoder_dim=256,  # decoder dimension
+        out_channel=1  # channel of restored image
+    ).cuda()
+
+
+    if reuse:
+        _, _, modelDict, _ = LoadModel(checkPath, model)
+
+    model.eval()
+    with torch.no_grad():
+        num = 0
+        pres = []
+        gts = []
+        for i, (spikes, gtImg) in enumerate(validData):
+            B, D, H, W = spikes.size()
+            spikes = spikes.cuda()
+            gtImg = gtImg.cuda()
+            predImg = model(spikes)
+            predImg = predImg.squeeze(1)
+
+            predImg = predImg.clamp(min=-1., max=1.)
+            predImg = predImg.detach().cpu().numpy()
+            gtImg = gtImg.clamp(min=-1., max=1.)
+            gtImg = gtImg.detach().cpu().numpy()
+
+            predImg = (predImg + 1.) / 2. * 255.
+            predImg = predImg.astype(np.uint8)
+            predImg = predImg[:, 3:-3]
+
+            gtImg = (gtImg + 1.) / 2. * 255.
+            gtImg = gtImg.astype(np.uint8)
+
+            pres.append(predImg)
+            gts.append(gtImg)
+        pres = np.concatenate(pres, axis=0)
+        gts = np.concatenate(gts, axis=0)
+
+        psnr = metrics.Cal_PSNR(pres, gts)
+        ssim = metrics.Cal_SSIM(pres, gts)
+        best_psnr, best_ssim, _ = metrics.GetBestMetrics()
+
+        B, H, W = pres.shape
+        divide_line = np.zeros((H, 4)).astype(np.uint8)
+        for pre, gt in zip(pres, gts):
+            num += 1
+            concatImg = np.concatenate([pre, divide_line, gt], axis=1)
+            concatImg = Image.fromarray(concatImg)
+            concatImg.save('EvalResults/test_%s.jpg' % (num))
+
+        print('*********************************************************')
+        print('PSNR: %s, SSIM: %s' % (psnr, ssim))
+

spikezoo/archs/spikeformer/recon_real_data.py

@@ -0,0 +1,97 @@
+import os
+os.environ['CUDA_VISIBLE_DEVICES'] = "0"
+import torch
+import numpy as np
+from PIL import Image
+from Metrics.Metrics import Metrics
+from Model.SpikeFormer import SpikeFormer
+from DataProcess.LoadSpike import load_spike_raw
+from utils import LoadModel
+import shutil
+from PIL import Image
+
+def PredictImg(model, inputs):
+    inputs = torch.FloatTensor(inputs)
+    inputs = inputs.cuda()
+
+    predImg = model(inputs).squeeze(dim=1)
+
+    predImg = predImg.clamp(min=-1., max=1.)
+    predImg = predImg.detach().cpu().numpy()
+    predImg = (predImg + 1.) / 2. * 255.
+    predImg = np.clip(predImg, 0., 255.)
+    predImg = predImg.astype(np.uint8)
+    predImg = predImg[:, 3:-3]
+
+    return predImg
+
+if __name__ == "__main__":
+
+    dataName = "reds"
+    spikeRadius = 32
+    spikeLen = 2 * spikeRadius + 1
+    stride = 32
+    batchSize = 8
+    reuse = True
+    checkPath = "best.pth"
+    sceneClass = {
+                   1:'ballon.dat', 2:'car-100kmh.dat',
+                   3:'forest.dat', 4:'railway.dat',
+                   5:'rotation1.dat', 6:'rotation2.dat',
+                   7:'train-350kmh.dat', 8:'viaduct-bridge.dat'
+    }
+    sceneName = sceneClass[2]
+    dataPath = "/home/storage1/Dataset/SpikeImageData/RealData/%s" %(sceneName)
+    resultPath = sceneName + "_stride_" + str(stride) + "/"
+    shutil.rmtree(resultPath) if os.path.exists(resultPath) else os.mkdir(resultPath)
+    spikes = load_spike_raw(dataPath)
+    totalLen = spikes.shape[0]
+    metrics = Metrics()
+    model = SpikeFormer(
+        inputDim=spikeLen,
+        dims=(32, 64, 160, 256),  # dimensions of each stage
+        heads=(1, 2, 5, 8),  # heads of each stage
+        ff_expansion=(8, 8, 4, 4),  # feedforward expansion factor of each stage
+        reduction_ratio=(8, 4, 2, 1),  # reduction ratio of each stage for efficient attention
+        num_layers=2,  # num layers of each stage
+        decoder_dim=256,  # decoder dimension
+        out_channel = 1  # channel of restored image
+    ).cuda()
+
+    if reuse:
+        _, _, modelDict, _ = LoadModel(checkPath, model)
+
+    model.eval()
+    with torch.no_grad():
+        num = 0
+        pres = []
+        batchFlag = 1
+        inputs = np.zeros((batchSize, spikeLen, 256, 400))  # 65
+        for i in range(32, totalLen - 32, stride):
+            batchFlag = 1
+            spike = spikes[i - spikeRadius: i + spikeRadius + 1]
+            spike = np.pad(spike, ((0, 0), (3, 3), (0, 0)), mode='constant')
+            spike = spike.astype(float) * 2 - 1
+            inputs[num % batchSize] = spike
+            num += 1
+
+            if num % batchSize == 0:
+                predImg = PredictImg(model, inputs)
+                inputs = np.zeros((batchSize, spikeLen, 256, 400))  # 65
+                pres.append(predImg)
+                batchFlag = 0
+
+        if batchFlag == 1:
+            imgNum = num % batchSize
+            inputs = inputs[0: imgNum]
+            predImg = PredictImg(model, inputs)
+            inputs = np.zeros((batchSize, spikeLen, 256, 400))
+            pres.append(predImg)
+
+        predImgs = np.concatenate(pres, axis=0)
+        count = 0
+        for img in predImgs:
+            count += 1
+            img = Image.fromarray(img)
+            img.save(resultPath + '%s.jpg' % (count))
+

spikezoo/archs/spikeformer/requirements.yml

@@ -0,0 +1,95 @@
+name: SpikeFormer
+channels:
+  - pytorch
+  - conda-forge
+  - defaults
+dependencies:
+  - _libgcc_mutex=0.1=main
+  - _openmp_mutex=5.1=1_gnu
+  - blas=1.0=mkl
+  - brotlipy=0.7.0=py38h0a891b7_1004
+  - bzip2=1.0.8=h7f98852_4
+  - ca-certificates=2022.07.19=h06a4308_0
+  - certifi=2022.6.15=py38h06a4308_0
+  - cffi=1.14.6=py38ha65f79e_0
+  - charset-normalizer=2.1.1=pyhd8ed1ab_0
+  - cloudpickle=2.0.0=pyhd3eb1b0_0
+  - cryptography=37.0.2=py38h2b5fc30_0
+  - cudatoolkit=11.6.0=hecad31d_10
+  - cytoolz=0.11.0=py38h7b6447c_0
+  - dask-core=2022.7.0=py38h06a4308_0
+  - ffmpeg=4.3=hf484d3e_0
+  - fftw=3.3.9=h27cfd23_1
+  - freetype=2.10.4=h0708190_1
+  - fsspec=2022.7.1=py38h06a4308_0
+  - gmp=6.2.1=h58526e2_0
+  - gnutls=3.6.13=h85f3911_1
+  - idna=3.3=pyhd8ed1ab_0
+  - imageio=2.9.0=pyhd3eb1b0_0
+  - intel-openmp=2021.4.0=h06a4308_3561
+  - jpeg=9e=h166bdaf_1
+  - lame=3.100=h7f98852_1001
+  - lcms2=2.12=hddcbb42_0
+  - ld_impl_linux-64=2.38=h1181459_1
+  - libffi=3.3=he6710b0_2
+  - libgcc-ng=11.2.0=h1234567_1
+  - libgfortran-ng=11.2.0=h00389a5_1
+  - libgfortran5=11.2.0=h1234567_1
+  - libgomp=11.2.0=h1234567_1
+  - libiconv=1.17=h166bdaf_0
+  - libpng=1.6.37=h21135ba_2
+  - libstdcxx-ng=11.2.0=h1234567_1
+  - libtiff=4.2.0=hf544144_3
+  - libwebp-base=1.2.2=h7f98852_1
+  - locket=1.0.0=py38h06a4308_0
+  - lz4-c=1.9.3=h9c3ff4c_1
+  - mkl=2021.4.0=h06a4308_640
+  - mkl-service=2.4.0=py38h95df7f1_0
+  - mkl_fft=1.3.1=py38h8666266_1
+  - mkl_random=1.2.2=py38h1abd341_0
+  - natsort=7.1.1=pyhd3eb1b0_0
+  - ncurses=6.3=h5eee18b_3
+  - nettle=3.6=he412f7d_0
+  - networkx=2.8.4=py38h06a4308_0
+  - numpy=1.21.5=py38h6c91a56_3
+  - numpy-base=1.21.5=py38ha15fc14_3
+  - olefile=0.46=pyh9f0ad1d_1
+  - openh264=2.1.1=h780b84a_0
+  - openjpeg=2.4.0=hb52868f_1
+  - openssl=1.1.1q=h7f8727e_0
+  - packaging=21.3=pyhd3eb1b0_0
+  - partd=1.2.0=pyhd3eb1b0_1
+  - pillow=8.2.0=py38ha0e1e83_1
+  - pip=22.1.2=py38h06a4308_0
+  - pycparser=2.21=pyhd8ed1ab_0
+  - pyopenssl=22.0.0=pyhd8ed1ab_0
+  - pyparsing=3.0.9=py38h06a4308_0
+  - pysocks=1.7.1=pyha2e5f31_6
+  - python=3.8.13=h12debd9_0
+  - python_abi=3.8=2_cp38
+  - pytorch=1.12.1=py3.8_cuda11.6_cudnn8.3.2_0
+  - pytorch-mutex=1.0=cuda
+  - pywavelets=1.3.0=py38h7f8727e_0
+  - pyyaml=6.0=py38h7f8727e_1
+  - readline=8.1.2=h7f8727e_1
+  - requests=2.28.1=pyhd8ed1ab_1
+  - scikit-image=0.19.2=py38h51133e4_0
+  - scipy=1.7.3=py38h6c91a56_2
+  - setuptools=63.4.1=py38h06a4308_0
+  - six=1.16.0=pyh6c4a22f_0
+  - sqlite=3.39.2=h5082296_0
+  - tifffile=2020.10.1=py38hdd07704_2
+  - tk=8.6.12=h1ccaba5_0
+  - toolz=0.11.2=pyhd3eb1b0_0
+  - torchaudio=0.12.1=py38_cu116
+  - torchvision=0.13.1=py38_cu116
+  - typing_extensions=4.3.0=pyha770c72_0
+  - urllib3=1.26.11=pyhd8ed1ab_0
+  - wheel=0.37.1=pyhd3eb1b0_0
+  - xz=5.2.5=h7f8727e_1
+  - yaml=0.2.5=h7b6447c_0
+  - zlib=1.2.12=h7f8727e_2
+  - zstd=1.5.0=ha95c52a_0
+  - pip:
+    - einops==0.4.1
+    - opencv-python==4.6.0.66