spikezoo 0.2.2__py3-none-any.whl → 0.2.3.2__py3-none-any.whl

spikezoo/archs/spikeformer/Model/SpikeFormer.py

@@ -1,230 +0,0 @@
-'''
-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

@@ -1,30 +0,0 @@
-# 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

@@ -1,87 +0,0 @@
-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

@@ -1,97 +0,0 @@
-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

@@ -1,95 +0,0 @@
-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