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

spikezoo/__init__.py

@@ -1,13 +1,29 @@
-from .utils.spike_utils import load_vidar_dat
+from .utils.spike_utils import *
 from .models import model_list
 from .datasets import dataset_list
 from .metrics import metric_all_names
 
-def get_datasets():
-    return dataset_list
+# 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"
 
-def get_models():
-    return model_list
+# DATASET NAME DEFINITION
+DATASETS = dataset_list
+class DATASET:
+    BASE = "base"
+    REDS_BASE = "reds_base"
+    REALWORLD = "realworld"
+    UHSR = "uhsr"
 
-def get_metrics():
-    return metric_all_names
+# METRIC NAME DEFINITION
+METRICS = metric_all_names

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/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/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/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/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/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/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/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/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/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/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)

spikezoo/datasets/base_dataset.py

@@ -6,7 +6,7 @@ import numpy as np
 from spikezoo.utils.spike_utils import load_vidar_dat
 import re
 from dataclasses import dataclass, replace
-from typing import Literal,Union
+from typing import Literal, Union
 import warnings
 import torch
 from tqdm import tqdm
@@ -19,7 +19,7 @@ class BaseDatasetConfig:
     "Dataset name."
     dataset_name: str = "base"
     "Directory specifying location of data."
-    root_dir: Union[str,Path] = Path(__file__).parent.parent / Path("data/base")
+    root_dir: Union[str, Path] = Path(__file__).parent.parent / Path("data/base")
     "Image width."
     width: int = 400
     "Image height."
@@ -46,6 +46,8 @@ class BaseDatasetConfig:
     use_cache: bool = False
     "Crop size."
     crop_size: tuple = (-1, -1)
+    "Rate. (-1 denotes variant)"
+    rate: float = 1
 
     # post process
     def __post_init__(self):
@@ -54,6 +56,7 @@ class BaseDatasetConfig:
         # 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):
@@ -80,7 +83,11 @@ class BaseDataset(Dataset):
         if self.cfg.use_aug == True and self.cfg.split == "train":
             spike, img = self.augmentor(spike, img)
 
-        batch = {"spike": spike, "img": img}
+        # rate
+        rate = self.cfg.rate
+
+        # ! spike and gt_img names are necessary
+        batch = {"spike": spike, "gt_img": img, "rate": rate}
         return batch
 
     # todo: To be overridden

spikezoo/datasets/realworld_dataset.py

@@ -4,7 +4,7 @@ from dataclasses import dataclass
 
 
 @dataclass
-class RealWorld_Config(BaseDatasetConfig):
+class RealWorldConfig(BaseDatasetConfig):
     dataset_name: str = "realworld"
     root_dir: Path = Path(__file__).parent.parent / Path("data/recVidarReal2019")
     width: int = 400
@@ -21,5 +21,3 @@ class RealWorld(BaseDataset):
     def prepare_data(self):
         self.spike_dir = self.cfg.root_dir
         self.spike_list = self.get_spike_files(self.spike_dir)
-        
-    

spikezoo/datasets/{reds_small_dataset.py → reds_base_dataset.py}

@@ -4,10 +4,11 @@ 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")
+class REDS_BASEConfig(BaseDatasetConfig):
+    dataset_name: str = "reds_base"
+    root_dir: Path = Path(__file__).parent.parent / Path("data/REDS_BASE")
     width: int = 400
     height: int = 250
     with_img: bool = True
@@ -15,13 +16,13 @@ class REDS_Small_Config(BaseDatasetConfig):
     spike_length_test: int = 301
     spike_dir_name: str = "spike"
     img_dir_name: str = "gt"
+    rate: float = 0.6
+
 
-class REDS_Small(BaseDataset):
+class REDS_BASE(BaseDataset):
     def __init__(self, cfg: BaseDatasetConfig):
-        super(REDS_Small, self).__init__(cfg)
+        super(REDS_BASE, 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]
-
+        self.img_list = [self.img_dir / Path(str(s.name).replace(".dat", ".png")) for s in self.spike_list]

spikezoo/datasets/reds_ssir_dataset.py

@@ -0,0 +1,181 @@
+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_SSIRConfig(BaseDatasetConfig):
+    dataset_name: str = "reds_ssir"
+    root_dir: Path = Path(__file__).parent.parent / Path("data/REDS_SSIR")
+    train_width: int = 96
+    train_height: int = 96
+    test_width: int = 1280
+    test_height: int = 720
+    width: int = -1
+    height: int = -1
+    with_img: bool = True
+    spike_length_train: int = 41
+    spike_length_test: int = 301
+
+    # post process
+    def __post_init__(self):
+        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}`."
+        # train/test split
+        if self.split == "train":
+            self.spike_length = self.spike_length_train
+            self.width = self.train_width
+            self.height = self.train_height
+        else:
+            self.spike_length = self.spike_length_test
+            self.width = self.test_width
+            self.height = self.test_height
+
+
+class REDS_SSIR(BaseDataset):
+    def __init__(self, cfg: BaseDatasetConfig):
+        super(REDS_SSIR, self).__init__(cfg)
+
+    def prepare_data(self):
+        """Specify the spike and image files to be loaded."""
+        # spike/imgs dir train/test
+        if self.cfg.split == "train":
+            self.img_dir = self.cfg.root_dir / Path("crop_mini/spike/train/interp_20_alpha_0.40")
+            self.spike_dir = self.cfg.root_dir / Path("crop_mini/image/train/train_orig")
+        else:
+            self.img_dir = self.cfg.root_dir / Path("imgs/val/val_orig")
+            self.spike_dir = self.cfg.root_dir / Path("spike/val/interp_20_alpha_0.40")
+        # get files
+        self.spike_list = self.get_spike_files(self.spike_dir)
+        self.img_list = []
+
+
+class sreds_train(torch.utils.data.Dataset):
+    def __init__(self, cfg):
+        self.cfg = cfg
+        self.pair_step = self.cfg["loader"]["pair_step"]
+        self.augmentor = Augmentor(crop_size=self.cfg["loader"]["crop_size"])
+        self.samples = self.collect_samples()
+        print("The samples num of training data: {:d}".format(len(self.samples)))
+
+    def confirm_exist(self, path_list_list):
+        for pl in path_list_list:
+            for p in pl:
+                if not osp.exists(p):
+                    return 0
+        return 1
+
+    def collect_samples(self):
+        spike_path = osp.join(
+            self.cfg["data"]["root"], "crop_mini", "spike", "train", "interp_{:d}_alpha_{:.2f}".format(self.cfg["data"]["interp"], self.cfg["data"]["alpha"])
+        )
+        image_path = osp.join(self.cfg["data"]["root"], "crop_mini", "image", "train", "train_orig")
+        scene_list = sorted(os.listdir(spike_path))
+        samples = []
+
+        for scene in scene_list:
+            spike_dir = osp.join(spike_path, scene)
+            image_dir = osp.join(image_path, scene)
+            spk_path_list = sorted(os.listdir(spike_dir))
+
+            spklen = len(spk_path_list)
+            seq_len = self.cfg["model"]["seq_len"] + 2
+            """
+            for st in range(0, spklen - ((spklen - self.pair_step) % seq_len) - seq_len, self.pair_step):
+                # 按照文件名称读取
+                spikes_path_list = [osp.join(spike_dir, spk_path_list[ii]) for ii in range(st, st+seq_len)]
+                images_path_list = [osp.join(image_dir, spk_path_list[ii][:-4]+'.png') for ii in range(st, st+seq_len)]
+                
+                if(self.confirm_exist([spikes_path_list, images_path_list])):
+                    s = {}
+                    s['spikes_paths'] = spikes_path_list
+                    s['images_paths'] = images_path_list
+                    samples.append(s)
+            """
+            # 按照文件名称读取
+            spikes_path_list = [osp.join(spike_dir, spk_path_list[ii]) for ii in range(spklen)]
+            images_path_list = [osp.join(image_dir, spk_path_list[ii][:-4] + ".png") for ii in range(spklen)]
+
+            if self.confirm_exist([spikes_path_list, images_path_list]):
+                s = {}
+                s["spikes_paths"] = spikes_path_list
+                s["images_paths"] = images_path_list
+                samples.append(s)
+
+        return samples
+
+    def _load_sample(self, s):
+        data = {}
+
+        data["spikes"] = [np.array(dat_to_spmat(p, size=(96, 96)), dtype=np.float32) for p in s["spikes_paths"]]
+        data["images"] = [read_img_gray(p) for p in s["images_paths"]]
+
+        data["spikes"], data["images"] = self.augmentor(data["spikes"], data["images"])
+        # print("data['spikes'][0].shape, data['images'][0].shape", data['spikes'][0].shape, data['images'][0].shape)
+
+        return data
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, index):
+        data = self._load_sample(self.samples[index])
+        return data
+
+
+class sreds_test(torch.utils.data.Dataset):
+    def __init__(self, cfg):
+        self.cfg = cfg
+        self.samples = self.collect_samples()
+        print("The samples num of testing data: {:d}".format(len(self.samples)))
+
+    def confirm_exist(self, path_list_list):
+        for pl in path_list_list:
+            for p in pl:
+                if not osp.exists(p):
+                    return 0
+        return 1
+
+    def collect_samples(self):
+        spike_path = osp.join(
+            self.cfg["data"]["root"], "spike", "val", "interp_{:d}_alpha_{:.2f}".format(self.cfg["data"]["interp"], self.cfg["data"]["alpha"])
+        )
+        image_path = osp.join(self.cfg["data"]["root"], "imgs", "val", "val_orig")
+        scene_list = sorted(os.listdir(spike_path))
+        samples = []
+
+        for scene in scene_list:
+            spike_dir = osp.join(spike_path, scene)
+            image_dir = osp.join(image_path, scene)
+            spk_path_list = sorted(os.listdir(spike_dir))
+
+            spklen = len(spk_path_list)
+            # seq_len = self.cfg['model']['seq_len']
+
+            # 按照文件名称读取
+            spikes_path_list = [osp.join(spike_dir, spk_path_list[ii]) for ii in range(spklen)]
+            images_path_list = [osp.join(image_dir, spk_path_list[ii][:-4] + ".png") for ii in range(spklen)]
+
+            if self.confirm_exist([spikes_path_list, images_path_list]):
+                s = {}
+                s["spikes_paths"] = spikes_path_list
+                s["images_paths"] = images_path_list
+                samples.append(s)
+
+        return samples
+
+    def _load_sample(self, s):
+        data = {}
+        data["spikes"] = [np.array(dat_to_spmat(p, size=(720, 1280)), dtype=np.float32) for p in s["spikes_paths"]]
+        data["images"] = [read_img_gray(p) for p in s["images_paths"]]
+        return data
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, index):
+        data = self._load_sample(self.samples[index])
+        return data