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

spikezoo/datasets/szdata_dataset.py

@@ -5,8 +5,9 @@ import cv2
 import torch
 import numpy as np
 
+# todo tobe evaluated
 @dataclass
-class SZData_Config(BaseDatasetConfig):
+class SZDataConfig(BaseDatasetConfig):
     dataset_name: str = "szdata"
     root_dir: Path = Path(__file__).parent.parent / Path("data/dataset")
     width: int = 400
@@ -21,17 +22,6 @@ 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
-    
+    def prepare_data(self):
+        super().prepare_data()
+        self.img_list = [self.img_dir / Path(str(s.name).replace('.dat','.png')) for s in self.spike_list]

spikezoo/datasets/uhsr_dataset.py

@@ -4,8 +4,9 @@ from dataclasses import dataclass
 import numpy as np
 import torch
 
+
 @dataclass
-class UHSR_Config(BaseDatasetConfig):
+class UHSRConfig(BaseDatasetConfig):
     dataset_name: str = "uhsr"
     root_dir: Path = Path(__file__).parent.parent / Path("data/U-CALTECH")
     width: int = 224
@@ -29,10 +30,10 @@ class UHSR(BaseDataset):
         files = path.glob("**/*.npz")
         return sorted(files)
 
-    def load_spike(self,idx):
+    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
+        return spike

spikezoo/models/__init__.py

@@ -1,6 +1,7 @@
 import importlib
 import inspect
 from spikezoo.models.base_model import BaseModel,BaseModelConfig
+from spikezoo.utils.other_utils import getattr_case_insensitive
 import os
 from pathlib import Path
 
@@ -17,8 +18,9 @@ def build_model_cfg(cfg: BaseModelConfig):
     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_name = cfg.model_name
+    model_name = model_name + 'Model' if model_name == "base" else model_name
+    model_cls: BaseModel = getattr_case_insensitive(module,model_name)
     model = model_cls(cfg)
     return model
 
@@ -30,8 +32,8 @@ def build_model_name(model_name: str):
     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_name = model_name + 'Model' if model_name == "base" else model_name
+    model_cls: BaseModel = getattr_case_insensitive(module,model_name)
+    model_cfg: BaseModelConfig = getattr_case_insensitive(module, model_name + 'config')()
     model = model_cls(model_cfg)
-    return model
+    return model

spikezoo/models/base_model.py

@@ -6,36 +6,44 @@ from dataclasses import dataclass, field
 from spikezoo.utils import load_network, download_file
 import os
 import time
-from typing import Dict
+from typing import Dict, Literal
 from torch.optim import Adam
 from torch.optim.lr_scheduler import CosineAnnealingLR
 import functools
+import torch.nn as nn
+from typing import Optional, Union, List
+from spikezoo.archs.base.nets import BaseNet 
 
 
 # todo private design
 @dataclass
 class BaseModelConfig:
-    # default params for BaseModel
+    # ------------- Not Recommended to Change -------------
     "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
+    "Spike input length. (local mode)"
+    model_length: int = 41
+    "Spike input length for different versions."
+    model_length_dict: dict = field(default_factory=lambda: {"v010": 41, "v023": 41})
     "Model require model parameters or not."
-    require_params: bool = False
-    "Model stored path."
+    require_params: bool = True
+    "Model parameters. (local mode)"
+    model_params: dict = field(default_factory=lambda: {})
+    "Model parameters for different versions."
+    model_params_dict: dict = field(default_factory=lambda: {"v010": {}, "v023": {}})
+    # ------------- Config -------------
+    "Load ckpt path. Used on the local mode."
     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/"
+    "Load pretrained weights or not. (default false, set to true during the evaluation mode.)"
+    load_state: bool = False
     "Multi-GPU setting."
     multi_gpu: bool = False
-    "Model parameters."
-    model_params: dict = field(default_factory=lambda: {})
+    "Base url."
+    base_url: str = "https://github.com/chenkang455/Spike-Zoo/releases/download"
 
 
 class BaseModel(nn.Module):
@@ -43,33 +51,69 @@ class BaseModel(nn.Module):
         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.model_half_win_length: int = cfg.model_win_length // 2
+        self.loss_func_cache = {}
 
-    # ! 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."""
+        return self.spk2img(spike)
+
+    # ! Might lead to low speed training on the BSF.
+    def spk2img(self, spike):
+        """A simple implementation for the spike-to-image conversion (**tailored for the evaluation mode**), 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):
+    def build_network(
+        self,
+        mode: Literal["debug", "train", "eval"] = "debug",
+        version: Literal["local", "v010", "v023"] = "local",
+    ):
         """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)
+        # load model config parameters
+        if version == "local":
+            model = model_cls(**self.cfg.model_params)
+            self.model_length = self.cfg.model_length
+            self.model_half_length = self.model_length // 2
+        else:
+            model = model_cls(**self.cfg.model_params_dict[version])
+            self.model_length = self.cfg.model_length_dict[version]
+            self.model_half_length = self.model_length // 2
+        model.train() if mode == "train" else model.eval()
+        # auto set the load_state to True under the eval mode
+        if mode == "eval" and self.cfg.load_state == False:
+            print(f"Method {self.cfg.model_name} on the evaluation mode, load_state is set to True automatically.")
+            self.cfg.load_state = True
+        # load model
         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)
+            # load from the url version
+            if version != "local":
+                load_folder = os.path.dirname(os.path.abspath(__file__))
+                ckpt_name = f"{self.cfg.model_name}.{get_suffix(self.cfg.model_name,version)}"
+                ckpt_path = os.path.join("weights",version,ckpt_name)
+                ckpt_path = os.path.join(load_folder, ckpt_path)
+                ckpt_path_url = os.path.join(self.cfg.base_url,get_url_version(version),ckpt_name)
+            elif version == "local":
+                ckpt_path = self.cfg.ckpt_path
+
+            # no ckpt found on the device, try to download from the url
+            if os.path.isfile(ckpt_path) == False and version != "local":
+                os.makedirs(os.path.dirname(ckpt_path), exist_ok=True)
+                download_file(ckpt_path_url, ckpt_path)
                 time.sleep(0.5)
-            model = load_network(weight_path, model)
-        return model
+            elif os.path.isfile(ckpt_path) == False and version == "local":
+                raise RuntimeError(
+                    f"For the method {self.cfg.model_name}, no ckpt can be found on the {ckpt_path} !!! Try set the version to get the model from the url."
+                )
+            model = load_network(ckpt_path, model)
+        # to device
+        model = model.to(self.device)
+        model = nn.DataParallel(model) if self.cfg.multi_gpu == True else model
+        self.net = model
+        return self
 
     def save_network(self, save_path):
         """Save the network."""
@@ -81,27 +125,22 @@ class BaseModel(nn.Module):
             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}."
+        assert spike_length >= self.model_length, f"Spike input is not long enough, given {spike_length} frames < {self.cfg.model_length}."
         # even length
-        if self.cfg.model_win_length == self.model_half_win_length * 2:
+        if self.model_length == self.model_half_length * 2:
             spike = spike[
                 :,
-                spike_mid - self.model_half_win_length : spike_mid + self.model_half_win_length,
+                spike_mid - self.model_half_length : spike_mid + self.model_half_length,
             ]
         # odd length
         else:
             spike = spike[
                 :,
-                spike_mid - self.model_half_win_length : spike_mid + self.model_half_win_length + 1,
+                spike_mid - self.model_half_length : spike_mid + self.model_half_length + 1,
             ]
         self.spike_size = (spike.shape[2], spike.shape[3])
         return spike
@@ -116,60 +155,77 @@ class BaseModel(nn.Module):
         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"]
+        rate = batch["rate"].view(-1, 1, 1, 1).float()
         # outputs
         outputs = {}
-        recon_img = self(spike)
-        outputs["recon_img"] = recon_img
+        recon_img = self.spk2img(spike)
+        outputs["recon_img"] = recon_img / rate
         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"]
+        visual_dict["recon_img"] = outputs["recon_img"]
+        visual_dict["gt_img"] = batch["gt_img"]
         return visual_dict
 
-    def get_loss_dict(self, outputs, batch):
+    def get_loss_dict(self, outputs, batch, loss_weight_dict):
         """Get the loss dict from the given input batch and outputs."""
         # data process
-        gt_img = batch["img"]
+        gt_img = batch["gt_img"]
         # recon image
         recon_img = outputs["recon_img"]
         # loss dict
         loss_dict = {}
-        loss_dict["l1"] = self.criterion(recon_img, gt_img)
+        for loss_name, weight in loss_weight_dict.items():
+            loss_dict[loss_name] = weight * self.get_loss_func(loss_name)(recon_img, gt_img)
+
+        # todo add your desired loss here by loss_dict["name"] = loss()
+
         loss_values_dict = {k: v.item() for k, v in loss_dict.items()}
-        return loss_dict,loss_values_dict
+        return loss_dict, loss_values_dict
+
+    def get_loss_func(self, name: Literal["l1", "l2"]):
+        """Get the loss function from the given loss name."""
+        if name not in self.loss_func_cache:
+            if name == "l1":
+                self.loss_func_cache[name] = nn.L1Loss()
+            elif name == "l2":
+                self.loss_func_cache[name] = nn.MSELoss()
+            else:
+                self.loss_func_cache[name] = lambda x, y: 0
+        loss_func = self.loss_func_cache[name]
+        return loss_func
 
     def get_paired_imgs(self, batch, outputs):
+        """Get paired images for the metric calculation."""
         recon_img = outputs["recon_img"]
-        img = batch["img"]
+        img = batch["gt_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):
+        """Feed the batch data to the given device."""
         batch = {k: v.to(self.device, non_blocking=True) if torch.is_tensor(v) else v for k, v in batch.items()}
         return batch
+
+
+# functions
+def get_suffix(model_name, version):
+    if version == "v010":
+        if model_name in ["ssml", "wgse"]:
+            return "pt"
+        else:
+            return "pth"
+    else:
+        return "pth"
+
+
+def get_url_version(version):
+    major = version[1]
+    minor = version[2]
+    patch = version[3]
+    return f"v{major}.{minor}.{patch}"

spikezoo/models/bsf_model.py

@@ -1,6 +1,12 @@
 import torch
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from spikezoo.models.base_model import BaseModel, BaseModelConfig
+from torch.optim import Adam
+import torch.optim.lr_scheduler as lr_scheduler
+import torch.nn as nn
+from spikezoo.pipeline import TrainPipelineConfig
+from typing import List
+from spikezoo.archs.bsf.models.bsf.bsf import BSF 
 
 
 @dataclass
@@ -9,9 +15,11 @@ class BSFConfig(BaseModelConfig):
     model_name: str = "bsf"
     model_file_name: str = "models.bsf.bsf"
     model_cls_name: str = "BSF"
-    model_win_length: int = 61
+    model_length: int = 61
+    model_length_dict: dict = field(default_factory=lambda: {"v010": 61, "v023": 41})
     require_params: bool = True
-    ckpt_path: str = "weights/bsf.pth"
+    model_params: dict = field(default_factory=lambda: {})
+    model_params_dict: dict = field(default_factory=lambda: {"v010": {"spike_dim": 61}, "v023": {"spike_dim": 41}})
 
 
 class BSF(BaseModel):

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

@@ -2,7 +2,8 @@ from dataclasses import dataclass
 from spikezoo.models.base_model import BaseModel, BaseModelConfig
 import torch
 import torch.nn.functional as F
-
+from dataclasses import field
+from spikezoo.archs.spikeclip.nets import LRN 
 
 @dataclass
 class SpikeCLIPConfig(BaseModelConfig):
@@ -10,9 +11,9 @@ class SpikeCLIPConfig(BaseModelConfig):
     model_name: str = "spikeclip"
     model_file_name: str = "nets"
     model_cls_name: str = "LRN"
-    model_win_length: int = 200
+    model_length: int = 200
+    model_length_dict: dict = field(default_factory=lambda: {"v010": 200, "v023": 200})
     require_params: bool = True
-    ckpt_path: str = "weights/spikeclip.pth"
 
 
 class SpikeCLIP(BaseModel):

spikezoo/models/spk2imgnet_model.py

@@ -1,7 +1,12 @@
 import torch
 from dataclasses import dataclass, field
 from spikezoo.models.base_model import BaseModel, BaseModelConfig
-
+from spikezoo.pipeline import TrainPipelineConfig
+import torch.nn as nn
+import torch.optim as optim
+import torch.optim.lr_scheduler as lr_scheduler
+from typing import List
+from spikezoo.archs.spk2imgnet.nets import SpikeNet 
 
 @dataclass
 class Spk2ImgNetConfig(BaseModelConfig):
@@ -9,22 +14,11 @@ class Spk2ImgNetConfig(BaseModelConfig):
     model_name: str = "spk2imgnet"
     model_file_name: str = "nets"
     model_cls_name: str = "SpikeNet"
-    model_win_length: int = 41
+    model_length: int = 41
+    model_length_dict: dict = field(default_factory=lambda: {"v010": 41, "v023": 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):
@@ -45,7 +39,7 @@ class Spk2ImgNet(BaseModel):
             image = image[:, :, :250, :]
         elif self.spike_size == (480, 854):
             image = image[:, :, :, :854]
-        # used on the REDS_small dataset.
+        # used on the REDS_BASE dataset.
         if self.cfg.light_correction == True:
             image = torch.clamp(image / 0.6, 0, 1)
         return image

spikezoo/models/ssir_model.py

@@ -1,5 +1,7 @@
 from dataclasses import dataclass
 from spikezoo.models.base_model import BaseModel, BaseModelConfig
+from dataclasses import field
+from spikezoo.archs.ssir.models.networks import SSIR 
 
 
 @dataclass
@@ -8,15 +10,11 @@ class SSIRConfig(BaseModelConfig):
     model_name: str = "ssir"
     model_file_name: str = "models.networks"
     model_cls_name: str = "SSIR"
-    model_win_length: int = 41
+    model_length: int = 41
+    model_length_dict: dict = field(default_factory=lambda: {"v010": 41, "v023": 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

@@ -1,5 +1,8 @@
 from dataclasses import dataclass
 from spikezoo.models.base_model import BaseModel, BaseModelConfig
+import torch
+from dataclasses import field
+from spikezoo.archs.ssml.model import DoubleNet 
 
 
 @dataclass
@@ -8,11 +11,50 @@ class SSMLConfig(BaseModelConfig):
     model_name: str = "ssml"
     model_file_name: str = "model"
     model_cls_name: str = "DoubleNet"
-    model_win_length: int = 41
+    model_length: int = 41
+    model_length_dict: dict = field(default_factory=lambda: {"v010": 41, "v023": 41})
+    tfp_label_length: int = 11
     require_params: bool = True
-    ckpt_path: str = 'weights/ssml.pt'
 
 
+# ! A simple version of SSML rather than the full version
 class SSML(BaseModel):
     def __init__(self, cfg: BaseModelConfig):
         super(SSML, self).__init__(cfg)
+
+    def get_outputs_dict(self, batch):
+        # data process
+        spike = batch["spike"]
+        spike = self.preprocess_spike(spike)
+        rate = batch["rate"].view(-1, 1, 1, 1).float()
+        # outputs
+        outputs = {}
+        bsn_pred, nbsn_pred = self.net(spike)
+        bsn_pred = self.postprocess_img(bsn_pred)
+        nbsn_pred = self.postprocess_img(nbsn_pred)
+        outputs["recon_img"] = nbsn_pred / rate
+        outputs["bsn_pred"] = bsn_pred / rate
+        # tfp-label
+        mid = spike.shape[1] // 2
+        tfp_label = torch.mean(spike[:, mid - self.cfg.tfp_label_length // 2 : mid + self.cfg.tfp_label_length // 2 + 1], dim=1, keepdim=True)
+        outputs["tfp_label"] = self.postprocess_img(tfp_label) / rate
+        return outputs
+
+    def get_visual_dict(self, batch, outputs):
+        visual_dict = super().get_visual_dict(batch, outputs)
+        visual_dict["bsn_pred"] = outputs["bsn_pred"]
+        visual_dict["tfp_label"] = outputs["tfp_label"]
+        return visual_dict
+
+    def get_loss_dict(self, outputs, batch, loss_weight_dict):
+        # recon image
+        recon_img = outputs["recon_img"]
+        bsn_pred = outputs["bsn_pred"]
+        tfp_label = outputs["tfp_label"]
+        # loss dict
+        loss_dict = {}
+        for loss_name, weight in loss_weight_dict.items():
+            loss_dict["bsn_loss_" + loss_name] = weight * self.get_loss_func(loss_name)(bsn_pred, tfp_label)
+            loss_dict["mutual_loss_" + loss_name] = 0.01 * weight * self.get_loss_func(loss_name)(recon_img, bsn_pred)
+        loss_values_dict = {k: v.item() for k, v in loss_dict.items()}
+        return loss_dict, loss_values_dict

spikezoo/models/stir_model.py

@@ -1,7 +1,14 @@
 import torch
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from spikezoo.models.base_model import BaseModel, BaseModelConfig
-
+from torch.optim import Adam
+import torch.optim.lr_scheduler as lr_scheduler
+import torch.nn as nn
+from spikezoo.pipeline import TrainPipelineConfig
+from typing import List
+from spikezoo.archs.stir.metrics.losses import compute_per_loss_single
+from spikezoo.archs.stir.models.Vgg19 import Vgg19
+from spikezoo.archs.stir.models.networks_STIR import STIR
 
 @dataclass
 class STIRConfig(BaseModelConfig):
@@ -9,9 +16,11 @@ class STIRConfig(BaseModelConfig):
     model_name: str = "stir"
     model_file_name: str = "models.networks_STIR"
     model_cls_name: str = "STIR"
-    model_win_length: int = 61
+    model_length: int = 61
+    model_length_dict: dict = field(default_factory=lambda: {"v010": 61, "v023": 41})
     require_params: bool = True
-    ckpt_path: str = "weights/stir.pth"
+    model_params: dict = field(default_factory=lambda: {})
+    model_params_dict: dict = field(default_factory=lambda: {"v010": {"spike_dim": 61}, "v023": {"spike_dim": 41}})
 
 
 class STIR(BaseModel):
@@ -29,9 +38,21 @@ class STIR(BaseModel):
         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
+
+    def get_outputs_dict(self, batch):
+        # data process
+        spike = batch["spike"]
+        rate = batch["rate"].view(-1, 1, 1, 1).float()
+        # outputs
+        outputs = {}
+        spike = self.preprocess_spike(spike)
+        # pyramid loss is omitted owing to limited performance gain.
+        img_pred_0, Fs_lv_0, Fs_lv_1, Fs_lv_2, Fs_lv_3, Fs_lv_4, Est = self.net(spike)
+        img_pred_0 = self.postprocess_img(img_pred_0)
+        outputs["recon_img"] = img_pred_0 / rate
+        return outputs

spikezoo/models/tfi_model.py

@@ -1,5 +1,6 @@
 from dataclasses import dataclass, field
 from spikezoo.models.base_model import BaseModel, BaseModelConfig
+from spikezoo.archs.tfi.nets import TFIModel
 
 
 @dataclass
@@ -8,7 +9,8 @@ class TFIConfig(BaseModelConfig):
     model_name: str = "tfi"
     model_file_name: str = "nets"
     model_cls_name: str = "TFIModel"
-    model_win_length: int = 41
+    model_length: int = 41
+    model_length_dict: dict = field(default_factory=lambda: {"v010": 41, "v023": 41})
     require_params: bool = False
     model_params: dict = field(default_factory=lambda: {"model_win_length": 41})
 

spikezoo/models/tfp_model.py

@@ -1,5 +1,6 @@
-from dataclasses import dataclass,field
+from dataclasses import dataclass, field
 from spikezoo.models.base_model import BaseModel, BaseModelConfig
+from spikezoo.archs.tfp.nets import TFPModel
 
 
 @dataclass
@@ -8,7 +9,8 @@ class TFPConfig(BaseModelConfig):
     model_name: str = "tfp"
     model_file_name: str = "nets"
     model_cls_name: str = "TFPModel"
-    model_win_length: int = 41
+    model_length: int = 41
+    model_length_dict: dict = field(default_factory=lambda: {"v010": 41, "v023": 41})
     require_params: bool = False
     model_params: dict = field(default_factory=lambda: {"model_win_length": 41})