spikezoo 0.2.2__tar.gz → 0.2.3__tar.gz

spikezoo-0.2.3/PKG-INFO

@@ -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}"
+}
+```

spikezoo-0.2.3/spikezoo/__init__.py

@@ -0,0 +1,29 @@
+from .utils.spike_utils import *
+from .models import model_list
+from .datasets import dataset_list
+from .metrics import metric_all_names
+
+# METHOD NAME DEFINITION
+METHODS = model_list
+class METHOD:
+    BASE = "base"
+    TFP = "tfp"
+    TFI = "tfi"
+    SPK2IMGNET = "spk2imgnet"
+    WGSE = "wgse"
+    SSML = "ssml"
+    BSF = "bsf"
+    STIR = "stir"
+    SSIR = "ssir"
+    SPIKECLIP = "spikeclip"
+
+# DATASET NAME DEFINITION
+DATASETS = dataset_list
+class DATASET:
+    BASE = "base"
+    REDS_BASE = "reds_base"
+    REALWORLD = "realworld"
+    UHSR = "uhsr"
+
+# METRIC NAME DEFINITION
+METRICS = metric_all_names

{spikezoo-0.2.2 → spikezoo-0.2.3}/spikezoo/archs/bsf/models/bsf/bsf.py

@@ -8,18 +8,18 @@ from .align import Multi_Granularity_Align
 class BasicModel(nn.Module):
     def __init__(self):
         super().__init__()
-    
+
     ####################################################################################
     ## Tools functions for neural networks
     def weight_parameters(self):
-        return [param for name, param in self.named_parameters() if 'weight' in name]
+        return [param for name, param in self.named_parameters() if "weight" in name]
 
     def bias_parameters(self):
-        return [param for name, param in self.named_parameters() if 'bias' in name]
+        return [param for name, param in self.named_parameters() if "bias" in name]
 
     def num_parameters(self):
         return sum([p.data.nelement() if p.requires_grad else 0 for p in self.parameters()])
-    
+
     def init_weights(self):
         for layer in self.named_modules():
             if isinstance(layer, nn.Conv2d):
@@ -33,12 +33,21 @@ class BasicModel(nn.Module):
                     nn.init.constant_(layer.bias, 0)
 
 
-def split_and_b_cat(x):
-    x0 = x[:, 10-10:10+10+1].clone()
-    x1 = x[:, 20-10:20+10+1].clone()
-    x2 = x[:, 30-10:30+10+1].clone()
-    x3 = x[:, 40-10:40+10+1].clone()
-    x4 = x[:, 50-10:50+10+1].clone()
+from typing import Literal
+
+
+def split_and_b_cat(x, spike_dim: Literal[41, 61] = 61):
+    if spike_dim == 61:
+        win_r = 10
+        win_step = 10
+    elif spike_dim == 41:
+        win_r = 6
+        win_step = 7
+    x0 = x[:, 0 : 2 * win_r + 1, :, :].clone()
+    x1 = x[:, win_step : win_step + 2 * win_r + 1, :, :].clone()
+    x2 = x[:, 2 * win_step : 2 * win_step + 2 * win_r + 1, :, :].clone()
+    x3 = x[:, 3 * win_step : 3 * win_step + 2 * win_r + 1, :, :].clone()
+    x4 = x[:, 4 * win_step : 4 * win_step + 2 * win_r + 1, :, :].clone()
     return torch.cat([x0, x1, x2, x3, x4], dim=0)
 
 
@@ -61,39 +70,42 @@ class Encoder(nn.Module):
             x = self.act(conv(x) + x)
         return x
 
+
 ##########################################################################
 class BSF(BasicModel):
-    def __init__(self, act=nn.ReLU()):
+    def __init__(self, spike_dim=61, act=nn.ReLU()):
         super().__init__()
+        self.spike_dim = spike_dim
         self.offset_groups = 4
         self.corr_max_disp = 3
-
-        self.rep = MODF(base_dim=64, act=act)
-
+        if spike_dim == 61:
+            self.rep = MODF(in_dim=21,base_dim=64, act=act)
+        elif spike_dim == 41:
+            self.rep = MODF(in_dim=13,base_dim=64, act=act)
         self.encoder = Encoder(base_dim=64, layers=4, act=act)
 
         self.align = Multi_Granularity_Align(base_dim=64, groups=self.offset_groups, act=act, sc=3)
 
         self.recons = nn.Sequential(
-            nn.Conv2d(64*5, 64*3, kernel_size=3, padding=1),
+            nn.Conv2d(64 * 5, 64 * 3, kernel_size=3, padding=1),
             act,
-            nn.Conv2d(64*3, 64, kernel_size=3, padding=1),
+            nn.Conv2d(64 * 3, 64, kernel_size=3, padding=1),
             act,
             nn.Conv2d(64, 1, kernel_size=3, padding=1),
         )
 
     def forward(self, input_dict):
-        dsft_dict = input_dict['dsft_dict']
-        dsft11 = dsft_dict['dsft11']
-        dsft12 = dsft_dict['dsft12']
-        dsft21 = dsft_dict['dsft21']
-        dsft22 = dsft_dict['dsft22']
+        dsft_dict = input_dict["dsft_dict"]
+        dsft11 = dsft_dict["dsft11"]
+        dsft12 = dsft_dict["dsft12"]
+        dsft21 = dsft_dict["dsft21"]
+        dsft22 = dsft_dict["dsft22"]
 
         dsft_b_cat = {
-            'dsft11': split_and_b_cat(dsft11),
-            'dsft12': split_and_b_cat(dsft12),
-            'dsft21': split_and_b_cat(dsft21),
-            'dsft22': split_and_b_cat(dsft22),
+            "dsft11": split_and_b_cat(dsft11, self.spike_dim),
+            "dsft12": split_and_b_cat(dsft12, self.spike_dim),
+            "dsft21": split_and_b_cat(dsft21, self.spike_dim),
+            "dsft22": split_and_b_cat(dsft22, self.spike_dim),
         }
 
         feat_b_cat = self.rep(dsft_b_cat)

{spikezoo-0.2.2 → spikezoo-0.2.3}/spikezoo/archs/bsf/models/bsf/rep.py

@@ -2,11 +2,11 @@ import torch
 import torch.nn as nn
 
 class MODF(nn.Module):
-    def __init__(self, base_dim=64, act=nn.ReLU()):
+    def __init__(self, in_dim = 21, base_dim=64, act=nn.ReLU()):
         super().__init__()
         self.base_dim = base_dim
 
-        self.conv1 = self._make_layer(input_dim=21, hidden_dim=self.base_dim, output_dim=self.base_dim, act=act)
+        self.conv1 = self._make_layer(input_dim=in_dim, hidden_dim=self.base_dim, output_dim=self.base_dim, act=act)
         self.conv_for_others = nn.ModuleList([
             self._make_layer(input_dim=self.base_dim, hidden_dim=self.base_dim, output_dim=self.base_dim, act=act) for ii in range(3)
         ])

{spikezoo-0.2.2 → spikezoo-0.2.3}/spikezoo/archs/spk2imgnet/nets.py

@@ -167,7 +167,7 @@ class FusionMaskV1(nn.Module):
 
 # current best model
 class SpikeNet(nn.Module):
-    def __init__(self, in_channels, features, out_channels, win_r, win_step):
+    def __init__(self, in_channels = 13, features = 64, out_channels = 1, win_r = 6, win_step = 7):
         super(SpikeNet, self).__init__()
         self.extractor = FeatureExtractor(
             in_channels=in_channels,

{spikezoo-0.2.2 → spikezoo-0.2.3}/spikezoo/archs/ssir/models/networks.py

@@ -56,6 +56,6 @@ class SSIR(BasicModel):
         out3 = self.pred3(x7)
 
         if self.training:
-            return [out3]
+            return out3
         else:
             return out3

{spikezoo-0.2.2 → spikezoo-0.2.3}/spikezoo/archs/ssml/model.py

@@ -272,18 +272,22 @@ class BSN(nn.Module):
             diff = W - H
             x0 = x0[:, :, (diff // 2):(diff // 2 + H), 0:W]
             
-        return x0,tfi_label,tfp_label   
+        return x0
 
 class DoubleNet(nn.Module):
     def __init__(self):
         super().__init__()
         self.nbsn = BSN(n_channels=41, n_output=1,blind=False)
-        # self.bsn = BSN(n_channels=41, n_output=1,blind=True)
+        self.bsn = BSN(n_channels=41, n_output=1,blind=True)
     
     def forward(self, x):  
-        out1,_,_ = self.nbsn(x)
-        
-        return out1
+        if self.training:
+            bsn_pred = self.bsn(x)
+            nbsn_pred = self.nbsn(x)
+            return bsn_pred,nbsn_pred
+        else:
+            nbsn_pred = self.nbsn(x)
+            return nbsn_pred
 
 if __name__ == '__main__':
     a=DoubleNet().cuda()

{spikezoo-0.2.2 → spikezoo-0.2.3}/spikezoo/archs/stir/metrics/losses.py

@@ -6,7 +6,7 @@ import torch.nn.functional as F
 
 import math
 
-from package_core.losses import *
+from ..package_core.package_core.losses import *
 
 def compute_l1_loss(img_list, gt):
     l1_loss = 0.0

{spikezoo-0.2.2 → spikezoo-0.2.3}/spikezoo/archs/stir/models/networks_STIR.py

@@ -292,16 +292,21 @@ class STIRDecorder(nn.Module):#second and third levels
 
 ##############################Our Model####################################
 class STIR(BasicModel):
-    def __init__(self, hidd_chs=8, win_r=6, win_step=7):
+    def __init__(self, spike_dim = 61,hidd_chs=8, win_r=6, win_step=7):
         super().__init__()
 
         self.init_chs = [16, 24, 32, 64, 96]
         self.hidd_chs = hidd_chs
+        self.spike_dim = spike_dim
         self.attn_num_splits = 1
 
         self.N_group = 3  
-        
-        dim_tfp = 16
+        if spike_dim == 61:
+            self.resnet =  ResidualBlock(in_channles=21, num_channles=11, use_1x1conv=True)
+            dim_tfp = 16 # 5 + num_channels
+        elif spike_dim == 41:
+            self.resnet =  ResidualBlock(in_channles=15, num_channles=11, use_1x1conv=True)
+            dim_tfp = 15  # 4 + num_channels
         self.encoder = ImageEncoder(in_chs=dim_tfp, init_chs=self.init_chs)
         
         self.transformer = CrossTransformerBlock(dim=self.init_chs[-1], num_heads=4, ffn_expansion_factor=2.66, bias=False, LayerNorm_type='WithBias')
@@ -314,14 +319,16 @@ class STIR(BasicModel):
         self.win_r = win_r
         self.win_step = win_step
 
-        self.resnet =  ResidualBlock(in_channles=21, num_channles=11, use_1x1conv=True)
-
     def forward(self, x):
         b,_,h,w=x.size()
-        
-        block1 = x[:, 0  : 21, :, :]
-        block2 = x[:, 20 : 41, :, :]
-        block3 = x[:, 40 : 61, :, :]
+        if self.spike_dim == 61:
+            block1 = x[:, 0  : 21, :, :]
+            block2 = x[:, 20 : 41, :, :]
+            block3 = x[:, 40 : 61, :, :]
+        elif self.spike_dim == 41:
+            block1 = x[:, 0  : 15, :, :]
+            block2 = x[:, 13 : 28, :, :]
+            block3 = x[:, 26 : 41, :, :]
 
         repre1 = TFP(block1, channel_step=2)#C: 5
         repre2 = TFP(block2, channel_step=2)

{spikezoo-0.2.2 → spikezoo-0.2.3}/spikezoo/archs/tfi/nets.py

@@ -6,7 +6,7 @@ import torch
 
 
 class TFIModel(nn.Module):
-    def __init__(self, model_win_length):
+    def __init__(self, model_win_length = 41):
         super(TFIModel, self).__init__()
         self.window = model_win_length
         self.hald_window = model_win_length // 2

{spikezoo-0.2.2 → spikezoo-0.2.3}/spikezoo/archs/tfp/nets.py

@@ -3,7 +3,7 @@ import torch
 
 
 class TFPModel(nn.Module):
-    def __init__(self, model_win_length):
+    def __init__(self, model_win_length = 41):
         self.window = model_win_length
         super(TFPModel, self).__init__()
 

{spikezoo-0.2.2 → spikezoo-0.2.3}/spikezoo/archs/wgse/dwtnets.py

@@ -94,15 +94,15 @@ class Dwt1dModule_Tcn(nn.Module):
 class Dwt1dResnetX_TCN(nn.Module):
     def __init__(
         self, 
-        wvlname='db1', 
-        J=3, 
-        yl_size=14, 
-        yh_size=[26, 18, 14], 
-        num_residual_blocks=2, 
+        wvlname='db8', 
+        J=5, 
+        yl_size=15, 
+        yh_size=[28, 21, 18, 16, 15], 
+        num_residual_blocks=3, 
         norm=None, 
         inc=41, 
         ks=3,
-        store_features=False
+        store_features=True
     ):
         super().__init__()
 

{spikezoo-0.2.2 → spikezoo-0.2.3}/spikezoo/datasets/__init__.py

@@ -4,28 +4,30 @@ import importlib, inspect
 import os
 import torch
 from typing import Literal
+from spikezoo.utils.other_utils import getattr_case_insensitive
 
 # todo auto detect/register datasets
 files_list = os.listdir(os.path.dirname(os.path.abspath(__file__)))
 dataset_list = [file.replace("_dataset.py", "") for file in files_list if file.endswith("_dataset.py")]
 
+
 # todo register function
 def build_dataset_cfg(cfg: BaseDatasetConfig, split: Literal["train", "test"] = "test"):
     """Build the dataset from the given dataset config."""
     # build new cfg according to split
-    cfg = replace(cfg,split = split,spike_length = cfg.spike_length_train if split == "train" else cfg.spike_length_test)
+    cfg = replace(cfg, split=split, spike_length=cfg.spike_length_train if split == "train" else cfg.spike_length_test)
     # dataset module
     module_name = cfg.dataset_name + "_dataset"
     assert cfg.dataset_name in dataset_list, f"Given dataset {cfg.dataset_name} not in our dataset list {dataset_list}."
     module_name = "spikezoo.datasets." + module_name
     module = importlib.import_module(module_name)
     # dataset,dataset_config
-    classes = sorted([name for name, obj in inspect.getmembers(module) if inspect.isclass(obj) and obj.__module__ == module.__name__])
-    dataset_cls: BaseDataset = getattr(module, classes[0])
+    dataset_name = cfg.dataset_name
+    dataset_name = dataset_name + "Dataset" if dataset_name == "base" else dataset_name
+    dataset_cls: BaseDataset = getattr_case_insensitive(module, dataset_name)
     dataset = dataset_cls(cfg)
     return dataset
 
-
 def build_dataset_name(dataset_name: str, split: Literal["train", "test"] = "test"):
     """Build the default dataset from the given name."""
     module_name = dataset_name + "_dataset"
@@ -33,21 +35,21 @@ def build_dataset_name(dataset_name: str, split: Literal["train", "test"] = "tes
     module_name = "spikezoo.datasets." + module_name
     module = importlib.import_module(module_name)
     # dataset,dataset_config
-    classes = sorted([name for name, obj in inspect.getmembers(module) if inspect.isclass(obj) and obj.__module__ == module.__name__])
-    dataset_cls: BaseDataset = getattr(module, classes[0])
-    dataset_cfg: BaseDatasetConfig = getattr(module, classes[1])(split=split)
+    dataset_name = dataset_name + "Dataset" if dataset_name == "base" else dataset_name
+    dataset_cls: BaseDataset = getattr_case_insensitive(module, dataset_name)
+    dataset_cfg: BaseDatasetConfig = getattr_case_insensitive(module, dataset_name + "config")(split=split)
     dataset = dataset_cls(dataset_cfg)
     return dataset
 
 
 # todo to modify according to the basicsr
-def build_dataloader(dataset: BaseDataset,cfg = None):
+def build_dataloader(dataset: BaseDataset, cfg=None):
     # train dataloader
     if dataset.cfg.split == "train":
         if cfg is None:
             return torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False, num_workers=1)
         else:
-            return torch.utils.data.DataLoader(dataset, batch_size=cfg.bs_train, shuffle=True, num_workers=cfg.num_workers,pin_memory=cfg.pin_memory)
+            return torch.utils.data.DataLoader(dataset, batch_size=cfg.bs_train, shuffle=True, num_workers=cfg.num_workers, pin_memory=cfg.pin_memory)
     # test dataloader
     elif dataset.cfg.split == "test":
         return torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False, num_workers=1)