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

spikezoo/utils/optimizer_utils.py

@@ -0,0 +1,22 @@
+# code borrow from https://github.com/nerfstudio-project/nerfstudio/blob/main/nerfstudio/engine/optimizers.py
+from dataclasses import dataclass
+import torch
+from typing import Any, Dict, List, Optional, Type,Tuple
+
+@dataclass
+class OptimizerConfig:
+    def setup(self, model_params) -> torch.optim.Optimizer:
+        """Returns the instantiated object using the config."""
+        kwargs = vars(self).copy()
+        kwargs.pop("_target")
+        return self._target(model_params, **kwargs)
+
+@dataclass
+class AdamOptimizerConfig(OptimizerConfig):
+    """Basic optimizer config with Adam"""
+    lr: float = 1e-4
+    betas: Tuple[float, float] = (0.9, 0.999)
+    eps: float = 1e-8
+    weight_decay: float = 0
+    _target: Type = torch.optim.Adam
+

spikezoo/utils/other_utils.py

@@ -3,7 +3,9 @@ from dataclasses import dataclass, field, asdict
 import requests
 from tqdm import tqdm
 import os
-
+import torch
+import numpy as np
+import random
 
 # log info
 def setup_logging(log_file):
@@ -36,9 +38,9 @@ def save_config(cfg, filename, mode="w"):
 def download_file(url, output_path):
     headers = {
         "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/90.0.4430.212 Safari/537.36",
-        "Accept": "*/*",  # 表示接受所有类型的响应
-        "Accept-Encoding": "gzip, deflate, br",  # 支持压缩格式
-        "Connection": "keep-alive"  # 保持连接
+        "Accept": "*/*",  
+        "Accept-Encoding": "gzip, deflate, br", 
+        "Connection": "keep-alive"  
     }
 
     try:
@@ -62,8 +64,31 @@ def download_file(url, output_path):
             raise RuntimeError(f"Files fail to download 😔😔😔. Try downloading it from {url} and move it to {output_path}.")
     
     except requests.exceptions.RequestException as e:
-        # If an error occurs, remove the partially downloaded file
         if os.path.exists(output_path):
             os.remove(output_path)
             print(f"Partial download failed. The incomplete file has been removed. 😔😔😔")
-        raise RuntimeError(f"Files fail to download 😔😔😔. Try downloading it from {url} and move it to {output_path}.")
+        raise RuntimeError(f"Files fail to download 😔😔😔. Try downloading it from {url} and move it to {output_path}.")
+    
+def check_file_exists(url):
+    response = requests.head(url)
+    if response.status_code == 200:
+        return True
+    else:
+        return False
+
+def getattr_case_insensitive(obj, name):
+    name = name.lower()
+    for attr in dir(obj):
+        if attr.lower() == name:
+            return getattr(obj, attr)
+    raise RuntimeError("No attr found!!")
+
+
+def set_random_seed(seed):
+    """Set random seeds."""
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    np.random.seed(seed)
+    random.seed(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.cuda.manual_seed(seed)

spikezoo/utils/scheduler_utils.py

@@ -0,0 +1,25 @@
+# code borrow from https://github.com/nerfstudio-project/nerfstudio/blob/main/nerfstudio/engine/optimizers.py
+from dataclasses import dataclass
+import torch
+from typing import Any, Dict, List, Optional, Type,Tuple
+
+@dataclass
+class SchedulerConfig:
+    def setup(self, optimizer) -> torch.optim.lr_scheduler.LRScheduler:
+        """Returns the instantiated object using the config."""
+        kwargs = vars(self).copy()
+        kwargs.pop("_target")
+        return self._target(optimizer, **kwargs)
+
+@dataclass
+class CosineAnnealingLRConfig(SchedulerConfig):
+    T_max: int
+    eta_min: float = 0
+    _target: Type = torch.optim.lr_scheduler.CosineAnnealingLR
+
+
+@dataclass
+class MultiStepSchedulerConfig(SchedulerConfig):
+    milestones: List[int]
+    gamma: float = 0.1
+    _target: Type = torch.optim.lr_scheduler.MultiStepLR

spikezoo/utils/spike_utils.py

@@ -2,23 +2,28 @@ import numpy as np
 import torch
 import torch.nn as nn
 import os
-from .vidar_loader import load_vidar_dat_cpp
 from typing import Literal
+import platform
+import cv2
+import imageio
 
-def load_vidar_dat(filename, height, width,remove_head=False, version:Literal['python','cpp'] = "python", out_format : Literal['array','tensor']="array",):
+_platform_check_done = False
+
+
+def load_vidar_dat(filename, height, width, remove_head=False, version: Literal["python", "cpp"] = "cpp", out_format: Literal["array", "tensor"] = "array"):
     """Load the spike stream from the .dat file."""
+    global _platform_check_done
     # Spike decode
-    if version == "python":
-        if isinstance(filename, str):
-            array = np.fromfile(filename, dtype=np.uint8)
-        elif isinstance(filename, (list, tuple)):
-            l = []
-            for name in filename:
-                a = np.fromfile(name, dtype=np.uint8)
-                l.append(a)
-            array = np.concatenate(l)
-        else:
-            raise NotImplementedError
+    if version == "cpp" and platform.system().lower() == "linux":
+        from .vidar_loader import load_vidar_dat_cpp
+
+        spikes = load_vidar_dat_cpp(filename, height, width)
+    else:
+        # todo double check
+        if version == "cpp" and platform.system().lower() != "linux" and _platform_check_done == False:
+            _platform_check_done = True
+            print("Cpp load version is only supported on the linux now. Auto transfer to the python version.")
+        array = np.fromfile(filename, dtype=np.uint8)
         len_per_frame = height * width // 8
         framecnt = len(array) // len_per_frame
         spikes = []
@@ -33,10 +38,6 @@ def load_vidar_dat(filename, height, width,remove_head=False, version:Literal['p
             spk = spk[:, :, :-16] if remove_head == True else spk
             spikes.append(spk)
         spikes = np.concatenate(spikes).astype(np.float32)
-    elif version == "cpp":
-        spikes = load_vidar_dat_cpp(filename, height, width)
-    else:
-        raise RuntimeError("Not recognized version.")
 
     # # Output format conversion
     format_dict = {"array": lambda x: x, "tensor": torch.from_numpy}
@@ -44,15 +45,14 @@ def load_vidar_dat(filename, height, width,remove_head=False, version:Literal['p
     return spikes
 
 
-def SpikeToRaw(save_path, SpikeSeq, filpud=True, delete_if_exists=True):
+def save_vidar_dat(save_path, SpikeSeq, filpud=True):
     """Save the spike sequence to the .dat file."""
-    if delete_if_exists:
-        if os.path.exists(save_path):
-            os.remove(save_path)
+    if os.path.exists(save_path):
+        os.remove(save_path)
     sfn, h, w = SpikeSeq.shape
     remainder = int((h * w) % 8)
     base = np.power(2, np.linspace(0, 7, 8))
-    fid = open(save_path, 'ab')
+    fid = open(save_path, "ab")
     for img_id in range(sfn):
         if filpud:
             spike = np.flipud(SpikeSeq[img_id, :, :])
@@ -61,20 +61,54 @@ def SpikeToRaw(save_path, SpikeSeq, filpud=True, delete_if_exists=True):
         if remainder == 0:
             spike = spike.flatten()
         else:
-            spike = np.concatenate([spike.flatten(), np.array([0]*(8-remainder))])
-        spike = spike.reshape([int(h*w/8), 8])
+            spike = np.concatenate([spike.flatten(), np.array([0] * (8 - remainder))])
+        spike = spike.reshape([int(h * w / 8), 8])
         data = spike * base
         data = np.sum(data, axis=1).astype(np.uint8)
         fid.write(data.tobytes())
     fid.close()
-    return
+
+
+def merge_vidar_dat(filename, dat_files, height, width, remove_head=False):
+    """Merge selected spike dat files."""
+    spikes = []
+    for dat_file in dat_files:
+        spike = load_vidar_dat(dat_file,height, width, remove_head)
+        spikes.append(spike)
+    spikes = np.concatenate(spikes, axis=0)
+    save_vidar_dat(filename, spikes)
+    return spikes
+
+def visual_vidar_dat(filename, spike, out_format: Literal["mp4", "gif"] = "gif", fps=15):
+    """Convert the spike stream to the video."""
+    _, height, width = spike.shape
+    if out_format == "mp4":
+        fourcc = cv2.VideoWriter_fourcc(*"mp4v")  # 或 'avc1'
+        mp4_video = cv2.VideoWriter(filename, fourcc, fps, (width, height))
+    elif out_format == "gif":
+        frames = []
+
+    for idx in range(len(spike)):
+        spk = spike[idx]
+        spk = (255 * spk).astype(np.uint8)
+        spk = spk[..., None].repeat(3, axis=-1)
+        if out_format == "mp4":
+            mp4_video.write(spk)
+        elif out_format == "gif":
+            frames.append(spk)
+
+    if out_format == "mp4":
+        mp4_video.release()
+    elif out_format == "gif":
+        imageio.mimsave(filename, frames, "GIF", fps=fps, loop=0)
+
 
 def video2spike_simulation(imgs, threshold=2.0):
     """Convert the images input to the spike stream."""
     imgs = np.array(imgs)
-    T,H, W = imgs.shape
+    T, H, W = imgs.shape
     spike = np.zeros([T, H, W], np.uint8)
-    integral = np.random.random(size=([H,W])) * threshold
+    integral = np.random.random(size=([H, W])) * threshold
     for t in range(0, T):
         integral += imgs[t]
         fire = (integral - threshold) >= 0
@@ -82,5 +116,3 @@ def video2spike_simulation(imgs, threshold=2.0):
         integral[fire_pos] -= threshold
         spike[t][fire_pos] = 1
     return spike
-
-

spikezoo-0.2.3.dist-info/METADATA

@@ -0,0 +1,263 @@
+Metadata-Version: 2.2
+Name: spikezoo
+Version: 0.2.3
+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
+
+<p align="center">
+    <br>
+    <img src="imgs/spike-zoo.png" width="500"/>
+    <br>
+<p>
+
+<h5 align="center">
+
+[![GitHub repo stars](https://img.shields.io/github/stars/chenkang455/Spike-Zoo?style=flat&logo=github&logoColor=whitesmoke&label=Stars)](https://github.com/chenkang455/Spike-Zoo/stargazers) [![GitHub Issues](https://img.shields.io/github/issues/chenkang455/Spike-Zoo?style=flat&logo=github&logoColor=whitesmoke&label=Stars)](https://github.com/chenkang455/Spike-Zoo/issues) <a href="https://badge.fury.io/py/spikezoo"><img src="https://badge.fury.io/py/spikezoo.svg" alt="PyPI version"></a> [![License](https://img.shields.io/badge/License-MIT-yellow)](https://github.com/chenkang455/Spike-Zoo)
+ 
+<p>
+
+<!-- <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 like training your method from scratch.
+* **24-07-19:** Release the `Spike-Zoo v0.1` code for base evaluation of SOTA methods.
+
+## 🍾 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 `0.2.3` from PyPI:
+```
+pip install spikezoo
+```
+*  Install the latest developing version `0.2.3` 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.
+```
+git clone https://github.com/chenkang455/Spike-Zoo
+cd Spike-Zoo
+python setup.py develop
+```
+
+### 2. Inference 
+Reconstructing images from the spike is super easy with Spike-Zoo. Try the following code of the single model:
+``` python
+from spikezoo.pipeline import Pipeline, PipelineConfig
+import spikezoo as sz
+pipeline = Pipeline(
+    cfg=PipelineConfig(save_folder="results",version="v023"),
+    model_cfg=sz.METHOD.BASE,
+    dataset_cfg=sz.DATASET.BASE 
+)
+```
+You can also run multiple models at once by changing the pipeline (version parameter corresponds to our released different versions in [Releases](https://github.com/chenkang455/Spike-Zoo/releases)):
+``` python
+import spikezoo as sz
+from spikezoo.pipeline import EnsemblePipeline, EnsemblePipelineConfig
+pipeline = EnsemblePipeline(
+    cfg=EnsemblePipelineConfig(save_folder="results",version="v023"),
+    model_cfg_list=[
+        sz.METHOD.BASE,sz.METHOD.TFP,sz.METHOD.TFI,sz.METHOD.SPK2IMGNET,sz.METHOD.WGSE,
+        sz.METHOD.SSML,sz.METHOD.BSF,sz.METHOD.STIR,sz.METHOD.SPIKECLIP,sz.METHOD.SSIR],
+    dataset_cfg=sz.DATASET.BASE,
+)
+```
+Having established our pipelines, we provide following functions to enjoy these spike-to-image models.
+
+* I. Obtain the restoration metric and save the recovered image from the given spike:
+``` python
+# 1. spike-to-image from the given dataset
+pipeline.infer_from_dataset(idx = 0)
+
+# 2. spike-to-image from the given .dat file
+pipeline.infer_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)
+pipeline.infer_from_spk(spike)
+```
+
+
+* II. Save all images from the given dataset.
+``` python
+pipeline.save_imgs_from_dataset()
+```
+
+* III. Calculate the metrics for the specified dataset.
+``` python
+pipeline.cal_metrics()
+```
+
+* IV. Calculate the parameters (params,flops,latency) based on the established pipeline.
+``` python
+pipeline.cal_params()
+```
+
+For detailed usage, welcome check [test_single.ipynb](examples/test/test_single.ipynb) and [test_ensemble.ipynb](examples/test/test_ensemble.ipynb).
+
+### 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_base_dataset import REDS_BASEConfig
+from spikezoo.models.base_model import BaseModelConfig
+pipeline = TrainPipeline(
+    cfg=TrainPipelineConfig(save_folder="results", epochs = 10),
+    dataset_cfg=REDS_BASEConfig(root_dir = "spikezoo/data/REDS_BASE"),
+    model_cfg=BaseModelConfig(),
+)
+pipeline.train()
+``` 
+We finish the training with one 4090 GPU in `2 minutes`, achieving `32.8dB` in PSNR and `0.92` in SSIM.
+
+> 🌟 We encourage users to develop their models with simple modifications to our framework, and the tutorial will be released soon. 
+
+We retrain all supported methods except `SPIKECLIP` on this REDS dataset (training scripts are placed on [examples/train_reds_base](examples/train_reds_base) and evaluation script is placed on [test_REDS_base.py](examples/test/test_REDS_base.py)), with our reported metrics as follows:
+
+| Method               | PSNR  | SSIM   | LPIPS   | NIQE    | BRISQUE  | PIQE  | Params (M) | FLOPs (G) | Latency (ms) |
+|----------------------|:-------:|:--------:|:---------:|:---------:|:----------:|:-------:|:------------:|:-----------:|:--------------:|
+| `TFI`                | 16.503 | 0.454  | 0.382   | 7.289   | 43.17    | 49.12 | 0.00       | 0.00      | 3.60         |
+| `TFP`                | 24.287 | 0.644  | 0.274   | 8.197   | 48.48    | 38.38 | 0.00       | 0.00      | 0.03         |
+| `SPIKECLIP`          | 21.873 | 0.578  | 0.333   | 7.802   | 42.08    | 54.01 | 0.19       | 23.69     | 1.27         |
+| `SSIR`               | 26.544 | 0.718  | 0.325   | 4.769   | 28.45    | 21.59 | 0.38       | 25.92     | 4.52         |
+| `SSML`               | 33.697 | 0.943  | 0.088   | 4.669   | 32.48    | 37.30 | 2.38       | 386.02    | 244.18       |
+| `BASE`               | 36.589 | 0.965  | 0.034   | 4.393   | 26.16    | 38.43 | 0.18       | 18.04     | 0.40         |
+| `STIR`               | 37.914 | 0.973  | 0.027   | 4.236   | 25.10    | 39.18 | 5.08       | 43.31     | 21.07        |
+| `WGSE`               | 39.036 | 0.978  | 0.023   | 4.231   | 25.76    | 44.11 | 3.81       | 415.26    | 73.62        |
+| `SPK2IMGNET`         | 39.154 | 0.978  | 0.022   | 4.243   | 25.20    | 43.09 | 3.90       | 1000.50   | 123.38       |
+| `BSF`                | 39.576 | 0.979  | 0.019   | 4.139   | 24.93    | 43.03 | 2.47       | 705.23    | 401.50       |
+
+### 4. Model Usage
+We also provide a direct interface for users interested in taking the spike-to-image model as a part of their work:
+
+```python 
+import spikezoo as sz
+from spikezoo.models.base_model import BaseModel, BaseModelConfig
+# input data
+spike = sz.load_vidar_dat("data/data.dat", width=400, height=250, out_format="tensor")
+spike = spike[None].cuda()
+print(f"Input spike shape: {spike.shape}")
+# net
+net = BaseModel(BaseModelConfig(model_params={"inDim": 41}))
+net.build_network(mode = "debug")
+# process
+recon_img = net(spike)
+print(recon_img.shape,recon_img.max(),recon_img.min())
+```
+For detailed usage, welcome check [test_model.ipynb](examples/test/test_model.ipynb).
+
+### 5. Spike Utility
+#### I. Faster spike loading interface
+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 [test_load_dat.py](examples/test_load_dat.py) show that the `cpp` version is more than 10 times faster than the `python` version.
+
+#### II. Spike simulation pipeline.
+We provide our overall spike simulation pipeline in [scripts](scripts/), try to modify the config in `run.sh` and run the command to start the simulation process:
+``` bash
+bash run.sh
+```
+
+#### III. Spike-related functions.
+For other spike-related functions, welcome check [spike_utils.py](spikezoo/utils/spike_utils.py)
+
+## 📅 TODO
+- [x] Support the overall pipeline for spike simulation. 
+- [ ] Provide the tutorials.
+- [ ] Support more training settings.
+- [ ] Support more spike-based image reconstruction methods and datasets. 
+
+## 🤗 Supports
+Run the following code to find our supported models, datasets and metrics:
+``` python
+import spikezoo as sz
+print(sz.METHODS)
+print(sz.DATASETS)
+print(sz.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 | 
+| `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 |
+| `base`,`spikeclip`  | Rethinking High-speed Image Reconstruction Framework with Spike Camera |
+
+**Supported Datasets:**
+|  Datasets   | Source  
+|  ----  | ----  | 
+| `reds_base`  | 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), [@shiyan_chen](https://github.com/hnmizuho) 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}"
+}
+```