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

spikezoo/archs/spikeformer/train.py

@@ -1,173 +0,0 @@
-import os
-os.environ['CUDA_VISIBLE_DEVICES'] = "0"
-import torch
-from torch import optim
-import numpy as np
-from DataProcess import DataLoader as dl
-from Model.SpikeFormer import SpikeFormer
-from Metrics.Metrics import Metrics
-from Model import Loss
-from utils import SaveModel, LoadModel
-from PIL import Image
-
-def eval(model, validData, epoch, optimizer, metrics):
-
-    model.eval()
-    print('Eval Epoch: %s' %(epoch))
-
-    with torch.no_grad():
-        pres = []
-        gts = []
-        for i, (spikes, gtImg) in enumerate(validData):
-
-            spikes = spikes.cuda()
-            gtImg = gtImg.cuda()
-            predImg = model(spikes)
-            predImg = predImg.squeeze(1)
-            predImg = predImg[:,3:-3,:]
-
-            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)
-
-            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()
-
-        SaveModel(epoch, (psnr, ssim), model, optimizer, saveRoot)
-        if psnr >= best_psnr and ssim >= best_ssim:
-            metrics.Update(psnr, ssim)
-            SaveModel(epoch, (psnr, ssim), model, optimizer, saveRoot, best=True)
-            with open('eval_best_log.txt', 'w') as f:
-                f.write('epoch: %s; psnr: %s, ssim: %s\n' %(epoch, psnr, ssim))
-            B, H, W = pres.shape
-            divide_line = np.zeros((H,4)).astype(np.uint8)
-            num = 0
-            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/valid_%s.jpg' % (num))
-
-        print('*********************************************************')
-        best_psnr, best_ssim, _ = metrics.GetBestMetrics()
-        print('Eval Epoch: %s, PSNR: %s, SSIM: %s, Best_PSNR: %s, Best_SSIM: %s'
-              %(epoch, psnr, ssim, best_psnr, best_ssim))
-
-    model.train()
-
-def Train(trainData, validData, model, optimizer, epoch, start_epoch, metrics, saveRoot, perIter):
-    avg_l2_loss = 0.
-    avg_vgg_loss = 0.
-    avg_edge_loss = 0.
-    avg_total_loss = 0.
-    l2loss = Loss.CharbonnierLoss()
-    vggloss = Loss.VGGLoss4('vgg19-low-level4.pth').cuda()
-    criterion_edge = Loss.EdgeLoss()
-    LAMBDA_L2 = 100.0
-    LAMBDA_VGG = 1.0
-    LAMBDA_EDGE = 5.0   
-    for i in range(start_epoch, epoch):
-        for iter, (spikes, gtImg) in enumerate(trainData):
-            spikes = spikes.cuda()
-            gtImg = gtImg.cuda()
-            predImg = model(spikes)
-            gtImg = gtImg.unsqueeze(1)
-            predImg = predImg[:,:,3:-3,:]
-
-            loss_vgg = vggloss(gtImg, predImg) * LAMBDA_VGG
-            loss_l2 = l2loss(gtImg, predImg) * LAMBDA_L2
-            loss_edge = criterion_edge(gtImg, predImg) * LAMBDA_EDGE
-
-            totalLoss = loss_l2 + loss_vgg + loss_edge
-
-            optimizer.zero_grad()
-            totalLoss.backward()
-            optimizer.step()
-
-            avg_l2_loss += loss_l2.detach().cpu()
-            avg_vgg_loss += loss_vgg.detach().cpu()
-            avg_edge_loss += loss_edge.detach().cpu()
-            avg_total_loss += totalLoss.detach().cpu()
-            if (iter + 1) % perIter == 0:
-                avg_l2_loss = avg_l2_loss / perIter
-                avg_vgg_loss = avg_vgg_loss / perIter
-                avg_edge_loss = avg_edge_loss / perIter
-                avg_total_loss = avg_total_loss / perIter
-                print('=============================================================')
-                print('Epoch: %s, Iter: %s' % (i, iter + 1))
-                print('L2Loss: %s; VggLoss: %s; EdgeLoss: %s; TotalLoss: %s' % (
-                    avg_l2_loss.item(), avg_vgg_loss.item(), avg_edge_loss.item(), avg_total_loss.item()))
-                avg_l2_loss = 0.
-                avg_vgg_loss = 0.
-                avg_edge_loss = 0.
-                avg_total_loss = 0.
-
-        if (i + 1) % 1 == 0:
-            eval(model, validData, i, optimizer, metrics)
-
-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 = 2
-    epoch = 200
-    start_epoch = 0
-    lr = 2e-4
-    saveRoot = "CheckPoints/"  # path to save the trained model
-    perIter = 20
-
-    reuse = False
-    reuseType = 'latest'  # 'latest' or 'best'
-    checkPath = os.path.join('CheckPoints', '%s.pth' % (reuseType))
-
-    trainContainer = dl.DataContainer(dataPath=dataPath, dataType='train',
-                                      spikeRadius=spikeRadius,
-                                      batchSize=batchSize)
-    trainData = trainContainer.GetLoader()
-
-    validContainer = dl.DataContainer(dataPath=dataPath, dataType='valid',
-                                      spikeRadius=spikeRadius,
-                                      batchSize=batchSize)
-    validData = validContainer.GetLoader()
-
-    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()
-
-    optimizer = optim.Adam(model.parameters(), lr=lr, betas=(0.9, 0.999), amsgrad=False)
-
-    if reuse:
-        preEpoch, prePerformance, modelDict, optDict = LoadModel(checkPath, model, optimizer)
-        start_epoch = preEpoch + 1
-        psnr, ssim = prePerformance[0], prePerformance[1]
-        metrics.Update(psnr, ssim)
-        for para in optimizer.param_groups:
-            para['lr'] = lr
-
-    model.train()
-
-    Train(trainData, validData, model, optimizer, epoch, start_epoch,
-                    metrics, saveRoot, perIter)

spikezoo/archs/spikeformer/utils.py

@@ -1,22 +0,0 @@
-import os
-import torch
-
-def SaveModel(epoch, bestPerformance, model, optimizer, saveRoot, best=False):
-    saveDict = {
-        'pre_epoch':epoch,
-        'performance':bestPerformance,
-        'model_state_dict':model.state_dict(),
-        'optimizer_state_dict':optimizer.state_dict()
-    }
-    savePath = os.path.join(saveRoot, '%s.pth' %('latest' if not best else 'best'))
-    torch.save(saveDict, savePath)
-
-def LoadModel(checkPath, model, optimizer=None):
-    stateDict = torch.load(checkPath)
-    pre_epoch = stateDict['pre_epoch']
-    model.load_state_dict(stateDict['model_state_dict'])
-    if optimizer is not None:
-        optimizer.load_state_dict(stateDict['optimizer_state_dict'])
-
-    return pre_epoch, stateDict['performance'], \
-           stateDict['model_state_dict'], stateDict['optimizer_state_dict']

spikezoo/models/spikeformer_model.py

@@ -1,50 +0,0 @@
-import torch
-from dataclasses import dataclass, field
-from spikezoo.models.base_model import BaseModel, BaseModelConfig
-
-
-@dataclass
-class SpikeFormerConfig(BaseModelConfig):
-    # default params for SpikeFormer
-    model_name: str = "spikeformer"
-    model_file_name: str = "Model.SpikeFormer"
-    model_cls_name: str = "SpikeFormer"
-    model_win_length: int = 65
-    require_params: bool = True
-    ckpt_path: str = "weights/spikeformer.pth"
-    model_params: dict = field(
-        default_factory=lambda: {
-            "inputDim": 65,
-            "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,
-            "decoder_dim": 256,
-            "out_channel": 1,
-        }
-    )
-
-
-class SpikeFormer(BaseModel):
-    def __init__(self, cfg: BaseModelConfig):
-        super(SpikeFormer, self).__init__(cfg)
-
-    def preprocess_spike(self, spike):
-        # length
-        spike = self.crop_spike_length(spike)
-        # size
-        if self.spike_size == (250, 400):
-            spike = torch.cat([spike[:, :, :3, :], spike, spike[:, :, -3:, :]], dim=2)
-        elif self.spike_size == (480, 854):
-            spike = torch.cat([spike, spike[:, :, :, -2:]], dim=3)
-        # input
-        spike = 2 * spike - 1
-        return spike
-
-    def postprocess_img(self, image):
-        if self.spike_size == (250, 400):
-            image = image[:, :, 3:-3, :]
-        elif self.spike_size == (480, 854):
-            image = image[:, :, :, :854]
-        return image

spikezoo-0.2.2.dist-info/METADATA

@@ -1,196 +0,0 @@
-Metadata-Version: 2.2
-Name: spikezoo
-Version: 0.2.2
-Summary: A deep learning toolbox for spike-to-image models.
-Home-page: https://github.com/chenkang455/Spike-Zoo
-Author: Kang Chen
-Author-email: mrchenkang@stu.pku.edu.cn
-Requires-Python: >=3.7
-Description-Content-Type: text/markdown
-License-File: LICENSE.txt
-Requires-Dist: torch
-Requires-Dist: requests
-Requires-Dist: numpy
-Requires-Dist: tqdm
-Requires-Dist: scikit-image
-Requires-Dist: lpips
-Requires-Dist: pyiqa
-Requires-Dist: opencv-python
-Requires-Dist: thop
-Requires-Dist: pytorch-wavelets
-Requires-Dist: pytz
-Requires-Dist: PyWavelets
-Requires-Dist: pandas
-Requires-Dist: pillow
-Requires-Dist: scikit-learn
-Requires-Dist: scipy
-Requires-Dist: spikingjelly
-Requires-Dist: setuptools
-Dynamic: author
-Dynamic: author-email
-Dynamic: description
-Dynamic: description-content-type
-Dynamic: home-page
-Dynamic: requires-dist
-Dynamic: requires-python
-Dynamic: summary
-
-<h2 align="center"> 
-  <a href="">⚡Spike-Zoo: A Toolbox for Spike-to-Image Reconstruction
-  </a>
-</h2>
-
-## 📖 About
-⚡Spike-Zoo is the go-to library for state-of-the-art pretrained **spike-to-image** models designed to reconstruct images from spike streams. Whether you're looking for a simple inference solution or aiming to train your own spike-to-image models, ⚡Spike-Zoo is a modular toolbox that supports both, with key features including:
-
-- Fast inference with pre-trained models.
-- Training support for custom-designed spike-to-image models.
-- Specialized functions for processing spike data.
-
-
-
-## 🚩 Updates/Changelog
-* **25-02-02:** Release the `Spike-Zoo v0.2` code, which supports more methods, provide more usages.
-* **24-08-26:** Update the `SpikeFormer` and `RSIR` methods, the `UHSR` dataset and the `piqe` non-reference metric.
-
-* **24-07-19:** Release the `Spike-Zoo v0.1` base code.
-
-## 🍾 Quick Start
-### 1. Installation
-For users focused on **utilizing pretrained models for spike-to-image conversion**, we recommend installing SpikeZoo using one of the following methods:
-
-* Install the last stable version from PyPI:
-```
-pip install spikezoo
-```
-*  Install the latest developing version from the source code:
-```
-git clone https://github.com/chenkang455/Spike-Zoo
-cd Spike-Zoo
-python setup.py install
-```
-
-For users interested in **training their own spike-to-image model based on our framework**, we recommend cloning the repository and modifying the related code directly.
-
-### 2. Inference 
-Reconstructing images from the spike input is super easy with Spike-Zoo. Try the following code of the single model:
-``` python
-from spikezoo.pipeline import Pipeline, PipelineConfig
-pipeline = Pipeline(
-    cfg = PipelineConfig(save_folder="results"),
-    model_cfg="spk2imgnet",
-    dataset_cfg="base"
-)
-```
-You can also run multiple models at once by changing the pipeline:
-``` python
-from spikezoo.pipeline import EnsemblePipeline, EnsemblePipelineConfig
-pipeline = EnsemblePipeline(
-    cfg = EnsemblePipelineConfig(save_folder="results"),
-    model_cfg_list=['tfp','tfi', 'spk2imgnet', 'wgse', 'ssml', 'bsf', 'stir',  'spikeclip','spikeformer'],
-    dataset_cfg="base"
-)
-```
-* Having established the pipeline, run the following code to obtain the metric and save the reconstructed image from the given spike:
-``` python
-# 1. spike-to-image from the given dataset
-pipeline.spk2img_from_dataset(idx = 0)
-
-# 2. spike-to-image from the given .dat file
-pipeline.spk2img_from_file(file_path = 'data/scissor.dat',width = 400,height=250)
-
-# 3. spike-to-image from the given spike
-import spikezoo as sz
-spike = sz.load_vidar_dat("data/scissor.dat",width = 400,height = 250,version='cpp')
-pipeline.spk2img_from_spk(spike)
-```
-For detailed usage, welcome check [test_single.ipynb](examples/test_single.ipynb) and [test_multi.ipynb](examples/test_multi.ipynb) 😊😊😊.
-
-* Save all images of the given dataset.
-``` python
-pipeline.save_imgs_from_dataset()
-```
-
-* Calculate the metrics for the specified dataset.
-``` python
-pipeline.cal_metrics()
-```
-
-* Calculate the parameters (params,flops,latency) based on the established pipeline.
-``` python
-pipeline.cal_params()
-```
-
-### 3. Training
-We provide a user-friendly code for training our provided base model (modified from the `SpikeCLIP`) for the classic `REDS` dataset introduced in `Spk2ImgNet`:
-``` python
-from spikezoo.pipeline import TrainPipelineConfig, TrainPipeline
-from spikezoo.datasets.reds_small_dataset import REDS_Small_Config
-pipeline = TrainPipeline(
-    cfg=TrainPipelineConfig(save_folder="results", epochs = 10),
-    dataset_cfg=REDS_Small_Config(root_dir = "path/REDS_Small"),
-    model_cfg="base",
-)
-pipeline.train()
-``` 
-We finish the training with one 4090 GPU in `2 minutes`, achieving `34.7dB` in PSNR and `0.94` in SSIM.
-
-> 🌟 We encourage users to develop their models using our framework, with the tutorial being released soon.
-
-### 4. Others
-We provide a faster `load_vidar_dat` function implemented with `cpp` (by [@zeal-ye](https://github.com/zeal-ye)):
-``` python
-import spikezoo as sz
-spike = sz.load_vidar_dat("data/scissor.dat",width = 400,height = 250,version='cpp')
-```
-🚀 Results on [examples/test_load_dat.py](examples/test_load_dat.py) show that the `cpp` version is more than 10 times faster than the `python` version.
-
-## 📅 TODO
-- [ ] Provide the tutorials.
-- [ ] Support more training settings.
-- [ ] Support more spike-based image reconstruction methods and datasets. 
-- [ ] Support the overall pipeline for spike simulation. 
-
-## 🤗 Supports
-Run the following code to find our supported models, datasets and metrics:
-``` python
-import spikezoo as sz
-print(sz.get_models())
-print(sz.get_datasets())
-print(sz.get_metrics())
-```
-**Supported Models:**
-|  Models   | Source  
-|  ----  | ----  | 
-| `tfp`,`tfi` | Spike camera and its coding methods | 
-| `spk2imgnet`  | Spk2ImgNet: Learning to Reconstruct Dynamic Scene from Continuous Spike Stream | 
-| `wgse`  | Learning Temporal-Ordered Representation for Spike Streams Based on Discrete Wavelet Transforms | 
-| `ssml`  | Self-Supervised Mutual Learning for Dynamic Scene Reconstruction of Spiking Camera | 
-| `spikeformer`  | SpikeFormer: Image Reconstruction from the Sequence of Spike Camera Based on Transformer |
-| `ssir`  | Spike Camera Image Reconstruction Using Deep Spiking Neural Networks |
-| `bsf`  | Boosting Spike Camera Image Reconstruction from a Perspective of Dealing with Spike Fluctuations |
-| `stir`  | Spatio-Temporal Interactive Learning for Efficient Image Reconstruction of Spiking Cameras |
-| `spikeclip`  | Rethinking High-speed Image Reconstruction Framework with Spike Camera |
-
-**Supported Datasets:**
-|  Datasets   | Source  
-|  ----  | ----  | 
-| `reds_small`  | Spk2ImgNet: Learning to Reconstruct Dynamic Scene from Continuous Spike Stream | 
-| `uhsr`  | Recognizing Ultra-High-Speed Moving Objects with Bio-Inspired Spike Camera | 
-| `realworld`  | `recVidarReal2019`,`momVidarReal2021` in [SpikeCV](https://github.com/Zyj061/SpikeCV) | 
-| `szdata`  | SpikeReveal: Unlocking Temporal Sequences from Real Blurry Inputs with Spike Streams | 
-
-
-## ✨‍ Acknowledgment
-Our code is built on the open-source projects of [SpikeCV](https://spikecv.github.io/), [IQA-Pytorch](https://github.com/chaofengc/IQA-PyTorch), [BasicSR](https://github.com/XPixelGroup/BasicSR) and [NeRFStudio](https://github.com/nerfstudio-project/nerfstudio).We appreciate the effort of the contributors to these repositories. Thanks for [@ruizhao26](https://github.com/ruizhao26) and [@Leozhangjiyuan](https://github.com/Leozhangjiyuan) for their help in building this project.
-
-## 📑 Citation
-If you find our codes helpful to your research, please consider to use the following citation:
-```
-@misc{spikezoo,
-  title={{Spike-Zoo}: Spike-Zoo: A Toolbox for Spike-to-Image Reconstruction},
-  author={Kang Chen and Zhiyuan Ye},
-  year={2025},
-  howpublished = "[Online]. Available: \url{https://github.com/chenkang455/Spike-Zoo}"
-}
-```