spikezoo 0.1__tar.gz

spikezoo-0.1/LICENSE.txt

@@ -0,0 +1,17 @@
+MIT License
+Copyright (c) 2018 YOUR NAME
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

spikezoo-0.1/PKG-INFO

@@ -0,0 +1,39 @@
+Metadata-Version: 2.2
+Name: spikezoo
+Version: 0.1
+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
+
+⚡Spike-Zoo is the go-to library for state-of-the-art pretrained **spike-to-image** models for reconstructing the image from the given spike stream. Whether you're looking for a **simple inference** solution or **training** your own spike-to-image models, ⚡Spike-Zoo is a modular toolbox that supports both. 
+

spikezoo-0.1/README.md

@@ -0,0 +1,2 @@
+⚡Spike-Zoo is the go-to library for state-of-the-art pretrained **spike-to-image** models for reconstructing the image from the given spike stream. Whether you're looking for a **simple inference** solution or **training** your own spike-to-image models, ⚡Spike-Zoo is a modular toolbox that supports both. 
+

spikezoo-0.1/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

spikezoo-0.1/setup.py

@@ -0,0 +1,23 @@
+from setuptools import find_packages
+from setuptools import setup
+
+with open("./requirements.txt", "r", encoding="utf-8") as fh:
+    install_requires = fh.read()
+
+with open("./README.md", "r", encoding="utf-8") as fh:
+    long_description = fh.read()
+
+setup(
+    install_requires=install_requires,
+    name="spikezoo",
+    version="0.1",
+    author="Kang Chen",
+    author_email="mrchenkang@stu.pku.edu.cn",
+    description="A deep learning toolbox for spike-to-image models.",
+    long_description=long_description,
+    long_description_content_type="text/markdown",
+    url="https://github.com/chenkang455/Spike-Zoo",
+    packages=find_packages(),
+    python_requires='>=3.7',
+    include_package_data=False
+)

spikezoo-0.1/spikezoo/datasets/__init__.py

@@ -0,0 +1,68 @@
+from spikezoo.datasets.base_dataset import BaseDataset, BaseDatasetConfig
+from dataclasses import replace
+import importlib, inspect
+import os
+import torch
+from typing import Literal
+
+# 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)
+    # 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 = 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"
+    assert dataset_name in dataset_list, f"Given dataset {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_cfg: BaseDatasetConfig = getattr(module, classes[1])(split=split)
+    dataset = dataset_cls(dataset_cfg)
+    return dataset
+
+
+# todo to modify according to the basicsr
+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)
+    # test dataloader
+    elif dataset.cfg.split == "test":
+        return torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False, num_workers=1)
+
+
+# dataset_size_dict = {}
+# for dataset in dataset_list:
+#     module_name = dataset + "_dataset"
+#     module_name = "spikezoo.datasets." + module_name
+#     module = importlib.import_module(module_name)
+#     classes = sorted([name for name, obj in inspect.getmembers(module) if inspect.isclass(obj)])
+#     dataset_cfg: BaseDatasetConfig = getattr(module, classes[1])
+#     dataset_size_dict[dataset] = (dataset_cfg.height, dataset_cfg.width)
+
+
+# def get_dataset_size(name):
+#     assert name in dataset_list, f"Given dataset {name} not in our dataset list {dataset_list}."
+#     return dataset_size_dict[name]

spikezoo-0.1/spikezoo/datasets/base_dataset.py

@@ -0,0 +1,157 @@
+import torch
+from torch.utils.data import Dataset
+from pathlib import Path
+import cv2
+import numpy as np
+from spikezoo.utils.spike_utils import load_vidar_dat
+import re
+from dataclasses import dataclass, replace
+from typing import Literal
+import warnings
+import torch
+from tqdm import tqdm
+from spikezoo.utils.data_utils import Augmentor
+
+
+@dataclass
+class BaseDatasetConfig:
+    # ------------- Not Recommended to Change -------------
+    "Dataset name."
+    dataset_name: str = "base"
+    "Directory specifying location of data."
+    root_dir: Path = Path(__file__).parent.parent / Path("data/base")
+    "Image width."
+    width: int = 400
+    "Image height."
+    height: int = 250
+    "Spike paried with the image or not."
+    with_img: bool = True
+    "Dataset spike length for the train data."
+    spike_length_train: int = -1
+    "Dataset spike length for the test data."
+    spike_length_test: int = -1
+    "Dataset spike length for the instantiation dataclass."
+    spike_length: int = -1
+    "Dir name for the spike."
+    spike_dir_name: str = "spike"
+    "Dir name for the image."
+    img_dir_name: str = "gt"
+
+    # ------------- 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}`."
+
+# todo cache mechanism
+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)
+
+    def __getitem__(self, idx: int):
+        # load data
+        if self.cfg.use_cache == True:
+            spike, img = self.cache_spkimg[idx]
+            spike = spike.to(torch.float32)
+        else:
+            spike = self.get_spike(idx)
+            img = self.get_img(idx)
+
+        # process data
+        if self.cfg.use_aug == True and self.cfg.split == "train":
+            spike, img = self.augmentor(spike, img)
+
+        batch = {"spike": spike, "img": img}
+        return batch
+
+    # 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_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_list = self.get_image_files(self.img_dir)
+
+    # todo: To be overridden
+    def get_spike_files(self, path: Path):
+        """Recognize spike files automatically (default .dat)."""
+        files = path.glob("**/*.dat")
+        return sorted(files)
+
+    # todo: To be overridden
+    def load_spike(self, idx):
+        """Load the spike stream from the given idx."""
+        spike_name = str(self.spike_list[idx])
+        spike = load_vidar_dat(
+            spike_name,
+            height=self.cfg.height,
+            width=self.cfg.width,
+            out_type="float",
+            out_format="tensor",
+        )
+        return spike
+
+    def get_spike(self, idx):
+        """Get and process the spike stream from the given idx."""
+        spike_length = self.cfg.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
+        # spike length process
+        if spike_length == -1:
+            spike = spike
+        elif spike_length % 2 == 1:
+            spike = spike[spike_mid - spike_length // 2 : spike_mid + spike_length // 2 + 1]
+        elif spike_length % 2 == 0:
+            spike = spike[spike_mid - spike_length // 2 : spike_mid + spike_length // 2]
+        return spike
+
+    def get_image_files(self, path: Path):
+        """Recognize image files automatically."""
+        files = [f for f in path.glob("**/*") if re.match(r".*\.(jpg|jpeg|png)$", f.name, re.IGNORECASE)]
+        return sorted(files)
+
+    # todo: To be overridden
+    def get_img(self, idx):
+        """Get the image from the given idx."""
+        if self.cfg.with_img:
+            img_name = str(self.img_list[idx])
+            img = cv2.imread(img_name)
+            img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+            img = (img / 255).astype(np.float32)
+            img = img[None]
+            img = torch.from_numpy(img)
+        else:
+            spike = self.get_spike(idx)
+            img = torch.mean(spike, dim=0, keepdim=True)
+        return img
+
+    def cache_data(self):
+        """Cache the data."""
+        self.cache_spkimg = []
+        for idx in tqdm(range(len(self.spike_list)), desc="Caching data", unit="sample"):
+            spike = self.get_spike(idx).to(torch.uint8)
+            img = self.get_img(idx)
+            self.cache_spkimg.append([spike, img])

spikezoo-0.1/spikezoo/datasets/realworld_dataset.py

@@ -0,0 +1,25 @@
+from pathlib import Path
+from spikezoo.datasets.base_dataset import BaseDataset, BaseDatasetConfig
+from dataclasses import dataclass
+
+
+@dataclass
+class RealWorld_Config(BaseDatasetConfig):
+    dataset_name: str = "realworld"
+    root_dir: Path = Path(__file__).parent.parent / Path("data/recVidarReal2019")
+    width: int = 400
+    height: int = 250
+    with_img: bool = False
+    spike_length_train: int = -1
+    spike_length_test: int = -1
+
+
+class RealWorld(BaseDataset):
+    def __init__(self, cfg: BaseDatasetConfig):
+        super(RealWorld, self).__init__(cfg)
+
+    def prepare_data(self):
+        self.spike_dir = self.cfg.root_dir
+        self.spike_list = self.get_spike_files(self.spike_dir)
+        
+    

spikezoo-0.1/spikezoo/datasets/reds_small_dataset.py

@@ -0,0 +1,27 @@
+from torch.utils.data import Dataset
+from pathlib import Path
+from spikezoo.datasets.base_dataset import BaseDataset, BaseDatasetConfig
+from dataclasses import dataclass
+import re
+
+@dataclass
+class REDS_Small_Config(BaseDatasetConfig):
+    dataset_name: str = "reds_small"
+    root_dir: Path = Path(__file__).parent.parent / Path("data/REDS_Small")
+    width: int = 400
+    height: int = 250
+    with_img: bool = True
+    spike_length_train: int = 41
+    spike_length_test: int = 301
+    spike_dir_name: str = "spike"
+    img_dir_name: str = "gt"
+
+class REDS_Small(BaseDataset):
+    def __init__(self, cfg: BaseDatasetConfig):
+        super(REDS_Small, self).__init__(cfg)
+
+    def prepare_data(self):
+        super().prepare_data()
+        if self.cfg.split == "train":
+            self.img_list = [self.img_dir / Path(str(s.name).replace('.dat','.png')) for s in self.spike_list]
+

spikezoo-0.1/spikezoo/datasets/szdata_dataset.py

@@ -0,0 +1,37 @@
+from pathlib import Path
+from spikezoo.datasets.base_dataset import BaseDataset, BaseDatasetConfig
+from dataclasses import dataclass
+import cv2
+import torch
+import numpy as np
+
+@dataclass
+class SZData_Config(BaseDatasetConfig):
+    dataset_name: str = "szdata"
+    root_dir: Path = Path(__file__).parent.parent / Path("data/dataset")
+    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"
+
+class SZData(BaseDataset):
+    def __init__(self, cfg: BaseDatasetConfig):
+        super(SZData, self).__init__(cfg)
+
+    def get_img(self, idx):
+        if self.cfg.with_img:
+            spike_name = self.spike_list[idx]
+            img_name = str(spike_name).replace(self.cfg.spike_dir_name,self.cfg.img_dir_name).replace(".dat",".png")
+            img = cv2.imread(img_name)
+            img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+            img = (img / 255).astype(np.float32)
+            img = img[None]
+            img = torch.from_numpy(img)
+        else:
+            spike = self.get_spike(idx)
+            img = torch.mean(spike, dim=0, keepdim=True)
+        return img
+    

spikezoo-0.1/spikezoo/datasets/uhsr_dataset.py

@@ -0,0 +1,38 @@
+from pathlib import Path
+from spikezoo.datasets.base_dataset import BaseDataset, BaseDatasetConfig
+from dataclasses import dataclass
+import numpy as np
+import torch
+
+@dataclass
+class UHSR_Config(BaseDatasetConfig):
+    dataset_name: str = "uhsr"
+    root_dir: Path = Path(__file__).parent.parent / Path("data/U-CALTECH")
+    width: int = 224
+    height: int = 224
+    with_img: bool = False
+    spike_length_train: int = 200
+    spike_length_test: int = 200
+    spike_dir_name: str = "spike"
+    img_dir_name: str = ""
+
+
+class UHSR(BaseDataset):
+    def __init__(self, cfg: BaseDatasetConfig):
+        super(UHSR, self).__init__(cfg)
+
+    def prepare_data(self):
+        self.spike_dir = self.cfg.root_dir / self.cfg.dataset_split
+        self.spike_list = self.get_spike_files(self.spike_dir)
+
+    def get_spike_files(self, path: Path):
+        files = path.glob("**/*.npz")
+        return sorted(files)
+
+    def load_spike(self,idx):
+        spike_name = str(self.spike_list[idx])
+        data = np.load(spike_name)
+        spike = data["spk"].astype(np.float32)
+        spike = torch.from_numpy(spike)
+        spike = spike[:, 13:237, 13:237]
+        return spike

spikezoo-0.1/spikezoo/metrics/__init__.py

@@ -0,0 +1,96 @@
+from skimage import metrics
+import torch
+import torch.hub
+from lpips.lpips import LPIPS
+import os
+import os
+import pyiqa
+import numpy as np
+from torchvision import transforms
+import torch.nn.functional as F
+
+# todo with the union type
+metric_pair_names = ["psnr", "ssim", "lpips", "mse"]
+metric_single_names = ["niqe", "brisque", "piqe"]
+metric_all_names = metric_pair_names + metric_single_names
+
+metric_single_list = {}
+
+metric_pair_list = {
+    "mse": metrics.mean_squared_error,
+    "ssim": metrics.structural_similarity,
+    "psnr": metrics.peak_signal_noise_ratio,
+    "lpips": None,
+}
+
+
+def cal_metric_single(img: torch.Tensor, metric_name="niqe"):
+    if metric_name not in metric_single_list.keys():
+        if metric_name in pyiqa.list_models():
+            iqa_metric = pyiqa.create_metric(metric_name, device=torch.device("cuda"))
+            metric_single_list.update({metric_name: iqa_metric})
+        else:
+            raise RuntimeError(f"Metric {metric_name} not recognized by the IQA lib.")
+    # image process
+    if img.dim() == 3:
+        img = img[None]
+    elif img.dim() == 2:
+        img = img[None, None]
+
+    # resize
+    if metric_name == "liqe_mix":
+        short_edge = 384
+        h, w = img.shape[2], img.shape[3]
+        if h < w:
+            new_h, new_w = short_edge, int(w * short_edge / h)
+        else:
+            new_h, new_w = int(h * short_edge / w), short_edge
+        img = F.interpolate(img, size=(new_h, new_w), mode="bilinear", align_corners=False)
+    return metric_single_list[metric_name](img).item()
+
+
+def cal_metric_pair(im1t: torch.Tensor, im2t: torch.Tensor, metric_name="mse"):
+    """
+    im1t, im2t: torch.tensors with batched imaged shape, range from (0, 1)
+    """
+    if metric_name not in metric_pair_list.keys():
+        raise RuntimeError(f"Metric {metric_name} not recognized")
+    if metric_name == "lpips" and metric_pair_list[metric_name] is None:
+        metric_pair_list[metric_name] = LPIPS().cuda() if im1t.is_cuda else LPIPS().cpu()
+    metric_method = metric_pair_list[metric_name]
+
+    # convert from [0, 1] to [-1, 1]
+    im1t = (im1t * 2 - 1).clamp(-1, 1)
+    im2t = (im2t * 2 - 1).clamp(-1, 1)
+
+    # [c,h,w] -> [1,c,h,w]
+    if im1t.dim() == 3:
+        im1t = im1t[None]
+        im2t = im2t[None]
+    elif im1t.dim() == 2:
+        im1t = im1t[None, None]
+        im2t = im2t[None, None]
+
+    # [1,h,w,3] -> [1,3,h,w]
+    if im1t.shape[-1] == 3:
+        im1t = im1t.permute(0, 3, 1, 2)
+        im2t = im2t.permute(0, 3, 1, 2)
+
+    # img array: [1,h,w,3] imgt tensor: [1,3,h,w]
+    im1 = im1t.permute(0, 2, 3, 1).detach().cpu().numpy()
+    im2 = im2t.permute(0, 2, 3, 1).detach().cpu().numpy()
+    batchsz, hei, wid, _ = im1.shape
+
+    # batch processing
+    values = []
+    for i in range(batchsz):
+        if metric_name in ["mse", "psnr"]:
+            value = metric_method(im1[i], im2[i])
+        elif metric_name in ["ssim"]:
+            value, ssimmap = metric_method(im1[i], im2[i], channel_axis=-1, data_range=2, full=True)
+        elif metric_name in ["lpips"]:
+            value = metric_method(im1t[i : i + 1], im2t[i : i + 1])[0, 0, 0, 0]
+            value = value.detach().cpu().float().item()
+        values.append(value)
+
+    return sum(values) / len(values)

spikezoo-0.1/spikezoo/models/__init__.py

@@ -0,0 +1,37 @@
+import importlib
+import inspect
+from spikezoo.models.base_model import BaseModel,BaseModelConfig
+import os
+from pathlib import Path
+
+current_file_path = Path(__file__).parent
+files_list = os.listdir(os.path.dirname(os.path.abspath(__file__)))
+model_list = [file.split("_")[0] for file in files_list if file.endswith("_model.py")]
+
+# todo register function
+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
+    classes = sorted([name for name, obj in inspect.getmembers(module) if inspect.isclass(obj) and obj.__module__ == module.__name__])
+    model_cls: BaseModel = getattr(module, classes[0])
+    model = model_cls(cfg)
+    return model
+
+def build_model_name(model_name: str):
+    """Build the default dataset from the given name."""
+    # model module name
+    module_name = model_name + "_model"
+    assert model_name in model_list, f"Given model {model_name} not in our model zoo {model_list}."
+    module_name = "spikezoo.models." + module_name
+    module = importlib.import_module(module_name)
+    # model,model_config
+    classes = sorted([name for name, obj in inspect.getmembers(module) if inspect.isclass(obj) and obj.__module__ == module.__name__])
+    model_cls: BaseModel = getattr(module, classes[0])
+    model_cfg: BaseModelConfig = getattr(module, classes[1])()
+    model = model_cls(model_cfg)
+    return model