spikezoo 0.1__py3-none-any.whl

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

spikezoo/models/base_model.py

@@ -0,0 +1,177 @@
+import torch
+import torch.nn as nn
+import importlib
+import inspect
+from dataclasses import dataclass, field
+from spikezoo.utils import load_network, download_file
+import os
+import time
+from typing import Dict
+from torch.optim import Adam
+from torch.optim.lr_scheduler import CosineAnnealingLR
+import functools
+
+
+# todo private design
+@dataclass
+class BaseModelConfig:
+    # default params for BaseModel
+    "Registerd model name."
+    model_name: str = "base"
+    "File name of the specified model."
+    model_file_name: str = "nets"
+    "Class name of the specified model in spikezoo/archs/base/{model_file_name}.py."
+    model_cls_name: str = "BaseNet"
+    "Spike input length for the specified model."
+    model_win_length: int = 41
+    "Model require model parameters or not."
+    require_params: bool = False
+    "Model stored path."
+    ckpt_path: str = ""
+    "Load pretrained weights or not."
+    load_state: bool = True
+    "Base url for storing pretrained models."
+    base_url: str = "https://github.com/chenkang455/Spike-Zoo/releases/download/v0.1/"
+    "Multi-GPU setting."
+    multi_gpu: bool = False
+    "Model parameters."
+    model_params: dict = field(default_factory=lambda: {})
+
+
+class BaseModel(nn.Module):
+    def __init__(self, cfg: BaseModelConfig):
+        super(BaseModel, self).__init__()
+        self.cfg = cfg
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.net = self.build_network().to(self.device)
+        self.net = nn.DataParallel(self.net) if cfg.multi_gpu == True else self.net
+        self.spike_size = None
+        self.model_half_win_length: int = cfg.model_win_length // 2
+
+    # ! Might lead to low speed training on the BSF.
+    def forward(self, spike):
+        """A simple implementation for the spike-to-image conversion, given the spike input and output the reconstructed image."""
+        spike = self.preprocess_spike(spike)
+        img = self.net(spike)
+        img = self.postprocess_img(img)
+        return img
+
+    def build_network(self):
+        """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)
+        model = model_cls(**self.cfg.model_params)
+        if self.cfg.load_state and self.cfg.require_params:
+            load_folder = os.path.dirname(os.path.abspath(__file__))
+            weight_path = os.path.join(load_folder, self.cfg.ckpt_path)
+            if os.path.exists(weight_path) == False:
+                os.makedirs(os.path.dirname(weight_path), exist_ok=True)
+                self.download_weight(weight_path)
+                time.sleep(0.5)
+            model = load_network(weight_path, model)
+        return model
+
+    def save_network(self, save_path):
+        """Save the network."""
+        network = self.net
+        if isinstance(network, nn.DataParallel):
+            network = network.module
+        state_dict = network.state_dict()
+        for key, param in state_dict.items():
+            state_dict[key] = param.cpu()
+        torch.save(state_dict, save_path)
+
+    def download_weight(self, weight_path):
+        """Download the pretrained weight from the given url."""
+        url = self.cfg.base_url + os.path.basename(self.cfg.ckpt_path)
+        download_file(url, weight_path)
+
+    def crop_spike_length(self, spike):
+        """Crop the spike length."""
+        spike_length = spike.shape[1]
+        spike_mid = spike_length // 2
+        assert spike_length >= self.cfg.model_win_length, f"Spike input is not long enough, given {spike_length} frames < {self.cfg.model_win_length}."
+        # even length
+        if self.cfg.model_win_length == self.model_half_win_length * 2:
+            spike = spike[
+                :,
+                spike_mid - self.model_half_win_length : spike_mid + self.model_half_win_length,
+            ]
+        # odd length
+        else:
+            spike = spike[
+                :,
+                spike_mid - self.model_half_win_length : spike_mid + self.model_half_win_length + 1,
+            ]
+        if self.spike_size == None:
+            self.spike_size = (spike.shape[2], spike.shape[3])
+        return spike
+
+    def preprocess_spike(self, spike):
+        """Preprocess the spike (length size)."""
+        spike = self.crop_spike_length(spike)
+        return spike
+
+    def postprocess_img(self, image):
+        """Postprocess the image."""
+        return image
+
+    # -------------------- Training Part --------------------
+    def setup_training(self, pipeline_cfg):
+        """Setup training optimizer and loss."""
+        from spikezoo.pipeline import TrainPipelineConfig
+
+        cfg: TrainPipelineConfig = pipeline_cfg
+        self.optimizer = Adam(self.net.parameters(), lr=cfg.lr, betas=(0.9, 0.99), weight_decay=0)
+        self.scheduler = CosineAnnealingLR(self.optimizer, T_max=cfg.epochs, eta_min=0)
+        self.criterion = nn.L1Loss()
+
+    def get_outputs_dict(self, batch):
+        """Get the output dict for the given input batch. (Designed for the training mode considering possible auxiliary output.)"""
+        # data process
+        spike = batch["spike"]
+        # outputs
+        outputs = {}
+        recon_img = self(spike)
+        outputs["recon_img"] = recon_img
+        return outputs
+
+    def get_visual_dict(self, batch, outputs):
+        """Get the visual dict from the given input batch and outputs."""
+        visual_dict = {}
+        visual_dict["recon"] = outputs["recon_img"]
+        visual_dict["img"] = batch["img"]
+        return visual_dict
+
+    def get_loss_dict(self, outputs, batch):
+        """Get the loss dict from the given input batch and outputs."""
+        # data process
+        gt_img = batch["img"]
+        # recon image
+        recon_img = outputs["recon_img"]
+        # loss dict
+        loss_dict = {}
+        loss_dict["l1"] = self.criterion(recon_img, gt_img)
+        loss_values_dict = {k: v.item() for k, v in loss_dict.items()}
+        return loss_dict,loss_values_dict
+
+    def get_paired_imgs(self, batch, outputs):
+        recon_img = outputs["recon_img"]
+        img = batch["img"]
+        return recon_img, img
+
+    def optimize_parameters(self, loss_dict):
+        """Optimize the parameters from the loss_dict."""
+        loss = functools.reduce(torch.add, loss_dict.values())
+        self.optimizer.zero_grad()
+        loss.backward()
+        self.optimizer.step()
+
+    def update_learning_rate(self):
+        """Update the learning rate."""
+        self.scheduler.step()
+
+    def feed_to_device(self, batch):
+        batch = {k: v.to(self.device, non_blocking=True) if torch.is_tensor(v) else v for k, v in batch.items()}
+        return batch