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

spikezoo/archs/ssml/test.py

@@ -0,0 +1,67 @@
+from json import load
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from model import DoubleNet
+import cv2
+
+def load_vidar_dat(filename, left_up=(0, 0), window=None, frame_cnt = None, **kwargs):
+    if isinstance(filename, str):
+        array = np.fromfile(filename, dtype=np.uint8)
+    elif isinstance(filename, (list, tuple)):
+        l = []
+        for name in filename:
+            a = np.fromfile(name, dtype=np.uint8)
+            l.append(a)
+        array = np.concatenate(l)
+    else:
+        raise NotImplementedError
+    
+    height = 250
+    width = 400
+
+    if window == None:
+        window = (height - left_up[0], width - left_up[0])
+
+    len_per_frame = height * width // 8
+    framecnt = frame_cnt if frame_cnt != None else len(array) // len_per_frame
+
+    spikes = []
+
+    for i in range(framecnt):
+        compr_frame = array[i * len_per_frame: (i + 1) * len_per_frame]
+        blist = []
+        for b in range(8):
+            blist.append(np.right_shift(np.bitwise_and(compr_frame, np.left_shift(1, b)), b))
+        
+        frame_ = np.stack(blist).transpose()
+        frame_ = np.flipud(frame_.reshape((height, width), order='C'))
+
+        if window is not None:
+            spk = frame_[left_up[0]:left_up[0] + window[0], left_up[1]:left_up[1] + window[1]]
+        else:
+            spk = frame_
+
+        spk = torch.from_numpy(spk.copy().astype(np.float32)).unsqueeze(dim=0)
+
+        spikes.append(spk)
+
+    return torch.cat(spikes)
+
+if __name__ == '__main__':
+    model = DoubleNet()
+    model_path = "./fin3g-best-lucky.pt"
+    model = nn.DataParallel(model)
+    model.load_state_dict(torch.load(model_path)) 
+    model = model.cuda()
+    
+    spike_path = "./rotation1.dat"
+    spike = load_vidar_dat(spike_path)[200:200+41].unsqueeze(0).cuda()
+
+    res = model(spike)
+    res = res[0].detach().cpu().permute(1,2,0).numpy()*255
+    res_path = "./res.png"
+    cv2.imwrite(res_path,res)
+
+    print("done.")

spikezoo/archs/stir/README.md

@@ -0,0 +1,65 @@
+<!---
+# Spatio-Temporal Interactive Learning for Efficient Image Reconstruction of Spiking Cameras
+
+This repository contains the source code for the paper: [Spatio-Temporal Interactive Learning for Efficient Image Reconstruction of Spiking Cameras (NeurIPS 2024)](https://openreview.net/pdf?id=S4ZqnMywcM). 
+The spiking camera is an emerging neuromorphic vision sensor that records high-speed motion scenes by asynchronously firing continuous binary spike streams. Prevailing image reconstruction methods, generating intermediate frames from these spike streams, often rely on complex step-by-step network architectures that overlook the intrinsic collaboration of spatio-temporal complementary information. In this paper, we propose an efficient spatio-temporal interactive reconstruction network to jointly perform inter-frame feature alignment and intra-frame feature filtering in a coarse-to-fine manner. Specifically, it starts by extracting hierarchical features from a concise hybrid spike representation, then refines the motion fields and target frames scale-by-scale, ultimately obtaining a full-resolution output. Meanwhile, we introduce a symmetric interactive attention block and a multi-motion field estimation block to further enhance the interaction capability of the overall network. Experiments on synthetic and real-captured data show that our approach exhibits excellent performance while maintaining low model complexity.
+
+<img src="picture/performance-speed.png" width="75%"/>
+<img src="picture/overview.png" width="80%"/>
+<img src="picture/results_visual.png" width="82%"/>
+-->
+## Installation
+You can choose cudatoolkit version to match your server. The code is tested with PyTorch 1.9.1 with CUDA 11.4.
+
+```shell
+conda create -n stir python==3.8.12
+conda activate stir
+# You can choose the PyTorch version you like, for example
+pip install torch==2.2.2 torchvision==0.17.2 torchaudio==2.0.2
+```
+
+Install the dependent packages:
+```
+pip install -r requirements.txt
+```
+
+Install core package
+```
+cd ./package_core
+python setup.py install
+```
+
+In our implementation, we borrowed the code framework of [SSIR](https://github.com/ruizhao26/SSIR):
+
+## Prepare the Data
+
+#### 1. Download and deploy the SREDS dataset to your local computer from [SSIR](https://github.com/ruizhao26/SSIR).
+
+#### 2. Set the path of the SREDS dataset in your serve
+
+Set that in `--data_root` when running train_STIR.sh or eval_SREDS.sh
+
+## Evaluate
+```
+sh eval_SREDS.sh
+```
+
+## Train
+```
+sh train_STIR.sh
+```
+<!---
+## Citations
+If you find our approach useful in your research or applications, please consider giving us a star 🌟 and citing it by the following BibTeX entry.
+```
+@article{fan2024spatio,
+	title={Spatio-Temporal Interactive Learning for Efficient Image Reconstruction of Spiking Cameras},
+	author={Fan, Bin and Yin, Jiaoyang and Dai, Yuchao and Xu, Chao and Huang, Tiejun and Shi, Boxin},
+	journal={Proceedings of the Advances in Neural Information Processing Systems (NeurIPS)},
+        volume={},
+	year={2024}
+}
+```
+-->
+## Statement
+This project is for research purpose only, please contact us for the licence of commercial use. For any other questions or discussion please contact: binfan@mail.nwpu.edu.cn

spikezoo/archs/stir/ckpt_outputs/Descriptions.txt

@@ -0,0 +1 @@
+This folder is used to store the trained model.

spikezoo/archs/stir/configs/STIR.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: 'STIR'
+  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: 1
+  mixed_precision: True
+  vis_freq: 20
+
+optimizer:
+  solver: Adam
+  momentum: 0.9
+  beta: 0.999

spikezoo/archs/stir/configs/utils.py

@@ -0,0 +1,155 @@
+import numpy as np
+import torch
+import torch.nn.functional as F
+import os
+import os.path as osp
+import random
+import cv2
+
+def set_seeds(_seed_):
+    random.seed(_seed_)
+    np.random.seed(_seed_)
+    torch.manual_seed(_seed_)  # use torch.manual_seed() to seed the RNG for all devices (both CPU and CUDA)
+    torch.cuda.manual_seed_all(_seed_)
+
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+    # os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8'
+    # set a debug environment variable CUBLAS_WORKSPACE_CONFIG to ":16:8" (may limit overall performance) or ":4096:8" (will increase library footprint in GPU memory by approximately 24MiB).
+    # torch.use_deterministic_algorithms(True)
+
+
+def make_dir(path):
+    if not osp.exists(path):
+        os.makedirs(path)
+    return
+
+
+def add_args_to_cfg(cfg, args, args_list):
+    for aa in args_list:
+        cfg['train'][aa] = eval('args.{:s}'.format(aa))
+    return cfg
+
+
+# class AverageMeter(object):
+#     """Computes and stores the average and current value"""
+#     def __init__(self, precision=3):
+#         self.precision = precision
+#         self.reset()
+
+#     def reset(self):
+#         self.val = 0
+#         self.avg = 0
+#         self.sum = 0
+#         self.count = 0
+
+#     def update(self, val, n=1):
+#         self.val = val
+#         self.sum += val * n
+#         self.count += n
+#         self.avg = self.sum / self.count
+
+#     def __repr__(self):
+#         return '{:.{}f} ({:.{}f})'.format(self.val, self.precision, self.avg, self.precision)
+
+
+class AverageMeter(object):
+    """Computes and stores the average and current value"""
+
+    def __init__(self, i=1, precision=3, names=None):
+        self.meters = i
+        self.precision = precision
+        self.reset(self.meters)
+        self.names = names
+        if names is not None:
+            assert self.meters == len(self.names)
+        else:
+            self.names = [''] * self.meters
+
+    def reset(self, i):
+        self.val = [0] * i
+        self.avg = [0] * i
+        self.sum = [0] * i
+        self.count = [0] * i
+
+    def update(self, val, n=1):
+        if not isinstance(val, list):
+            val = [val]
+        if not isinstance(n, list):
+            n = [n] * self.meters
+        assert (len(val) == self.meters and len(n) == self.meters)
+        for i in range(self.meters):
+            self.count[i] += n[i]
+        for i, v in enumerate(val):
+            self.val[i] = v
+            self.sum[i] += v * n[i]
+            self.avg[i] = self.sum[i] / self.count[i]
+
+    def __repr__(self):
+        # val = ' '.join(['{} {:.{}f}'.format(n, v, self.precision) for n, v in
+        #                 zip(self.names, self.val)])
+        # avg = ' '.join(['{} {:.{}f}'.format(n, a, self.precision) for n, a in
+        #                 zip(self.names, self.avg)])
+        out = ' '.join(['{} {:.{}f} ({:.{}f})'.format(n, v, self.precision, a, self.precision) for n, v, a in
+                        zip(self.names, self.val, self.avg)])
+        # return '{} ({})'.format(val, avg)
+        return '{}'.format(out)
+
+
+def normalize_image_torch(image, percentile_lower=1, percentile_upper=99):
+    b, c, h, w = image.shape
+    image_reshape = image.reshape([b, c, h*w])
+    mini = torch.quantile(image_reshape, 0.01, dim=2, keepdim=True).unsqueeze_(dim=3)
+    maxi = torch.quantile(image_reshape, 0.99, dim=2, keepdim=True).unsqueeze_(dim=3)
+    # if mini == maxi:
+    #     return 0 * image + 0.5  # gray image
+    return torch.clip((image - mini) / (maxi - mini + 1e-5), 0, 1)
+
+def normalize_image_torch2(image):
+    return torch.clip(image, 0, 1)
+
+# --------------------------------------------
+# Torch to Numpy 0~255
+# --------------------------------------------
+def torch2numpy255(im):
+    im = im[0, 0].detach().cpu().numpy()
+    im = (im * 255).astype(np.float64)
+    return im
+
+def torch2torch255(im):
+    return im * 255.0
+
+class InputPadder:
+    """ Pads images such that dimensions are divisible by padsize """
+    def __init__(self, dims, padsize=16):
+        self.ht, self.wd = dims[-2:]
+        pad_ht = (((self.ht // padsize) + 1) * padsize - self.ht) % padsize
+        pad_wd = (((self.wd // padsize) + 1) * padsize - self.wd) % padsize
+        #self._pad = [pad_wd//2, pad_wd - pad_wd//2, 0, pad_ht]
+        self._pad = [pad_wd//2, pad_wd - pad_wd//2, pad_ht//2, pad_ht - pad_ht//2]
+
+    def pad(self, *inputs):
+        return [F.pad(x, self._pad, mode='replicate') for x in inputs]
+
+    def unpad(self,x):
+        ht, wd = x.shape[-2:]
+        c = [self._pad[2], ht-self._pad[3], self._pad[0], wd-self._pad[1]]
+        return x[..., c[0]:c[1], c[2]:c[3]]
+    
+
+
+def vis_img(vis_path: str, img: torch.Tensor, vis_name: str = 'vis'):
+    ww = 0
+    rows = []
+    for ii in range(4):
+        cur_row = []
+        for jj in range(img.shape[0]//4):
+            cur_img = img[ww, 0].detach().cpu().numpy() * 255
+            cur_img = cur_img.astype(np.uint8)
+            cur_row.append(cur_img)
+            ww += 1
+        cur_row_cat = np.concatenate(cur_row, axis=1)
+        rows.append(cur_row_cat)
+    out_img = np.concatenate(rows, axis=0)
+    cv2.imwrite(osp.join(vis_path, vis_name+'.png'), out_img)
+    return

spikezoo/archs/stir/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/stir/datasets/dataset_sreds.py

@@ -0,0 +1,180 @@
+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)
+            '''
+            # 按照文件名称读取
+            spikes_path_list = [osp.join(spike_dir, spk_path_list[ii]) for ii in range(spklen)]
+            images_path_list = [osp.join(image_dir, spk_path_list[ii][:-4]+'.png') for ii in range(spklen)]
+
+            if(self.confirm_exist([spikes_path_list, images_path_list])):
+                s = {}
+                s['spikes_paths'] = spikes_path_list
+                s['images_paths'] = images_path_list
+                samples.append(s)
+                
+        return samples
+
+    def _load_sample(self, s):
+        data = {}
+
+        data['spikes'] = [np.array(dat_to_spmat(p, size=(96, 96)), dtype=np.float32) for p in s['spikes_paths']]
+        data['images'] = [read_img_gray(p) for p in s['images_paths']]
+
+        data['spikes'], data['images'] = self.augmentor(data['spikes'], data['images'])
+        # print("data['spikes'][0].shape, data['images'][0].shape", data['spikes'][0].shape, data['images'][0].shape)
+
+        return data
+    
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, index):
+        data = self._load_sample(self.samples[index])
+        return data
+
+
+class sreds_test(torch.utils.data.Dataset):
+    def __init__(self, cfg):
+        self.cfg = cfg
+        self.samples = self.collect_samples()
+        print('The samples num of testing data: {:d}'.format(len(self.samples)))
+    
+    def confirm_exist(self, path_list_list):
+        for pl in path_list_list:
+            for p in pl:
+                if not osp.exists(p):
+                    return 0
+        return 1
+    
+    def collect_samples(self):
+        spike_path = osp.join(self.cfg['data']['root'], 'spike', 'val', 'interp_{:d}_alpha_{:.2f}'.format(self.cfg['data']['interp'], self.cfg['data']['alpha']))
+        image_path = osp.join(self.cfg['data']['root'], 'imgs', 'val', 'val_orig')
+        scene_list = sorted(os.listdir(spike_path))
+        samples = []
+
+        for scene in scene_list:
+            spike_dir = osp.join(spike_path, scene)
+            image_dir = osp.join(image_path, scene)
+            spk_path_list = sorted(os.listdir(spike_dir))
+            
+            spklen = len(spk_path_list)
+            # seq_len = self.cfg['model']['seq_len']
+
+            # 按照文件名称读取
+            spikes_path_list = [osp.join(spike_dir, spk_path_list[ii]) for ii in range(spklen)]
+            images_path_list = [osp.join(image_dir, spk_path_list[ii][:-4]+'.png') for ii in range(spklen)]
+
+            if(self.confirm_exist([spikes_path_list, images_path_list])):
+                s = {}
+                s['spikes_paths'] = spikes_path_list
+                s['images_paths'] = images_path_list
+                samples.append(s)
+
+        return samples
+
+    def _load_sample(self, s):
+        data = {}
+        data['spikes'] = [np.array(dat_to_spmat(p, size=(720, 1280)), dtype=np.float32) for p in s['spikes_paths']]
+        data['images'] = [read_img_gray(p) for p in s['images_paths']]
+        return data
+    
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, index):
+        data = self._load_sample(self.samples[index])
+        return data

spikezoo/archs/stir/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

spikezoo/archs/stir/eval_SREDS.sh

@@ -0,0 +1,5 @@
+python3 main.py \
+--data_root /data/local_userdata/fanbin/REDS_dataset/REDS120fps \
+--arch STIR \
+--pretrained ./ckpt/STIR_pretrain.pth \
+--eval