spikezoo 0.1.2__py3-none-any.whl → 0.2__py3-none-any.whl

spikezoo/archs/spk2imgnet/train.py

@@ -0,0 +1,189 @@
+import argparse
+import glob
+import os
+import re
+from collections import OrderedDict
+
+import torch.optim as optim
+from torch.autograd import Variable
+from torch.utils.data import DataLoader
+from dataset import *
+from nets import *
+from utils import *
+
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+
+parser = argparse.ArgumentParser(description="Spk2ImgNet")
+parser.add_argument(
+    "--preprocess", type=bool, default=False, help="run prepare_data or not"
+)
+parser.add_argument("--batchSize", type=int, default=16, help="Trainning batch size")
+parser.add_argument(
+    "--num_of_layers", type=int, default=17, help="Number of total layers"
+)
+parser.add_argument("--epochs", type=int, default=61, help="Number of trainning epochs")
+parser.add_argument(
+    "--milestone",
+    type=int,
+    default=20,
+    help="When to decay learning rate: should be less than epochs",
+)
+parser.add_argument(
+    "--lr",
+    type=float,
+    default=1e-4,
+    help="Initial learning rate; should be less than epochs",
+)
+parser.add_argument(
+    "--outf",
+    type=str,
+    default="./ckpt2",
+    help="path of log files",
+)
+parser.add_argument(
+    "--load_model", type=bool, default=False, help="load model from net.pth"
+)
+opt = parser.parse_args()
+
+if not os.path.exists(opt.outf):
+    os.mkdir(opt.outf)
+
+
+def find_last_checkpoint(save_dir):
+    file_list = glob.glob(os.path.join(save_dir, "model_*.pth"))
+    if file_list:
+        epoch_exist = []
+        for file_ in file_list:
+            result = re.findall(".*model_(.*).pth.*", file_)
+            epoch_exist.append(int(result[0]))
+        initial_epoch = max(epoch_exist)
+    else:
+        initial_epoch = 0
+    return initial_epoch
+
+
+def main():
+    # Load dataset
+    print("Loading dataset ...\n")
+    dataset_train = Dataset("train")
+    loader_train = DataLoader(
+        dataset=dataset_train, num_workers=4, batch_size=opt.batchSize, shuffle=True
+    )
+    print("# of training samples: %d\n" % int(len(dataset_train)))
+    '''
+    dataset_val = Dataset("val_stack")
+    loader_val = DataLoader(
+        dataset=dataset_val, num_workers=4, batch_size=opt.batchSize, shuffle=False
+    ) 
+    '''
+    # Build model
+    model = SpikeNet(in_channels=13, features=64, out_channels=1, win_r=6, win_step=7)
+    if not opt.load_model:
+        initial_epoch = 0
+        print("haha")
+    else:
+        # load model
+        initial_epoch = find_last_checkpoint(save_dir=opt.outf)
+        print("load model from model.pth")
+        state_dict = torch.load(
+            os.path.join(opt.outf, "model_%03d.pth" % initial_epoch)
+        )
+        new_state_dict = OrderedDict()
+        for k, v in state_dict.items():
+            name = k[7:]
+            new_state_dict[name] = v
+        model.load_state_dict(new_state_dict)
+    criterion = nn.L1Loss(size_average=True)
+    # Move to GPU
+    device_ids = [0]
+    model = nn.DataParallel(model).cuda()
+    criterion = criterion.cuda()
+    # Optimazer
+    optimizer = optim.Adam(model.parameters(), lr=opt.lr)
+    # training
+    model.train()
+    step = 0
+    for epoch in range(initial_epoch, opt.epochs):
+        avg_psnr = 0
+        if epoch < opt.milestone:
+            current_lr = opt.lr
+        else:
+            current_lr = opt.lr / 10.0
+        # set learning rate
+        for param_group in optimizer.param_groups:
+            param_group["lr"] = current_lr
+        print("learning rate %f" % current_lr)
+        # train
+        for i, (inputs, gt) in enumerate(loader_train, 0):
+            # print(inputs.shape)
+            inputs = Variable(inputs).cuda()
+            gt = Variable(gt).cuda()
+            # training step
+            model.train()
+            model.zero_grad()
+            optimizer.zero_grad()
+            rec, est0, est1, est2, est3, est4 = model(inputs)
+            est0 = est0 / 0.6
+            est1 = est1 / 0.6
+            est2 = est2 / 0.6
+            est3 = est3 / 0.6
+            est4 = est4 / 0.6
+            rec = rec / 0.6
+            loss = criterion(gt[:, 2:3, :, :], rec)
+            for slice_id in range(4):
+                loss = loss + 0.02 * (
+                    criterion(gt[:, 0:1, :, :], est0[:, slice_id : slice_id + 1, :, :])
+                    + criterion(
+                        gt[:, 1:2, :, :], est1[:, slice_id : slice_id + 1, :, :]
+                    )
+                    + criterion(
+                        gt[:, 2:3, :, :], est2[:, slice_id : slice_id + 1, :, :]
+                    )
+                    + criterion(
+                        gt[:, 3:4, :, :], est3[:, slice_id : slice_id + 1, :, :]
+                    )
+                    + criterion(
+                        gt[:, 4:5, :, :], est4[:, slice_id : slice_id + 1, :, :]
+                    )
+                )
+            loss.backward()
+            optimizer.step()
+            rec = torch.clamp(rec, 0, 1)
+            # print(rec)
+            psnr_train = batch_psnr(rec, gt[:, 2:3, :, :], 1.0)
+            # print(gt[:,2:3,:,:])
+            avg_psnr += psnr_train
+            if i % 10 == 0:
+                print(
+                    "[epoch %d][%d | %d] loss: %.4f PSNR_train: %.4f"
+                    % (epoch + 1, i + 1, len(loader_train), loss.item(), psnr_train)
+                )
+            step += 1
+        avg_psnr = avg_psnr / len(loader_train)
+        print("avg_psnr: %.2f" % avg_psnr)
+
+        if epoch % 5 == 0:
+            '''
+            # validate
+            model.eval()
+            psnr_val = 0
+            for i, (inputs, gt) in enumerate(loader_val, 0):
+                inputs = Variable(inputs).cuda()
+                gt = Variable(gt).cuda()
+                rec = model(inputs)
+                rec = rec / 0.6
+                rec = torch.clamp(rec, 0, 1)
+                psnr_val += batch_psnr(rec, gt, 1.0)
+            print("[epoch %d] PSNR_val: %.4f" % (epoch + 1, psnr_val / len(loader_val)))
+            '''
+            # save model
+            torch.save(
+                model.state_dict(),
+                os.path.join(opt.outf, "model_%03d.pth" % (epoch + 1)),
+            )
+
+
+if __name__ == "__main__":
+    if opt.preprocess:
+        prepare_data(data_path="./Spk2ImgNet_train/train2/", patch_size=40, stride=40, h5_name='train')
+    main()

spikezoo/archs/spk2imgnet/utils.py

@@ -0,0 +1,64 @@
+import math
+
+import numpy as np
+import torch.nn as nn
+from skimage.metrics import peak_signal_noise_ratio
+
+
+def weights_init_kaiming(m):
+    classname = m.__class__.__name__
+    if classname.find("Conv") != -1:
+        nn.init.kaiming_normal_(m.weight.data, a=0, mode="fan_in")
+    elif classname.find("Linear") != -1:
+        nn.init.kaiming_normal_(m.weight.data, a=0, mode="fan_in")
+    elif classname.find("BatchNorm") != -1:
+        # nn.init.uniform(m.weight.data, 1.0, 0.02)
+        m.weight.data.normal_(mean=0, std=math.sqrt(2.0 / 9.0 / 64.0)).clamp_(
+            -0.025, 0.025
+        )
+        nn.init.constant(m.bias.data, 0.0)
+
+
+def batch_psnr(img, imclean, data_range):
+    img = img.data.cpu().numpy().astype(np.float32)
+    imclean = imclean.data.cpu().numpy().astype(np.float32)
+    psnr = peak_signal_noise_ratio(img, imclean, data_range=data_range)
+    """
+    PSNR = 0
+    for i in range(Img.shape[0]):
+        PSNR += compare_psnr(imclean[i,:,:,:], img[i,:,:,:], data_range=data_range)
+    return (PSNR/Img.shape[0])
+    """
+    return psnr
+
+
+def data_augmentation(image, mode):
+    out = np.transpose(image, (1, 2, 0))
+    if mode == 0:
+        # original
+        out = out
+    elif mode == 1:
+        # flip up and down
+        out = np.flipud(out)
+    elif mode == 2:
+        # rotate counterwise 90 degree
+        out = np.rot90(out)
+    elif mode == 3:
+        # rotate 90 degree and flip up and down
+        out = np.rot90(out)
+        out = np.flipud(out)
+    elif mode == 4:
+        # rotate 180 degree
+        out = np.rot90(out, k=2)
+    elif mode == 5:
+        # rotate 180 degree and flip
+        out = np.rot90(out, k=2)
+        out = np.flipud(out)
+    elif mode == 6:
+        # rotate 270 degree
+        out = np.rot90(out, k=3)
+    elif mode == 7:
+        # rotate 270 degree and flip
+        out = np.rot90(out, k=3)
+        out = np.flipud(out)
+    return np.transpose(out, (2, 0, 1))

spikezoo/archs/ssir/README.md

@@ -0,0 +1,87 @@
+## [TCSVT 2023] Spike Camera Image Reconstruction Using Deep Spiking Neural Networks
+
+<h4 align="center"> Rui Zhao<sup>1</sup>, Ruiqin Xiong<sup>1</sup>, Jian Zhang<sup>2</sup>, Zhaofei Yu<sup>1</sup>, Shuyuan Zhu<sup>3</sup>, Lei Ma <sup>1</sup>, Tiejun Huang<sup>1</sup> </h4>
+<h4 align="center">1. National Engineering Research Center of Visual Technology, School of Computer Science, Peking University<br>
+2. School of Electronic and Computer Engineering, Peking University Shenzhen Graduate School<br>
+3.  School of Information and Communication Engineering, UESTC</h4><br>
+
+This repository contains the official source code for our paper:
+
+Spike Camera Image Reconstruction Using Deep Spiking Neural Networks
+
+TCSVT 2023
+
+[Paper](https://ieeexplore.ieee.org/document/10288531)
+
+
+
+## Environment
+
+You can choose cudatoolkit version to match your server. The code is tested on PyTorch 2.0.1+cuda12.0.
+
+```shell
+conda create -n ssir python==3.10
+conda activate ssir
+# You can choose the PyTorch version you like, we recommand version >= 1.10.1
+# For example
+pip install torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2
+pip install -r requirements.txt
+```
+
+## Prepare the Data
+
+#### 1. Download and deploy the SREDS dataset
+
+[BaiduNetDisk](https://pan.baidu.com/s/1clA43FcxjOibL1zGTaU82g) (Password: 2728)
+
+`train.tar` corresponds to the training data, and `test.tar` corresponds to the testing data.
+
+Move the above two `.tar` file to the `data root` directory and extract to the current directory
+
+```
+file directory:
+train:
+your_data_root/crop_mini/spike/...
+your_data_root/crop_mini/image/...
+test:
+your_data_root/spike/...
+your_data_root/imgs/...
+```
+
+#### 2. Set the path of RSSF dataset in your serve
+
+In the line25 of `main.py` or set that in command line when running main.py
+
+## Evaluate
+```shell
+cd shells
+bash eval_SREDS.sh
+```
+
+## Train
+```shell
+cd shells
+bash train_SSIR.sh
+```
+We recommended to redirect the output logs by adding
+`>> SSIR.txt 2>&1` 
+to the last of the above command for management.
+
+
+## Citation
+
+If you find this code useful in your research, please consider citing our paper.
+
+```
+@article{zhao2023spike,
+  title={Spike Camera Image Reconstruction Using Deep Spiking Neural Networks},
+  author={Zhao, Rui and Xiong, Ruiqin and Zhang, Jian and Yu, Zhaofei and Zhu, Shuyuan and Ma, Lei and Huang, Tiejun},
+  journal={IEEE Transactions on Circuits and Systems for Video Technology (TCSVT)},
+  year={2023},
+}
+```
+
+If you have any questions, please contact:  
+ruizhao@stu.pku.edu.cn
+
+ 

spikezoo/archs/ssir/configs/SSIR.yml

@@ -0,0 +1,37 @@
+data:
+  interp: 20
+  alpha: 0.4 
+
+seed: 6666
+
+loader:
+  # crop_size: [128, 128]
+  crop_size: [96, 96]
+  pair_step: 4
+
+model:
+  arch: 'sunet'
+  seq_len: 8
+  flow_weight_decay: 0.0004
+  flow_bias_decay: 0.0
+  #########################
+  kwargs:
+    activation_type: 'lif'
+    mp_activation_type: 'amp_lif'
+    spike_connection: 'concat'
+    num_encoders: 3
+    num_resblocks: 1
+    v_threshold: 1.0
+    v_reset: None
+    tau: 2.0
+
+
+train:
+  print_freq: 100
+  mixed_precision: True
+  vis_freq: 20
+
+optimizer:
+  solver: Adam
+  momentum: 0.9
+  beta: 0.999

spikezoo/archs/ssir/configs/yml_parser.py

@@ -0,0 +1,78 @@
+import numpy as np
+import torch
+import yaml
+
+
+class YAMLParser:
+    """ 
+        Modified from code from tudelft ssl-evflow
+    """
+
+    def __init__(self, config):
+        self.reset_config()
+        self.parse_config(config)
+        # self.init_seeds()
+
+    def parse_config(self, file):
+        with open(file) as fid:
+            yaml_config = yaml.load(fid, Loader=yaml.FullLoader)
+        self.parse_dict(yaml_config)
+
+    @property
+    def config(self):
+        return self._config
+
+    @property
+    def device(self):
+        return self._device
+
+    @property
+    def loader_kwargs(self):
+        return self._loader_kwargs
+
+    def reset_config(self):
+        self._config = {}
+
+    def update(self, config):
+        self.reset_config()
+        self.parse_config(config)
+
+    def parse_dict(self, input_dict, parent=None):
+        if parent is None:
+            parent = self._config
+        for key, val in input_dict.items():
+            if isinstance(val, dict):
+                if key not in parent.keys():
+                    parent[key] = {}
+                self.parse_dict(val, parent[key])
+            else:
+                parent[key] = val
+
+    @staticmethod
+    def worker_init_fn(worker_id):
+        np.random.seed(np.random.get_state()[1][0] + worker_id)
+
+    # def init_seeds(self):
+    #     torch.manual_seed(self._config["loader"]["seed"])
+    #     if torch.cuda.is_available():
+    #         torch.cuda.manual_seed(self._config["loader"]["seed"])
+    #         torch.cuda.manual_seed_all(self._config["loader"]["seed"])
+
+    def merge_configs(self, run):
+        """
+        Overwrites mlflow metadata with configs.
+        """
+
+        # parse mlflow settings
+        config = {}
+        for key in run.keys():
+            if len(run[key]) > 0 and run[key][0] == "{":  # assume dictionary
+                config[key] = eval(run[key])
+            else:  # string
+                config[key] = run[key]
+
+        # overwrite with config settings
+        self.parse_dict(self._config, config)
+        self.combine_entries(config)
+
+        return config

spikezoo/archs/ssir/datasets/dataset_sreds.py

@@ -0,0 +1,170 @@
+import os
+import os.path as osp
+import random
+import numpy as np
+import torch
+import torch.utils.data as data
+from datasets.ds_utils import *
+import time
+
+
+class Augmentor:
+    def __init__(self, crop_size):
+        # spatial augmentation params
+        self.crop_size = crop_size
+
+    def augment_img(self, img, mode=0):
+        '''Kai Zhang (github: https://github.com/cszn)
+        W x H x C or W x H
+        '''
+        if mode == 0:
+            return img
+        elif mode == 1:
+            return np.flipud(np.rot90(img))
+        elif mode == 2:
+            return np.flipud(img)
+        elif mode == 3:
+            return np.rot90(img, k=3)
+        elif mode == 4:
+            return np.flipud(np.rot90(img, k=2))
+        elif mode == 5:
+            return np.rot90(img)
+        elif mode == 6:
+            return np.rot90(img, k=2)
+        elif mode == 7:
+            return np.flipud(np.rot90(img, k=3))
+
+    def spatial_transform(self, spk_list, img_list):
+        mode = random.randint(0, 7)
+
+        for ii, spk in enumerate(spk_list):
+            spk = np.transpose(spk, [1,2,0])
+            spk = self.augment_img(spk, mode=mode)
+            spk_list[ii] = np.transpose(spk, [2,0,1])
+
+        for ii, img in enumerate(img_list):
+            img = np.transpose(img, [1,2,0])
+            img = self.augment_img(img, mode=mode)
+            img_list[ii] = np.transpose(img, [2,0,1])
+
+        return spk_list, img_list
+
+    def __call__(self, spk_list, img_list):
+        spk_list, img_list = self.spatial_transform(spk_list, img_list)
+        spk_list = [np.ascontiguousarray(spk) for spk in spk_list]
+        img_list = [np.ascontiguousarray(img) for img in img_list]
+        return spk_list, 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)
+        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

spikezoo/archs/ssir/datasets/ds_utils.py

@@ -0,0 +1,66 @@
+import numpy as np
+import os
+import cv2
+import os.path as osp
+
+def RawToSpike(video_seq, h, w, flipud=True):
+    video_seq = np.array(video_seq).astype(np.uint8)
+    img_size = h*w
+    img_num = len(video_seq)//(img_size//8)
+    SpikeMatrix = np.zeros([img_num, h, w], np.uint8)
+    pix_id = np.arange(0,h*w)
+    pix_id = np.reshape(pix_id, (h, w))
+    comparator = np.left_shift(1, np.mod(pix_id, 8))
+    byte_id = pix_id // 8
+
+    for img_id in np.arange(img_num):
+        id_start = img_id*img_size//8
+        id_end = id_start + img_size//8
+        cur_info = video_seq[id_start:id_end]
+        data = cur_info[byte_id]
+        result = np.bitwise_and(data, comparator)
+        if flipud:
+            SpikeMatrix[img_id, :, :] = np.flipud((result == comparator))
+        else:
+            SpikeMatrix[img_id, :, :] = (result == comparator)
+
+    return SpikeMatrix
+
+
+def SpikeToRaw(SpikeSeq, save_path):
+    """
+        SpikeSeq: Numpy array (sfn x h x w)
+        save_path: full saving path (string)
+        Rui Zhao
+    """
+    sfn, h, w = SpikeSeq.shape
+    base = np.power(2, np.linspace(0, 7, 8))
+    fid = open(save_path, 'ab')
+    for img_id in range(sfn):
+        # 模拟相机的倒像
+        spike = np.flipud(SpikeSeq[img_id, :, :])
+        # numpy按自动按行排，数据也是按行存的
+        spike = spike.flatten()
+        spike = spike.reshape([int(h*w/8), 8])
+        data = spike * base
+        data = np.sum(data, axis=1).astype(np.uint8)
+        fid.write(data.tobytes())
+
+    fid.close()
+
+    return
+
+
+def dat_to_spmat(dat_path,  size=[720, 1280]):
+    f = open(dat_path, 'rb')
+    video_seq = f.read()
+    video_seq = np.frombuffer(video_seq, 'b')
+    sp_mat = RawToSpike(video_seq, size[0], size[1])
+    return sp_mat
+
+
+def read_img_gray(file_path):
+    im = cv2.imread(file_path).astype(np.float32) / 255.0
+    im = cv2.cvtColor(im, cv2.COLOR_BGR2GRAY)
+    im = np.expand_dims(im, axis=0)
+    return im