spikezoo 0.1__py3-none-any.whl

spikezoo/models/bsf_model.py

@@ -0,0 +1,90 @@
+import torch
+from dataclasses import dataclass
+from spikezoo.models.base_model import BaseModel, BaseModelConfig
+
+
+@dataclass
+class BSFConfig(BaseModelConfig):
+    # default params for BSF
+    model_name: str = "bsf"
+    model_file_name: str = "models.bsf.bsf"
+    model_cls_name: str = "BSF"
+    model_win_length: int = 61
+    require_params: bool = True
+    ckpt_path: str = "weights/bsf.pth"
+
+
+class BSF(BaseModel):
+    def __init__(self, cfg: BaseModelConfig):
+        super(BSF, self).__init__(cfg)
+
+    def preprocess_spike(self, spike):
+        # length
+        spike = self.crop_spike_length(spike)
+        # size
+        if self.spike_size == (250, 400):
+            spike = torch.cat([spike, spike[:, :, -2:]], dim=2)
+        elif self.spike_size == (480, 854):
+            spike = torch.cat([spike, spike[:, :, :, -2:]], dim=3)
+        # dsft
+        dsft = self.compute_dsft_core(spike)
+        dsft_dict = self.convert_dsft4(dsft, spike)
+        input_dict = {
+            "dsft_dict": dsft_dict,
+            "spikes": spike,
+        }
+        return input_dict
+
+    def postprocess_img(self, image):
+        if self.spike_size == (250, 400):
+            image = image[:, :, :250, :]
+        elif self.spike_size == (480, 854):
+            image = image[:, :, :, :854]
+        return image
+
+    def compute_dsft_core(self, spike):
+        bs, T, H, W = spike.shape
+        time = spike * torch.arange(T, device="cuda").reshape(1, T, 1, 1)
+        l_idx, _ = time.cummax(dim=1)
+        time[time == 0] = T
+        r_idx, _ = torch.flip(time, [1]).cummin(dim=1)
+        r_idx = torch.flip(r_idx, [1])
+        r_idx = torch.cat([r_idx[:, 1:, :, :], torch.ones([bs, 1, H, W], device="cuda") * T], dim=1)
+        res = r_idx - l_idx
+        res = torch.clip(res, 0)
+        return res
+
+    def convert_dsft4(self, dsft, spike):
+        b, T, h, w = spike.shape
+        dmls1 = -1 * torch.ones(spike.shape, device=spike.device, dtype=torch.float32)
+        dmrs1 = -1 * torch.ones(spike.shape, device=spike.device, dtype=torch.float32)
+        flag = -2 * torch.ones([b, h, w], device=spike.device, dtype=torch.float32)
+        for ii in range(T - 1, 0 - 1, -1):
+            flag += spike[:, ii] == 1
+            copy_pad_coord = flag < 0
+            dmls1[:, ii][copy_pad_coord] = dsft[:, ii][copy_pad_coord]
+            if ii < T - 1:
+                update_coord = (spike[:, ii + 1] == 1) * (~copy_pad_coord)
+                dmls1[:, ii][update_coord] = dsft[:, ii + 1][update_coord]
+                non_update_coord = (spike[:, ii + 1] != 1) * (~copy_pad_coord)
+                dmls1[:, ii][non_update_coord] = dmls1[:, ii + 1][non_update_coord]
+        flag = -2 * torch.ones([b, h, w], device=spike.device, dtype=torch.float32)
+        for ii in range(0, T, 1):
+            flag += spike[:, ii] == 1
+            copy_pad_coord = flag < 0
+            dmrs1[:, ii][copy_pad_coord] = dsft[:, ii][copy_pad_coord]
+            if ii > 0:
+                update_coord = (spike[:, ii] == 1) * (~copy_pad_coord)
+                dmrs1[:, ii][update_coord] = dsft[:, ii - 1][update_coord]
+                non_update_coord = (spike[:, ii] != 1) * (~copy_pad_coord)
+                dmrs1[:, ii][non_update_coord] = dmrs1[:, ii - 1][non_update_coord]
+        dsft12 = dsft + dmls1
+        dsft21 = dsft + dmrs1
+        dsft22 = dsft + dmls1 + dmrs1
+        dsft_dict = {
+            "dsft11": dsft,
+            "dsft12": dsft12,
+            "dsft21": dsft21,
+            "dsft22": dsft22,
+        }
+        return dsft_dict

spikezoo/models/spcsnet_model.py

@@ -0,0 +1,19 @@
+from dataclasses import dataclass, field
+from spikezoo.models.base_model import BaseModel, BaseModelConfig
+from typing import List
+
+
+@dataclass
+class SPCSNetConfig(BaseModelConfig):
+    # default params for WGSE
+    model_name: str = "spcsnet"
+    model_file_name: str = "models"
+    model_cls_name: str = "SPCS_Net"
+    model_win_length: int = 41
+    require_params: bool = True
+    ckpt_path: str = 'weights/spcsnet.pth'
+
+
+class SPCSNet(BaseModel):
+    def __init__(self, cfg: BaseModelConfig):
+        super(SPCSNet, self).__init__(cfg)

spikezoo/models/spikeclip_model.py

@@ -0,0 +1,32 @@
+from dataclasses import dataclass
+from spikezoo.models.base_model import BaseModel, BaseModelConfig
+import torch
+import torch.nn.functional as F
+
+
+@dataclass
+class SpikeCLIPConfig(BaseModelConfig):
+    # default params for SpikeCLIP
+    model_name: str = "spikeclip"
+    model_file_name: str = "nets"
+    model_cls_name: str = "LRN"
+    model_win_length: int = 200
+    require_params: bool = True
+    ckpt_path: str = "weights/spikeclip.pth"
+
+
+class SpikeCLIP(BaseModel):
+    def __init__(self, cfg: BaseModelConfig):
+        super(SpikeCLIP, self).__init__(cfg)
+
+    def preprocess_spike(self, spike):
+        # length
+        spike = self.crop_spike_length(spike)
+        # voxel
+        voxel = torch.sum(spike.reshape(-1, 50, 4, spike.shape[-2], spike.shape[-1]), axis=2)  # [200,224,224] -> [50,224,224]
+        voxel = F.pad(voxel, pad=(20, 20, 20, 20), mode="reflect", value=0)
+        return voxel
+
+    def postprocess_img(self, image):
+        image = image[:, :, 20:-20, 20:-20]
+        return image

spikezoo/models/spikeformer_model.py

@@ -0,0 +1,50 @@
+import torch
+from dataclasses import dataclass, field
+from spikezoo.models.base_model import BaseModel, BaseModelConfig
+
+
+@dataclass
+class SpikeFormerConfig(BaseModelConfig):
+    # default params for SpikeFormer
+    model_name: str = "spikeformer"
+    model_file_name: str = "Model.SpikeFormer"
+    model_cls_name: str = "SpikeFormer"
+    model_win_length: int = 65
+    require_params: bool = True
+    ckpt_path: str = "weights/spikeformer.pth"
+    model_params: dict = field(
+        default_factory=lambda: {
+            "inputDim": 65,
+            "dims": (32, 64, 160, 256),
+            "heads": (1, 2, 5, 8),
+            "ff_expansion": (8, 8, 4, 4),
+            "reduction_ratio": (8, 4, 2, 1),
+            "num_layers": 2,
+            "decoder_dim": 256,
+            "out_channel": 1,
+        }
+    )
+
+
+class SpikeFormer(BaseModel):
+    def __init__(self, cfg: BaseModelConfig):
+        super(SpikeFormer, self).__init__(cfg)
+
+    def preprocess_spike(self, spike):
+        # length
+        spike = self.crop_spike_length(spike)
+        # size
+        if self.spike_size == (250, 400):
+            spike = torch.cat([spike[:, :, :3, :], spike, spike[:, :, -3:, :]], dim=2)
+        elif self.spike_size == (480, 854):
+            spike = torch.cat([spike, spike[:, :, :, -2:]], dim=3)
+        # input
+        spike = 2 * spike - 1
+        return spike
+
+    def postprocess_img(self, image):
+        if self.spike_size == (250, 400):
+            image = image[:, :, 3:-3, :]
+        elif self.spike_size == (480, 854):
+            image = image[:, :, :, :854]
+        return image

spikezoo/models/spk2imgnet_model.py

@@ -0,0 +1,51 @@
+import torch
+from dataclasses import dataclass, field
+from spikezoo.models.base_model import BaseModel, BaseModelConfig
+
+
+@dataclass
+class Spk2ImgNetConfig(BaseModelConfig):
+    # default params for Spk2ImgNet
+    model_name: str = "spk2imgnet"
+    model_file_name: str = "nets"
+    model_cls_name: str = "SpikeNet"
+    model_win_length: int = 41
+    require_params: bool = True
+    ckpt_path: str = "weights/spk2imgnet.pth"
+    light_correction: bool = False
+
+    # model params
+    model_params: dict = field(
+        default_factory=lambda: {
+            "in_channels": 13,
+            "features": 64,
+            "out_channels": 1,
+            "win_r": 6,
+            "win_step": 7,
+        }
+    )
+
+
+class Spk2ImgNet(BaseModel):
+    def __init__(self, cfg: BaseModelConfig):
+        super(Spk2ImgNet, self).__init__(cfg)
+
+    def preprocess_spike(self, spike):
+        # length
+        spike = self.crop_spike_length(spike)
+        # size
+        if self.spike_size == (250, 400):
+            spike = torch.cat([spike, spike[:, :, -2:]], dim=2)
+        elif self.spike_size == (480, 854):
+            spike = torch.cat([spike, spike[:, :, :, -2:]], dim=3)
+        return spike
+
+    def postprocess_img(self, image):
+        if self.spike_size == (250, 400):
+            image = image[:, :, :250, :]
+        elif self.spike_size == (480, 854):
+            image = image[:, :, :, :854]
+        # used on the REDS_small dataset.
+        if self.cfg.light_correction == True:
+            image = torch.clamp(image / 0.6, 0, 1)
+        return image

spikezoo/models/ssir_model.py

@@ -0,0 +1,22 @@
+from dataclasses import dataclass
+from spikezoo.models.base_model import BaseModel, BaseModelConfig
+
+
+@dataclass
+class SSIRConfig(BaseModelConfig):
+    # default params for SSIR
+    model_name: str = "ssir"
+    model_file_name: str = "models.networks"
+    model_cls_name: str = "SSIR"
+    model_win_length: int = 41
+    require_params: bool = True
+    ckpt_path: str = "weights/ssir.pth"
+
+
+class SSIR(BaseModel):
+    def __init__(self, cfg: BaseModelConfig):
+        super(SSIR, self).__init__(cfg)
+
+    def postprocess_img(self, image):
+        # image = image[0]
+        return image

spikezoo/models/ssml_model.py

@@ -0,0 +1,18 @@
+from dataclasses import dataclass
+from spikezoo.models.base_model import BaseModel, BaseModelConfig
+
+
+@dataclass
+class SSMLConfig(BaseModelConfig):
+    # default params for SSML
+    model_name: str = "ssml"
+    model_file_name: str = "model"
+    model_cls_name: str = "DoubleNet"
+    model_win_length: int = 41
+    require_params: bool = True
+    ckpt_path: str = 'weights/ssml.pt'
+
+
+class SSML(BaseModel):
+    def __init__(self, cfg: BaseModelConfig):
+        super(SSML, self).__init__(cfg)

spikezoo/models/stir_model.py

@@ -0,0 +1,37 @@
+import torch
+from dataclasses import dataclass
+from spikezoo.models.base_model import BaseModel, BaseModelConfig
+
+
+@dataclass
+class STIRConfig(BaseModelConfig):
+    # default params for SSIR
+    model_name: str = "stir"
+    model_file_name: str = "models.networks_STIR"
+    model_cls_name: str = "STIR"
+    model_win_length: int = 61
+    require_params: bool = True
+    ckpt_path: str = "weights/stir.pth"
+
+
+class STIR(BaseModel):
+    def __init__(self, cfg: BaseModelConfig):
+        super(STIR, self).__init__(cfg)
+
+    def preprocess_spike(self, spike):
+        # length
+        spike = self.crop_spike_length(spike)
+        # size
+        if self.spike_size == (250, 400):
+            spike = torch.cat([spike, spike[:, :, -6:]], dim=2)
+        elif self.spike_size == (480, 854):
+            spike = torch.cat([spike, spike[:, :, :, -10:]], dim=3)
+        return spike
+
+    def postprocess_img(self, image):
+        # recon, Fs_lv_0, Fs_lv_1, Fs_lv_2, Fs_lv_3, Fs_lv_4, Est = image
+        if self.spike_size == (250, 400):
+            image = image[:, :, :250, :]
+        elif self.spike_size == (480, 854):
+            image = image[:, :, :, :854]
+        return image

spikezoo/models/tfi_model.py

@@ -0,0 +1,18 @@
+from dataclasses import dataclass, field
+from spikezoo.models.base_model import BaseModel, BaseModelConfig
+
+
+@dataclass
+class TFIConfig(BaseModelConfig):
+    # default params for TFI
+    model_name: str = "tfi"
+    model_file_name: str = "nets"
+    model_cls_name: str = "TFIModel"
+    model_win_length: int = 41
+    require_params: bool = False
+    model_params: dict = field(default_factory=lambda: {"model_win_length": 41})
+
+
+class TFI(BaseModel):
+    def __init__(self, cfg: BaseModelConfig):
+        super(TFI, self).__init__(cfg)

spikezoo/models/tfp_model.py

@@ -0,0 +1,18 @@
+from dataclasses import dataclass,field
+from spikezoo.models.base_model import BaseModel, BaseModelConfig
+
+
+@dataclass
+class TFPConfig(BaseModelConfig):
+    # default params for TFP
+    model_name: str = "tfp"
+    model_file_name: str = "nets"
+    model_cls_name: str = "TFPModel"
+    model_win_length: int = 41
+    require_params: bool = False
+    model_params: dict = field(default_factory=lambda: {"model_win_length": 41})
+
+
+class TFP(BaseModel):
+    def __init__(self, cfg: BaseModelConfig):
+        super(TFP, self).__init__(cfg)

spikezoo/models/wgse_model.py

@@ -0,0 +1,31 @@
+from dataclasses import dataclass, field
+from spikezoo.models.base_model import BaseModel, BaseModelConfig
+from typing import List
+
+
+@dataclass
+class WGSEConfig(BaseModelConfig):
+    # default params for WGSE
+    model_name: str = "wgse"
+    model_file_name: str = "dwtnets"
+    model_cls_name: str = "Dwt1dResnetX_TCN"
+    model_win_length: int = 41
+    require_params: bool = True
+    ckpt_path: str = "weights/wgse.pt"
+    model_params: dict = field(
+        default_factory=lambda: {
+            "wvlname": "db8",
+            "J": 5,
+            "yl_size": "15",
+            "yh_size": [28, 21, 18, 16, 15],
+            "num_residual_blocks": 3,
+            "norm": None,
+            "ks": 3,
+            "store_features": True,
+        }
+    )
+
+
+class WGSE(BaseModel):
+    def __init__(self, cfg: BaseModelConfig):
+        super(WGSE, self).__init__(cfg)

spikezoo/pipeline/__init__.py

@@ -0,0 +1,4 @@
+from .base_pipeline import Pipeline, PipelineConfig
+from .ensemble_pipeline import EnsemblePipelineConfig, EnsemblePipeline
+from .train_pipeline import TrainPipelineConfig, TrainPipeline
+

spikezoo/pipeline/base_pipeline.py

@@ -0,0 +1,267 @@
+import torch
+from dataclasses import dataclass, field
+import os
+from spikezoo.utils.img_utils import tensor2npy, AverageMeter
+from spikezoo.utils.spike_utils import load_vidar_dat
+from spikezoo.metrics import cal_metric_pair, cal_metric_single
+import numpy as np
+import cv2
+from pathlib import Path
+from enum import Enum, auto
+from typing import Literal
+from spikezoo.metrics import metric_pair_names, metric_single_names, metric_all_names
+from thop import profile
+import time
+from datetime import datetime
+from spikezoo.utils import setup_logging, save_config
+from tqdm import tqdm
+from spikezoo.models import build_model_cfg, build_model_name, BaseModel, BaseModelConfig
+from spikezoo.datasets import build_dataset_cfg, build_dataset_name, BaseDataset, BaseDatasetConfig, build_dataloader
+from typing import Optional, Union, List
+
+
+@dataclass
+class PipelineConfig:
+    "Evaluate metrics or not."
+    save_metric: bool = True
+    "Save recoverd images or not."
+    save_img: bool = True
+    "Normalizing recoverd images or not."
+    save_img_norm: bool = False
+    "Normalizing gt or not."
+    gt_img_norm: bool = False
+    "Save folder for the code running result."
+    save_folder: str = ""
+    "Saved experiment name."
+    exp_name: str = ""
+    "Metric names for evaluation."
+    metric_names: List[str] = field(default_factory=lambda: ["psnr", "ssim"])
+    "Different modes for the pipeline."
+    _mode: Literal["single_mode", "multi_mode", "train_mode"] = "single_mode"
+
+
+class Pipeline:
+    def __init__(
+        self,
+        cfg: PipelineConfig,
+        model_cfg: Union[str, BaseModelConfig],
+        dataset_cfg: Union[str, BaseDatasetConfig],
+    ):
+        self.cfg = cfg
+        self._setup_model_data(model_cfg, dataset_cfg)
+        self._setup_pipeline()
+
+    def _setup_model_data(self, model_cfg, dataset_cfg):
+        """Model and Data setup."""
+        # model
+        self.model: BaseModel = build_model_name(model_cfg) if isinstance(model_cfg, str) else build_model_cfg(model_cfg)
+        self.model = self.model.eval()
+        torch.set_grad_enabled(False)
+        # dataset
+        self.dataset: BaseDataset = build_dataset_name(dataset_cfg) if isinstance(dataset_cfg, str) else build_dataset_cfg(dataset_cfg)
+        self.dataloader = build_dataloader(self.dataset)
+        # device
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    def _setup_pipeline(self):
+        """Pipeline setup."""
+        # save folder
+        self.thistime = datetime.now().strftime("%Y-%m-%d_%H-%M-%S-%f")[:23]
+        self.save_folder = Path(__file__).parent.parent / Path(f"results") if len(self.cfg.save_folder) == 0 else self.cfg.save_folder
+        mode_name = "train" if self.cfg._mode == "train_mode" else "detect"
+        self.save_folder = (
+            self.save_folder / Path(f"{mode_name}/{self.thistime}")
+            if len(self.cfg.exp_name) == 0
+            else self.save_folder / Path(f"{mode_name}/{self.cfg.exp_name}")
+        )
+        save_folder = self.save_folder
+        os.makedirs(str(save_folder), exist_ok=True)
+        # logger result
+        self.logger = setup_logging(save_folder / Path("result.log"))
+        self.logger.info(f"Info logs are saved on the {save_folder}/result.log")
+        # pipeline config
+        save_config(self.cfg, save_folder / Path("cfg_pipeline.log"))
+        # model config
+        if self.cfg._mode == "single_mode":
+            save_config(self.model.cfg, save_folder / Path("cfg_model.log"))
+        elif self.cfg._mode == "multi_mode":
+            for model in self.model_list:
+                save_config(model.cfg, save_folder / Path("cfg_model.log"), mode="a")
+        # dataset config
+        save_config(self.dataset.cfg, save_folder / Path("cfg_dataset.log"))
+
+    def spk2img_from_dataset(self, idx=0):
+        """Func---Save the recoverd image and calculate the metric from the given dataset."""
+        # save folder
+        save_folder = self.save_folder / Path(f"spk2img_from_dataset/{self.dataset.cfg.dataset_name}_dataset/{self.dataset.cfg.split}/{idx:06d}")
+        os.makedirs(str(save_folder), exist_ok=True)
+
+        # data process
+        batch = self.dataset[idx]
+        spike, img = batch["spike"], batch["img"]
+        spike = spike[None].to(self.device)
+        if self.dataset.cfg.with_img == True:
+            img = img[None].to(self.device)
+        else:
+            img = None
+        return self._spk2img(spike, img, save_folder)
+
+    def spk2img_from_file(self, file_path, height, width, img_path=None, remove_head=False):
+        """Func---Save the recoverd image and calculate the metric from the given input file."""
+        # save folder
+        save_folder = self.save_folder / Path(f"spk2img_from_file/{os.path.basename(file_path)}")
+        os.makedirs(str(save_folder), exist_ok=True)
+
+        # load spike from .dat
+        if file_path.endswith(".dat"):
+            spike = load_vidar_dat(file_path, height, width, remove_head)
+        # load spike from .npz from UHSR
+        elif file_path.endswith("npz"):
+            spike = np.load(file_path)["spk"].astype(np.float32)[:, 13:237, 13:237]
+        else:
+            raise RuntimeError("Not recognized spike input file.")
+        # load img from .png/.jpg image file
+        if img_path is not None:
+            img = cv2.imread(img_path)
+            if img.ndim == 3:
+                img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+            img = (img / 255).astype(np.float32)
+            img = torch.from_numpy(img)[None, None].to(self.device)
+        else:
+            img = img_path
+        spike = torch.from_numpy(spike)[None].to(self.device)
+        return self._spk2img(spike, img, save_folder)
+
+    def spk2img_from_spk(self, spike, img=None):
+        """Func---Save the recoverd image and calculate the metric from the given spike stream."""
+        # save folder
+        save_folder = self.save_folder / Path(f"spk2img_from_spk/{self.thistime}")
+        os.makedirs(str(save_folder), exist_ok=True)
+
+        # spike process
+        if isinstance(spike, np.ndarray):
+            spike = torch.from_numpy(spike)
+        spike = spike.to(self.device)
+        # [c,h,w] -> [1,c,w,h]
+        if spike.dim() == 3:
+            spike = spike[None]
+        spike = spike.float()
+        # img process
+        if img is not None:
+            if isinstance(img, np.ndarray):
+                img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) if img.ndim == 3 else img
+                img = (img / 255).astype(np.float32)
+                img = torch.from_numpy(img)[None, None].to(self.device)
+            else:
+                raise RuntimeError("Not recognized image input type.")
+        return self._spk2img(spike, img, save_folder)
+
+    def save_imgs_from_dataset(self):
+        """Func---Save all images from the given dataset."""
+        for idx in range(len(self.dataset)):
+            self.spk2img_from_dataset(idx=idx)
+
+    # TODO: To be overridden
+    def cal_params(self):
+        """Func---Calculate the parameters/flops/latency of the given method."""
+        self._cal_prams_model(self.model)
+
+    # TODO: To be overridden
+    def cal_metrics(self):
+        """Func---Calculate the metric of the given method."""
+        self._cal_metrics_model(self.model)
+
+    # TODO: To be overridden
+    def _spk2img(self, spike, img, save_folder):
+        """Spike-to-image: spike:[bs,c,h,w] (0-1), img:[bs,1,h,w] (0-1)"""
+        return self._spk2img_model(self.model, spike, img, save_folder)
+
+    def _spk2img_model(self, model, spike, img, save_folder):
+        """Spike-to-image from the given model."""
+        # spike2image conversion
+        model_name = model.cfg.model_name
+        recon_img = model(spike)
+        recon_img_copy = recon_img.clone()
+        # normalization
+        recon_img, img = self._post_process_img(model_name, recon_img, img)
+        # metric
+        if self.cfg.save_metric == True:
+            self.logger.info(f"----------------------Method: {model_name.upper()}----------------------")
+            # paired metric
+            for metric_name in metric_all_names:
+                if img is not None and metric_name in metric_pair_names:
+                    self.logger.info(f"{metric_name.upper()}: {cal_metric_pair(recon_img,img,metric_name)}")
+                elif metric_name in metric_single_names:
+                    self.logger.info(f"{metric_name.upper()}: {cal_metric_single(recon_img,metric_name)}")
+                else:
+                    self.logger.info(f"{metric_name.upper()} not calculated since no ground truth provided.")
+
+        # visual
+        if self.cfg.save_img == True:
+            recon_img = tensor2npy(recon_img[0, 0])
+            cv2.imwrite(f"{save_folder}/{model.cfg.model_name}.png", recon_img)
+            if img is not None:
+                img = tensor2npy(img[0, 0])
+                cv2.imwrite(f"{save_folder}/sharp_img.png", img)
+        self.logger.info(f"Images are saved on the {save_folder}")
+
+        return recon_img_copy
+
+    def _post_process_img(self, model_name, recon_img, gt_img):
+        """Post process the reconstructed image."""
+        # TFP and TFI algorithms are normalized automatically, others are normalized based on the self.cfg.use_norm
+        if model_name in ["tfp", "tfi", "spikeformer", "spikeclip"]:
+            recon_img = (recon_img - recon_img.min()) / (recon_img.max() - recon_img.min())
+        elif self.cfg.save_img_norm == True:
+            recon_img = (recon_img - recon_img.min()) / (recon_img.max() - recon_img.min())
+        recon_img = recon_img.clip(0, 1)
+        gt_img = (gt_img - gt_img.min()) / (gt_img.max() - gt_img.min()) if self.cfg.gt_img_norm == True and gt_img is not None else gt_img
+        return recon_img, gt_img
+
+    def _cal_metrics_model(self, model: BaseModel):
+        """Calculate the metrics for the given model."""
+        # metrics construct
+        model_name = model.cfg.model_name
+        metrics_dict = {}
+        for metric_name in self.cfg.metric_names:
+            if (self.dataset.cfg.with_img == True) or (metric_name in metric_single_names):
+                metrics_dict[metric_name] = AverageMeter()
+
+        # metrics calculate
+        for batch_idx, batch in enumerate(tqdm(self.dataloader)):
+            batch = model.feed_to_device(batch)
+            outputs = model.get_outputs_dict(batch)
+            recon_img, img = model.get_paired_imgs(batch, outputs)
+            recon_img, img = self._post_process_img(model_name, recon_img, img)
+            for metric_name in metrics_dict.keys():
+                if metric_name in metric_pair_names:
+                    metrics_dict[metric_name].update(cal_metric_pair(recon_img, img, metric_name))
+                elif metric_name in metric_single_names:
+                    metrics_dict[metric_name].update(cal_metric_single(recon_img, metric_name))
+
+        # metrics self.logger.info
+        self.logger.info(f"----------------------Method: {model_name.upper()}----------------------")
+        for metric_name in metrics_dict.keys():
+            self.logger.info(f"{metric_name.upper()}: {metrics_dict[metric_name].avg}")
+
+    def _cal_prams_model(self, model):
+        """Calculate the parameters for the given model."""
+        network = model.net
+        model_name = model.cfg.model_name.upper()
+        # params
+        params = sum(p.numel() for p in network.parameters())
+        # latency
+        spike = torch.zeros((1, 200, 250, 400)).cuda()
+        start_time = time.time()
+        for _ in range(100):
+            model(spike)
+        latency = (time.time() - start_time) / 100
+        # flop # todo thop bug for BSF
+        flops, _ = profile((model), inputs=(spike,))
+        re_msg = (
+            "Total params: %.4fM" % (params / 1e6),
+            "FLOPs:" + str(flops / 1e9) + "{}".format("G"),
+            "Latency: {:.6f} seconds".format(latency),
+        )
+        self.logger.info(f"----------------------Method: {model_name}----------------------")
+        self.logger.info(re_msg)