spikezoo 0.2.3.5__tar.gz → 0.2.3.6__tar.gz

spikezoo-0.2.3.6/PKG-INFO

@@ -0,0 +1,151 @@
+Metadata-Version: 2.2
+Name: spikezoo
+Version: 0.2.3.6
+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">
+    <img src="imgs/spike-zoo.png" width="300"/>
+<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>  <a href='https://spike-zoo.readthedocs.io/zh-cn/latest/index.html'><img src='https://readthedocs.com/projects/plenoptix-nerfstudio/badge/?version=latest' alt='Documentation Status' /></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>
+</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.
+
+> We are highly looking forward to your advice on our project. We welcome any issues or code contributions and will respond within one day.
+
+## 🚩 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.5` from PyPI:
+```
+pip install spikezoo
+```
+*  Install the latest developing version `0.2.3.6` from the source code **(recommended)**:
+```
+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 
+)
+pipeline.infer_from_dataset(idx = 0)
+```
+
+
+### 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.
+
+## 📚 How to navigate the documentation
+
+| **Link** | **Description** |
+| --- | --- |
+| [Quick Start](https://spike-zoo.readthedocs.io/zh-cn/latest/%E5%BF%AB%E9%80%9F%E5%BC%80%E5%A7%8B.html) | Learn how to quickly get started with the Spike-Zoo repository for inference and training. |
+| [Dataset](https://spike-zoo.readthedocs.io/zh-cn/latest/%E6%95%B0%E6%8D%AE%E9%9B%86.html) | Learn the parameter configuration of datasets and how to construct them. |
+| [Model](https://spike-zoo.readthedocs.io/zh-cn/latest/%E6%A8%A1%E5%9E%8B.html) | Learn the parameter configuration of models and how to construct them. |
+| [Pipeline](https://spike-zoo.readthedocs.io/zh-cn/latest/%E5%A4%84%E7%90%86%E7%AE%A1%E7%BA%BF.html) | Learn how to configure and construct the processing pipeline for models. |
+| [Released Version](https://spike-zoo.readthedocs.io/zh-cn/latest/%E5%8F%91%E8%A1%8C%E7%89%88%E6%9C%AC%E4%BB%8B%E7%BB%8D.html) | Introduces the differences between different release versions of pre-trained weights. |
+| [Examples](https://spike-zoo.readthedocs.io/zh-cn/latest/%E4%BD%BF%E7%94%A8%E4%BE%8B%E5%AD%90.html) | Complete code examples for using Spike-Zoo. |
+| [Supports](https://spike-zoo.readthedocs.io/zh-cn/latest/%E6%94%AF%E6%8C%81%E8%8C%83%E5%9B%B4.html) | Learn about the datasets and models supported by Spike-Zoo. |
+
+
+## 📅 TODO
+- [x] Support the overall pipeline for spike simulation. 
+- [x] Provide the tutorials.
+- [ ] Support more training settings.
+- [ ] Support more spike-based image reconstruction methods and datasets. 
+
+## ✨‍ 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 [@zhiwen_huang](https://github.com/hzw-abc), [@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}: A Toolbox for Spike-to-Image Reconstruction},
+  author={Kang Chen and Zhiyuan Ye and Tiejun Huang and Zhaofei Yu},
+  year={2025},
+  howpublished = {\url{https://github.com/chenkang455/Spike-Zoo}},
+}
+```

spikezoo-0.2.3.6/spikezoo/archs/yourmodel/arch/net.py

@@ -0,0 +1,35 @@
+import torch.nn as nn
+
+def conv_layer(inDim, outDim, ks, s, p, norm_layer="none"):
+    ## convolutional layer
+    conv = nn.Conv2d(inDim, outDim, kernel_size=ks, stride=s, padding=p)
+    relu = nn.ReLU(True)
+    assert norm_layer in ("batch", "instance", "none")
+    if norm_layer == "none":
+        seq = nn.Sequential(*[conv, relu])
+    else:
+        if norm_layer == "instance":
+            norm = nn.InstanceNorm2d(outDim, affine=False, track_running_stats=False)  # instance norm
+        else:
+            momentum = 0.1
+            norm = nn.BatchNorm2d(outDim, momentum=momentum, affine=True, track_running_stats=True)
+        seq = nn.Sequential(*[conv, norm, relu])
+    return seq
+
+
+class YourNet(nn.Module):
+    """Borrow the structure from the SpikeCLIP. (https://arxiv.org/abs/2501.04477)"""
+
+    def __init__(self, inDim=41):
+        super(YourNet, self).__init__()
+        norm = "none"
+        outDim = 1
+        convBlock1 = conv_layer(inDim, 64, 3, 1, 1)
+        convBlock2 = conv_layer(64, 128, 3, 1, 1, norm)
+        convBlock3 = conv_layer(128, 64, 3, 1, 1, norm)
+        convBlock4 = conv_layer(64, 16, 3, 1, 1, norm)
+        conv = nn.Conv2d(16, outDim, 3, 1, 1)
+        self.seq = nn.Sequential(*[convBlock1, convBlock2, convBlock3, convBlock4, conv])
+
+    def forward(self, x):
+        return self.seq(x)

{spikezoo-0.2.3.5 → spikezoo-0.2.3.6}/spikezoo/datasets/__init__.py

@@ -1,4 +1,5 @@
 from spikezoo.datasets.base_dataset import BaseDataset, BaseDatasetConfig
+
 from dataclasses import replace
 import importlib, inspect
 import os
@@ -12,23 +13,24 @@ dataset_list = [file.replace("_dataset.py", "") for file in files_list if file.e
 
 
 # todo register function
-def build_dataset_cfg(cfg: BaseDatasetConfig, split: Literal["train", "test"] = "test"):
+def build_dataset_cfg(cfg: BaseDatasetConfig):
     """Build the dataset from the given dataset config."""
-    # build new cfg according to split
-    cfg = replace(cfg, split=split)
     # 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
-    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)
+    if cfg.dataset_cls_local == None:
+        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
+        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)
+    else:
+        dataset_cls = cfg.dataset_cls_local
     dataset = dataset_cls(cfg)
     return dataset
 
-def build_dataset_name(dataset_name: str, split: Literal["train", "test"] = "test"):
+def build_dataset_name(dataset_name: str):
     """Build the default dataset from the given name."""
     module_name = dataset_name + "_dataset"
     assert dataset_name in dataset_list, f"Given dataset {dataset_name} not in our dataset list {dataset_list}."
@@ -37,22 +39,19 @@ def build_dataset_name(dataset_name: str, split: Literal["train", "test"] = "tes
     # dataset,dataset_config
     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_cfg: BaseDatasetConfig = getattr_case_insensitive(module, dataset_name + "config")()
     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, cfg):
     # 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)
+    if dataset.split == "train" and cfg._mode == "train_mode":
+        return torch.utils.data.DataLoader(dataset, batch_size=cfg.bs_train, shuffle=True, num_workers=cfg.nw_train, 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)
+    else:
+        return torch.utils.data.DataLoader(dataset, batch_size=cfg.bs_test, shuffle=False, num_workers=cfg.nw_test,pin_memory=False)
 
 
 # dataset_size_dict = {}

{spikezoo-0.2.3.5 → spikezoo-0.2.3.6}/spikezoo/datasets/base_dataset.py

@@ -11,6 +11,7 @@ import warnings
 import torch
 from tqdm import tqdm
 from spikezoo.utils.data_utils import Augmentor
+from typing import Optional
 
 
 @dataclass
@@ -36,24 +37,18 @@ class BaseDatasetConfig:
     img_dir_name: str = "gt"
     "Rate. (-1 denotes variant)"
     rate: float = 0.6
-    
+
     # ------------- Config -------------
-    "Dataset split: train/test. Default set as the 'test' for evaluation."
-    split: Literal["train", "test"] = "test"
     "Use the data augumentation technique or not."
     use_aug: bool = False
     "Use cache mechanism."
     use_cache: bool = False
     "Crop size."
     crop_size: tuple = (-1, -1)
-
-
-    # post process
-    def __post_init__(self):
-        self.spike_length = self.spike_length_train if self.split == "train" else self.spike_length_test
-        self.root_dir = Path(self.root_dir) if isinstance(self.root_dir, str) else self.root_dir
-        # todo try download
-        assert self.root_dir.exists(), f"No files found in {self.root_dir} for the specified dataset `{self.dataset_name}`."
+    "Load the dataset from local or spikezoo lib."
+    dataset_cls_local: Optional[Dataset] = None
+    "Spike load version. [python,cpp]"
+    spike_load_version: Literal["python", "cpp"] = "python"
 
 
 # todo cache mechanism
@@ -61,10 +56,6 @@ class BaseDataset(Dataset):
     def __init__(self, cfg: BaseDatasetConfig):
         super(BaseDataset, self).__init__()
         self.cfg = cfg
-        self.augmentor = Augmentor(self.cfg.crop_size) if self.cfg.use_aug == True and self.cfg.split == "train" else -1
-        self.prepare_data()
-        self.cache_data() if cfg.use_cache == True else -1
-        warnings.warn("Lengths of the image list and the spike list should be equal.") if len(self.img_list) != len(self.spike_list) else -1
 
     def __len__(self):
         return len(self.spike_list)
@@ -79,7 +70,7 @@ class BaseDataset(Dataset):
             img = self.get_img(idx)
 
         # process data
-        if self.cfg.use_aug == True and self.cfg.split == "train":
+        if self.cfg.use_aug == True and self.split == "train":
             spike, img = self.augmentor(spike, img)
 
         # rate
@@ -89,15 +80,29 @@ class BaseDataset(Dataset):
         batch = {"spike": spike, "gt_img": img, "rate": rate}
         return batch
 
+    def build_source(self, split: Literal["train", "test"] = "test"):
+        """Build the dataset source and prepare to be loaded files."""
+        # spike length
+        self.split = split
+        self.spike_length = self.cfg.spike_length_train if self.split == "train" else self.cfg.spike_length_test
+        # root dir
+        self.cfg.root_dir = Path(self.cfg.root_dir) if isinstance(self.cfg.root_dir, str) else self.cfg.root_dir
+        assert self.cfg.root_dir.exists(), f"No files found in {self.cfg.root_dir} for the specified dataset `{self.cfg.dataset_name}`."
+        # prepare
+        self.augmentor = Augmentor(self.cfg.crop_size) if self.cfg.use_aug == True and self.split == "train" else -1
+        self.prepare_data()
+        self.cache_data() if self.cfg.use_cache == True else -1
+        warnings.warn("Lengths of the image list and the spike list should be equal.") if len(self.img_list) != len(self.spike_list) else -1
+
     # todo: To be overridden
     def prepare_data(self):
         """Specify the spike and image files to be loaded."""
         # spike
-        self.spike_dir = self.cfg.root_dir / self.cfg.split / self.cfg.spike_dir_name
+        self.spike_dir = self.cfg.root_dir / self.split / self.cfg.spike_dir_name
         self.spike_list = self.get_spike_files(self.spike_dir)
         # gt
         if self.cfg.with_img == True:
-            self.img_dir = self.cfg.root_dir / self.cfg.split / self.cfg.img_dir_name
+            self.img_dir = self.cfg.root_dir / self.split / self.cfg.img_dir_name
             self.img_list = self.get_image_files(self.img_dir)
 
     # todo: To be overridden
@@ -115,12 +120,13 @@ class BaseDataset(Dataset):
             height=self.cfg.height,
             width=self.cfg.width,
             out_format="tensor",
+            version=self.cfg.spike_load_version
         )
         return spike
 
     def get_spike(self, idx):
         """Get and process the spike stream from the given idx."""
-        spike_length = self.cfg.spike_length
+        spike_length = self.spike_length
         spike = self.load_spike(idx)
         assert spike.shape[0] >= spike_length, f"Given spike length {spike.shape[0]} smaller than the required length {spike_length}"
         spike_mid = spike.shape[0] // 2

spikezoo-0.2.3.5/spikezoo/datasets/realworld_dataset.py → spikezoo-0.2.3.6/spikezoo/datasets/realdata_dataset.py

@@ -4,21 +4,19 @@ from dataclasses import dataclass
 
 
 @dataclass
-class RealWorldConfig(BaseDatasetConfig):
-    dataset_name: str = "realworld"
-    root_dir: Path = Path(__file__).parent.parent / Path("data/recVidarReal2019")
+class RealDataConfig(BaseDatasetConfig):
+    dataset_name: str = "realdata"
+    root_dir: Path = Path(__file__).parent.parent / Path("data/realdata")
     width: int = 400
     height: int = 250
     with_img: bool = False
     spike_length_train: int = -1
     spike_length_test: int = -1
     rate: float = 1
-    
 
-
-class RealWorld(BaseDataset):
+class RealData(BaseDataset):
     def __init__(self, cfg: BaseDatasetConfig):
-        super(RealWorld, self).__init__(cfg)
+        super(RealData, self).__init__(cfg)
 
     def prepare_data(self):
         self.spike_dir = self.cfg.root_dir

{spikezoo-0.2.3.5 → spikezoo-0.2.3.6}/spikezoo/datasets/reds_base_dataset.py

@@ -8,7 +8,7 @@ import re
 @dataclass
 class REDS_BASEConfig(BaseDatasetConfig):
     dataset_name: str = "reds_base"
-    root_dir: Path = Path(__file__).parent.parent / Path("data/REDS_BASE")
+    root_dir: Path = Path(__file__).parent.parent / Path("data/reds_base")
     width: int = 400
     height: int = 250
     with_img: bool = True

{spikezoo-0.2.3.5 → spikezoo-0.2.3.6}/spikezoo/datasets/szdata_dataset.py

@@ -9,7 +9,7 @@ import numpy as np
 @dataclass
 class SZDataConfig(BaseDatasetConfig):
     dataset_name: str = "szdata"
-    root_dir: Path = Path(__file__).parent.parent / Path("data/dataset")
+    root_dir: Path = Path(__file__).parent.parent / Path("data/szdata")
     width: int = 400
     height: int = 250
     with_img: bool = True

{spikezoo-0.2.3.5 → spikezoo-0.2.3.6}/spikezoo/datasets/uhsr_dataset.py

@@ -7,7 +7,7 @@ import torch
 @dataclass
 class UHSRConfig(BaseDatasetConfig):
     dataset_name: str = "uhsr"
-    root_dir: Path = Path(__file__).parent.parent / Path("data/U-CALTECH")
+    root_dir: Path = Path(__file__).parent.parent / Path("data/u_caltech")
     width: int = 224
     height: int = 224
     with_img: bool = False

spikezoo-0.2.3.6/spikezoo/datasets/yourdataset_dataset.py

@@ -0,0 +1,23 @@
+from torch.utils.data import Dataset
+from pathlib import Path
+from dataclasses import dataclass
+from typing import Literal, Union
+from typing import Optional
+from spikezoo.datasets.base_dataset import BaseDatasetConfig,BaseDataset
+
+@dataclass
+class YourDatasetConfig(BaseDatasetConfig):
+    dataset_name: str = "yourdataset"
+    root_dir: Union[str, Path] = Path(__file__).parent.parent / Path("data/your_data_path")
+    width: int = 400
+    height: int = 250
+    with_img: bool = True
+    spike_length_train: int = -1
+    spike_length_test: int = -1
+    spike_dir_name: str = "spike_data"
+    img_dir_name: str = "sharp_data"
+    rate: float = 1
+
+class YourDataset(BaseDataset):
+    def __init__(self, cfg: BaseDatasetConfig):
+        super(YourDataset, self).__init__(cfg)

{spikezoo-0.2.3.5 → spikezoo-0.2.3.6}/spikezoo/models/__init__.py

@@ -1,16 +1,6 @@
 import importlib
 import inspect
 from spikezoo.models.base_model import BaseModel,BaseModelConfig
-from spikezoo.models.tfp_model import TFPModel,TFPConfig
-from spikezoo.models.tfi_model import TFIModel,TFIConfig
-from spikezoo.models.spk2imgnet_model import Spk2ImgNet,Spk2ImgNetConfig
-from spikezoo.models.wgse_model import WGSE,WGSEConfig
-from spikezoo.models.ssml_model import SSML,SSMLConfig
-from spikezoo.models.bsf_model import BSF,BSFConfig
-from spikezoo.models.stir_model import STIR,STIRConfig
-from spikezoo.models.ssir_model import SSIR,SSIRConfig
-from spikezoo.models.spikeclip_model import SpikeCLIP,SpikeCLIPConfig
-
 
 from spikezoo.utils.other_utils import getattr_case_insensitive
 import os
@@ -24,14 +14,17 @@ model_list = [file.split("_")[0] for file in files_list if file.endswith("_model
 def build_model_cfg(cfg: BaseModelConfig):
     """Build the model from the given model config."""
     # model module name
-    module_name = cfg.model_name + "_model"
-    assert cfg.model_name in model_list, f"Given model {cfg.model_name} not in our model zoo {model_list}."
-    module_name = "spikezoo.models." + module_name
-    module = importlib.import_module(module_name)
-    # model,model_config
-    model_name = cfg.model_name
-    model_name = model_name + 'Model' if model_name == "base" else model_name
-    model_cls: BaseModel = getattr_case_insensitive(module,model_name)
+    if cfg.model_cls_local == None:
+        module_name = cfg.model_name + "_model"
+        assert cfg.model_name in model_list, f"Given model {cfg.model_name} not in our model zoo {model_list}."
+        module_name = "spikezoo.models." + module_name
+        module = importlib.import_module(module_name)
+        # model,model_config
+        model_name = cfg.model_name
+        model_name = model_name + 'Model' if model_name == "base" else model_name
+        model_cls: BaseModel = getattr_case_insensitive(module,model_name)
+    else:
+        model_cls: BaseModel = cfg.model_cls_local
     model = model_cls(cfg)
     return model
 

{spikezoo-0.2.3.5 → spikezoo-0.2.3.6}/spikezoo/models/base_model.py

@@ -44,7 +44,10 @@ class BaseModelConfig:
     multi_gpu: bool = False
     "Base url."
     base_url: str = "https://github.com/chenkang455/Spike-Zoo/releases/download"
-
+    "Load the model from local class or spikezoo lib. (None)"
+    model_cls_local: Optional[nn.Module] = None
+    "Load the arch from local class or spikezoo lib. (None)"
+    arch_cls_local: Optional[nn.Module] = None
 
 class BaseModel(nn.Module):
     def __init__(self, cfg: BaseModelConfig):
@@ -71,8 +74,11 @@ class BaseModel(nn.Module):
     ):
         """Build the network and load the pretrained weight."""
         # network
-        module = importlib.import_module(f"spikezoo.archs.{self.cfg.model_name}.{self.cfg.model_file_name}")
-        model_cls = getattr(module, self.cfg.model_cls_name)
+        if self.cfg.arch_cls_local == None:
+            module = importlib.import_module(f"spikezoo.archs.{self.cfg.model_name}.{self.cfg.model_file_name}")
+            model_cls = getattr(module, self.cfg.model_cls_name)
+        else:
+            model_cls = self.cfg.arch_cls_local
         # load model config parameters
         if version == "local":
             model = model_cls(**self.cfg.model_params)
@@ -129,7 +135,7 @@ class BaseModel(nn.Module):
         """Crop the spike length."""
         spike_length = spike.shape[1]
         spike_mid = spike_length // 2
-        assert spike_length >= self.model_length, f"Spike input is not long enough, given {spike_length} frames < {self.cfg.model_length}."
+        assert spike_length >= self.model_length, f"Spike input is not long enough, given {spike_length} frames < {self.cfg.model_length} required by the {self.cfg.model_name}."
         # even length
         if self.model_length == self.model_half_length * 2:
             spike = spike[

spikezoo-0.2.3.6/spikezoo/models/yourmodel_model.py

@@ -0,0 +1,22 @@
+from torch.utils.data import Dataset
+from pathlib import Path
+from dataclasses import dataclass
+from typing import Literal, Union
+from typing import Optional
+from spikezoo.models.base_model import BaseModel, BaseModelConfig
+from dataclasses import field
+import torch.nn as nn
+
+
+@dataclass
+class YourModelConfig(BaseModelConfig):
+    model_name: str = "yourmodel"  # 需与文件名保持一致
+    model_file_name: str = "arch.net"  # archs路径下的模块路径
+    model_cls_name: str = "YourNet"  # 模型类名
+    model_length: int = 41
+    require_params: bool = True
+    model_params: dict = field(default_factory=lambda: {"inDim": 41})
+
+class YourModel(BaseModel):
+    def __init__(self, cfg: BaseModelConfig):
+        super(YourModel, self).__init__(cfg)