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

spikezoo/archs/bsf/models/bsf/bsf.py

@@ -0,0 +1,105 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .rep import MODF
+from .align import Multi_Granularity_Align
+
+
+class BasicModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+    
+    ####################################################################################
+    ## Tools functions for neural networks
+    def weight_parameters(self):
+        return [param for name, param in self.named_parameters() if 'weight' in name]
+
+    def bias_parameters(self):
+        return [param for name, param in self.named_parameters() if 'bias' in name]
+
+    def num_parameters(self):
+        return sum([p.data.nelement() if p.requires_grad else 0 for p in self.parameters()])
+    
+    def init_weights(self):
+        for layer in self.named_modules():
+            if isinstance(layer, nn.Conv2d):
+                nn.init.kaiming_normal_(layer.weight)
+                if layer.bias is not None:
+                    nn.init.constant_(layer.bias, 0)
+
+            elif isinstance(layer, nn.ConvTranspose2d):
+                nn.init.kaiming_normal_(layer.weight)
+                if layer.bias is not None:
+                    nn.init.constant_(layer.bias, 0)
+
+
+def split_and_b_cat(x):
+    x0 = x[:, 10-10:10+10+1].clone()
+    x1 = x[:, 20-10:20+10+1].clone()
+    x2 = x[:, 30-10:30+10+1].clone()
+    x3 = x[:, 40-10:40+10+1].clone()
+    x4 = x[:, 50-10:50+10+1].clone()
+    return torch.cat([x0, x1, x2, x3, x4], dim=0)
+
+
+class Encoder(nn.Module):
+    def __init__(self, base_dim=64, layers=4, act=nn.ReLU()):
+        super().__init__()
+        self.conv_list = nn.ModuleList()
+        for ii in range(layers):
+            self.conv_list.append(
+                nn.Sequential(
+                    nn.Conv2d(base_dim, base_dim, kernel_size=3, padding=1),
+                    act,
+                    nn.Conv2d(base_dim, base_dim, kernel_size=3, padding=1),
+                )
+            )
+        self.act = act
+
+    def forward(self, x):
+        for conv in self.conv_list:
+            x = self.act(conv(x) + x)
+        return x
+
+##########################################################################
+class BSF(BasicModel):
+    def __init__(self, act=nn.ReLU()):
+        super().__init__()
+        self.offset_groups = 4
+        self.corr_max_disp = 3
+
+        self.rep = MODF(base_dim=64, act=act)
+
+        self.encoder = Encoder(base_dim=64, layers=4, act=act)
+
+        self.align = Multi_Granularity_Align(base_dim=64, groups=self.offset_groups, act=act, sc=3)
+
+        self.recons = nn.Sequential(
+            nn.Conv2d(64*5, 64*3, kernel_size=3, padding=1),
+            act,
+            nn.Conv2d(64*3, 64, kernel_size=3, padding=1),
+            act,
+            nn.Conv2d(64, 1, kernel_size=3, padding=1),
+        )
+
+    def forward(self, input_dict):
+        dsft_dict = input_dict['dsft_dict']
+        dsft11 = dsft_dict['dsft11']
+        dsft12 = dsft_dict['dsft12']
+        dsft21 = dsft_dict['dsft21']
+        dsft22 = dsft_dict['dsft22']
+
+        dsft_b_cat = {
+            'dsft11': split_and_b_cat(dsft11),
+            'dsft12': split_and_b_cat(dsft12),
+            'dsft21': split_and_b_cat(dsft21),
+            'dsft22': split_and_b_cat(dsft22),
+        }
+
+        feat_b_cat = self.rep(dsft_b_cat)
+        feat_b_cat = self.encoder(feat_b_cat)
+        feat_list = feat_b_cat.chunk(5, dim=0)
+        feat_list_align = self.align(feat_list=feat_list)
+        out = self.recons(torch.cat(feat_list_align, dim=1))
+
+        return out

spikezoo/archs/bsf/models/bsf/dsft_convert.py

@@ -0,0 +1,96 @@
+import torch
+
+
+def convert_dsft4(dsft, spike):
+    '''
+    input: Pytorch Tensor
+        dsft: dsft(1,1)  b x T x h x w
+        spike: 01 spike  b x T x h x w
+    output: Pytorch Tensor
+        dsft_dict: {dsft(1,1), dsft(1,2), dsft(2,1), dsft(2,2)}
+    '''
+
+    b, T, h, w = spike.shape
+
+    ## dsft_mask_left_shift  -- abbr. -->  dmls1,  (right-shift: dmrs1)
+    dmls1 = -1 * torch.ones(spike.shape, device=spike.device, dtype=torch.float32)
+    dmrs1 = -1 * torch.ones(spike.shape, device=spike.device, dtype=torch.float32)
+
+    ## for dmls1
+    # flag的用途是为了边界的copy-padding
+    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:
+            ## dmls1的数据该更新的情况
+            update_coord = (spike[:,ii+1]==1) * (~copy_pad_coord)
+            dmls1[:,ii][update_coord] = dsft[:,ii+1][update_coord]
+
+            ## dmls1的数据不该更新，该继承之前的数的情况
+            non_update_coord = (spike[:,ii+1]!=1) * (~copy_pad_coord)
+            dmls1[:,ii][non_update_coord] = dmls1[:, ii+1][non_update_coord]
+    
+    
+    ## for dmrs1
+    # flag的用途是为了边界的copy-padding
+    flag = -2 * torch.ones([b, h, w], device=spike.device, dtype=torch.float32)
+    for ii in range(0, T, 1):
+        flag += (spike[:,ii]==1)
+
+        ## for 边界的 copy-padding
+        copy_pad_coord = (flag < 0)
+        dmrs1[:,ii][copy_pad_coord] = dsft[:,ii][copy_pad_coord]
+
+        if ii > 0:
+            ## dmrs1的数据该更新的情况
+            update_coord = (spike[:,ii]==1) * (~copy_pad_coord)
+            dmrs1[:,ii][update_coord] = dsft[:,ii-1][update_coord]
+
+            ## dmrs1的数据不该更新，该继承之前的数的情况
+            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
+
+
+
+if __name__ == '__main__':
+    # spike = [0,0,1,0,1,0,0,0,0,1,0,0,1,0,1,0,0,0,1,0,0]
+    # dsft  = [2,2,2,2,2,5,5,5,5,5,3,3,3,2,2,4,4,4,4,4,4]
+
+    spike = [0,0,1,0,0,1,0,0,0,1,0,0,0,1,0,0,1,0,1,0,0,1,0,0,0,1,0,0,1,0,0,1,0]
+    dsft  = [3,3,3,3,3,4,4,4,4,4,4,4,4,3,3,3,2,2,3,3,3,4,4,4,4,3,3,3,3,3,3,3,3]
+
+    spike = torch.tensor(spike, device='cpu', dtype=torch.float32)[None,:,None,None]
+    dsft  = torch.tensor(dsft , device='cpu', dtype=torch.float32)[None,:,None,None]
+
+    dsft_dict = convert_dsft4(dsft=dsft, spike=spike)
+    dsft_11 = dsft_dict['dsft11']
+    dsft_12 = dsft_dict['dsft12']
+    dsft_21 = dsft_dict['dsft21']
+    dsft_22 = dsft_dict['dsft22']
+
+    print(dsft_11[0,:,0,0])
+    print()
+    print(dsft_12[0,:,0,0])
+    print()
+    print(dsft_21[0,:,0,0])
+    print()
+    print(dsft_22[0,:,0,0])

spikezoo/archs/bsf/models/bsf/rep.py

@@ -0,0 +1,44 @@
+import torch
+import torch.nn as nn
+
+class MODF(nn.Module):
+    def __init__(self, base_dim=64, act=nn.ReLU()):
+        super().__init__()
+        self.base_dim = base_dim
+
+        self.conv1 = self._make_layer(input_dim=21, hidden_dim=self.base_dim, output_dim=self.base_dim, act=act)
+        self.conv_for_others = nn.ModuleList([
+            self._make_layer(input_dim=self.base_dim, hidden_dim=self.base_dim, output_dim=self.base_dim, act=act) for ii in range(3)
+        ])
+        self.conv_fuse = self._make_layer(input_dim=self.base_dim*3, hidden_dim=self.base_dim, output_dim=self.base_dim, act=act)
+
+    def _make_layer(self, input_dim, hidden_dim, output_dim, act):
+        layer = nn.Sequential(
+            nn.Conv2d(input_dim, hidden_dim, kernel_size=3, padding=1),
+            act,
+            nn.Conv2d(hidden_dim, output_dim, kernel_size=3, padding=1),
+        )
+        return layer
+
+    def forward(self, dsft_dict):
+        d11 = 1.0 / dsft_dict['dsft11']
+        d12 = 2.0 / dsft_dict['dsft12']
+        d21 = 2.0 / dsft_dict['dsft21']
+        d22 = 3.0 / dsft_dict['dsft22']
+
+        d_list = [d11, d12, d21, d22]
+        feat_batch_cat = self.conv1(torch.cat(d_list, dim=0))
+        feat_list = feat_batch_cat.chunk(4, dim=0)
+
+        feat_11 = feat_list[0]
+        feat_others_list = feat_list[1:]
+        feat_others_list_processed = []
+        for ii in range(3):
+            feat_others_list_processed.append(self.conv_for_others[ii](feat_others_list[ii]))
+
+
+        other_feat = torch.cat(feat_others_list_processed, dim=1)
+        other_feat_res = self.conv_fuse(other_feat)
+
+        return feat_11 + other_feat_res
+

spikezoo/archs/bsf/models/get_model.py

@@ -0,0 +1,7 @@
+from .bsf.bsf import BSF
+
+def get_model(args):
+    if args.arch.upper() == 'BSF'.upper():
+        model = BSF()
+
+    return model

spikezoo/archs/bsf/prepare_data/DSFT.py

@@ -0,0 +1,62 @@
+import numpy as np
+import torch
+
+class DSFT:
+    def __init__(self, spike_h, spike_w, device):
+        self.spike_h = spike_h
+        self.spike_w = spike_w
+        self.device = device
+
+
+    def spikes2images(self, spikes, max_search_half_window=20):
+        '''
+        将spikes整体转换为一段DSFT
+        
+        输入：
+        spikes: T x H x W 的numpy张量, 类型: 整型与浮点皆可
+        max_search_half_window: 对于要转换为图像的时刻点而言, 左右各参考的最大脉冲帧数量，超过这个数字就不搜了
+
+        输出：
+        ImageMatrix: T' x H x W 的numpy张量, 其中T' = T - (2 x max_search_half_window)
+        类型: uint8, 取值范围: 0 ~ 255
+        '''
+
+        T = spikes.shape[0]
+        T_im = T - 2*max_search_half_window
+
+        if T_im < 0:
+            raise ValueError('The length of spike stream {:d} is not enough for max_search half window length {:d}'.format(T, max_search_half_window))
+        
+        spikes = torch.from_numpy(spikes).to(self.device).float()
+        ImageMatrix = torch.zeros([T_im, self.spike_h, self.spike_w]).to(self.device)
+
+        pre_idx = -1 * torch.ones([T, self.spike_h, self.spike_w]).float().to(self.device)
+        cur_idx = -1 * torch.ones([T, self.spike_h, self.spike_w]).float().to(self.device)
+        
+        for ii in range(T):
+            if ii > 0:
+                pre_idx[ii] = cur_idx[ii-1]
+                cur_idx[ii] = cur_idx[ii-1]
+            cur_spk = spikes[ii]
+            cur_idx[ii][cur_spk==1] = ii
+
+        diff = cur_idx - pre_idx
+
+
+        interval = -1 * torch.ones([T, self.spike_h, self.spike_w]).float().to(self.device)
+        for ii in range(T-1, 0-1, -1):
+            interval[ii][diff[ii]!=0] = diff[ii][diff[ii]!=0]
+            if ii < T-1:
+                interval[ii][diff[ii]==0] = interval[ii+1][diff[ii]==0]
+        
+        # boundary
+        interval[interval==-1] = 255
+        interval[pre_idx==-1] = 255
+        
+        # for uint8
+        interval = torch.clip(interval, 0, 255)
+        
+        ImageMatrix = interval[max_search_half_window:-max_search_half_window].cpu().detach().numpy().astype(np.uint8)
+        
+
+        return ImageMatrix

spikezoo/archs/bsf/prepare_data/crop_dataset_train.py

@@ -0,0 +1,135 @@
+import os
+import os.path as osp
+import argparse
+import cv2
+import numpy as np
+from io_utils import *
+import h5py
+from tqdm import *
+from DSFT import DSFT
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--root", type=str, default="/data/rzhao/REDS120fps")
+parser.add_argument("--output_path", type=str, default="/data/rzhao/REDS120fps/crop")
+###### 参数
+parser.add_argument("--eta", type=float, default=1.0)
+parser.add_argument("--gamma", type=int, default=60)
+parser.add_argument("--alpha", type=float, default=0.7)
+
+parser.add_argument("--cu", '-c', type=str, default='0')
+
+parser.add_argument("--crop_image", action='store_true')
+args = parser.parse_args()
+
+os.environ['CUDA_VISIBLE_DEVICES'] = args.cu
+
+
+if __name__ == '__main__':
+    imgs_path = osp.join(args.root, 'imgs', 'train')
+    spks_path = osp.join(args.root, 'spikes', 'train', 
+                        "eta_{:.2f}_gamma_{:d}_alpha_{:.1f}".format(args.eta, args.gamma, args.alpha))
+
+    scene_list = sorted(os.listdir(spks_path))    
+    for scene in tqdm(scene_list):
+        scene_imgs_path = osp.join(imgs_path, scene)
+        scene_spks_path = osp.join(spks_path, scene)
+
+
+        if not args.crop_image:
+            # read all the dat files
+            dat_path = sorted(os.listdir(scene_spks_path))
+            spks_list = []
+            #### abandon 00000000.dat, corresponding to the following spike_idx_offset
+            for dat_name in dat_path[1:]:
+                spks_list.append(dat_to_spmat(dat_path=osp.join(scene_spks_path, dat_name), size=(720, 1280)))
+            spikes = np.concatenate(spks_list, axis=0)
+
+            # spikes -> DSFT(max_search_half_win=80)
+            dsft_solver = DSFT(spike_h=720, spike_w=1280, device='cuda')
+            dsft = dsft_solver.spikes2images(spikes, max_search_half_window=100)
+
+
+        # crop Image
+        if args.crop_image:
+            imgs_list = []
+            for im_idx in range(11, 28+1):
+                    img = cv2.imread(osp.join(scene_imgs_path, '{:08d}.png'.format(im_idx)))
+                    # 1. central crop
+                    crop_img = img[32:-32, 128:-128]
+                    # 2. crop
+                    for ii in range(3):
+                        for jj in range(4):
+                            if (ii != 2) and (jj != 3):
+                                cur_img = crop_img[256*ii:256*(ii+1), 256*jj:256*(jj+1)]
+                            elif ii != 2:
+                                cur_img = crop_img[256*ii:256*(ii+1), -256:]
+                            elif jj != 3:
+                                cur_img = crop_img[-256:, 256*jj:256*(jj+1)]
+                            else:
+                                cur_img = crop_img[-256:, -256:]
+                            cur_save_root = osp.join(args.output_path, 'train', 'imgs', scene, '{:02}'.format(ii*4+jj))
+                            os.makedirs(cur_save_root, exist_ok=True)
+                            cur_save_path = osp.join(cur_save_root, '{:08d}.png'.format(im_idx))
+                            if osp.exists(cur_save_path):
+                                os.remove(cur_save_path)
+                            cv2.imwrite(cur_save_path, cur_img)
+            continue
+
+        # crop spikes
+        # since 00000000.dat is abandoned
+        spike_idx_offset = 10
+        # 1. central crop
+        spikes = spikes[:, 32:-32, 128:-128]
+        # 2. crop
+        for spk_idx in range(11, 28+1):
+            crop_spike = spikes[spk_idx*10-spike_idx_offset : spk_idx*10-spike_idx_offset+10]
+            for ii in range(3):
+                for jj in range(4):
+                    if (ii != 2) and (jj != 3):
+                        cur_spk = crop_spike[:, 256*ii:256*(ii+1), 256*jj:256*(jj+1)]
+                    elif ii != 2:
+                        cur_spk = crop_spike[:, 256*ii:256*(ii+1), -256:]
+                    elif jj != 3:
+                        cur_spk = crop_spike[:, -256:, 256*jj:256*(jj+1)]
+                    else:
+                        cur_spk = crop_spike[:, -256:, -256:]
+
+                    cur_save_root = osp.join(args.output_path, 'train', 
+                                        "eta_{:.2f}_gamma_{:d}_alpha_{:.1f}".format(args.eta, args.gamma, args.alpha),
+                                        scene, '{:02}'.format(ii*4+jj), 'spikes')
+
+                    os.makedirs(cur_save_root, exist_ok=True)
+                    cur_save_path = osp.join(cur_save_root,'{:08d}.dat'.format(spk_idx))
+                    if osp.exists(cur_save_path):
+                        os.remove(cur_save_path)
+                    SpikeToRaw(SpikeSeq=cur_spk, save_path=cur_save_path)
+
+
+        # crop dsft
+        dsft_idx_offset = 10 + 100
+        # 1. central crop
+        dsft = dsft[:, 32:-32, 128:-128]
+        # 2. crop
+        for dsft_idx in range(11, 28+1):
+            crop_dsft = dsft[dsft_idx*10-dsft_idx_offset : dsft_idx*10-dsft_idx_offset+10]
+            for ii in range(3):
+                for jj in range(4):
+                    if (ii != 2) and (jj != 3):
+                        cur_dsft = crop_dsft[:, 256*ii:256*(ii+1), 256*jj:256*(jj+1)]
+                    elif ii != 2:
+                        cur_dsft = crop_dsft[:, 256*ii:256*(ii+1), -256:]
+                    elif jj != 3:
+                        cur_dsft = crop_dsft[:, -256:, 256*jj:256*(jj+1)]
+                    else:
+                        cur_dsft = crop_dsft[:, -256:, -256:]
+
+                    cur_save_root = osp.join(args.output_path, 'train', 
+                                        "eta_{:.2f}_gamma_{:d}_alpha_{:.1f}".format(args.eta, args.gamma, args.alpha),
+                                        scene, '{:02}'.format(ii*4+jj), 'dsft')
+                    os.makedirs(cur_save_root, exist_ok=True)
+                    cur_save_path = osp.join(cur_save_root, '{:08d}.h5'.format(dsft_idx))
+                    if osp.exists(cur_save_path):
+                        os.remove(cur_save_path)
+                    f = h5py.File(cur_save_path, 'w')
+                    f['dsft'] = cur_dsft
+                    f.close()

spikezoo/archs/bsf/prepare_data/crop_dataset_val.py

@@ -0,0 +1,139 @@
+import os
+import os.path as osp
+import argparse
+import cv2
+import numpy as np
+from io_utils import *
+import h5py
+from tqdm import *
+from DSFT import DSFT
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--root", type=str, default="/data/rzhao/REDS120fps")
+parser.add_argument("--output_path", type=str, default="/data/rzhao/REDS120fps/crop")
+###### 参数
+parser.add_argument("--eta", type=float, default=1.00)
+parser.add_argument("--gamma", type=int, default=60)
+parser.add_argument("--alpha", type=float, default=0.7)
+
+parser.add_argument("--cu", '-c', type=str, default='0')
+
+parser.add_argument("--crop_image", action='store_true')
+args = parser.parse_args()
+
+os.environ['CUDA_VISIBLE_DEVICES'] = args.cu
+
+if __name__ == '__main__':
+    imgs_path = osp.join(args.root, 'imgs', 'val')
+    spks_path = osp.join(args.root, 'spikes', 'val', 
+                        "eta_{:.2f}_gamma_{:d}_alpha_{:.1f}".format(args.eta, args.gamma, args.alpha))
+
+    scene_list = sorted(os.listdir(spks_path))    
+    for scene in tqdm(scene_list):
+        scene_imgs_path = osp.join(imgs_path, scene)
+        scene_spks_path = osp.join(spks_path, scene)
+
+        if not args.crop_image:
+            # read all the dat files
+            dat_path = sorted(os.listdir(scene_spks_path))
+            spks_list = []
+            #### abandon 00000000.dat, corresponding to the following spike_idx_offset
+            for dat_name in dat_path[1:]:
+                spks_list.append(dat_to_spmat(dat_path=osp.join(scene_spks_path, dat_name), size=(720, 1280)))
+            spikes = np.concatenate(spks_list, axis=0)
+
+            # spikes -> DSFT(max_search_half_win=80)
+            dsft_solver = DSFT(spike_h=720, spike_w=1280, device='cuda')
+            dsft = dsft_solver.spikes2images(spikes, max_search_half_window=100)
+
+
+        # crop Image
+        if args.crop_image:
+            imgs_list = []
+            for im_idx in range(11, 28+1):
+                    img = cv2.imread(osp.join(scene_imgs_path, '{:08d}.png'.format(im_idx)))
+                    # 1. central crop
+                    crop_img = img[32:-32, 128:-128]
+                    for sub_scene_idx in range(4):
+                        cur_scene = '{:s}_{:d}'.format(scene, sub_scene_idx)
+                        if sub_scene_idx == 0:
+                            cur_crop_img = crop_img[:384, :512]
+                        elif sub_scene_idx == 1:
+                            cur_crop_img = crop_img[-384:, :512]
+                        elif sub_scene_idx == 2:
+                            cur_crop_img = crop_img[:384, -512:]
+                        elif sub_scene_idx == 3:
+                            cur_crop_img = crop_img[-384:, -512:]
+
+                        cur_save_root = osp.join(args.output_path, 'val_small', 'imgs', cur_scene)
+                        os.makedirs(cur_save_root, exist_ok=True)
+                        cur_save_path = osp.join(cur_save_root, '{:08d}.png'.format(im_idx))
+                        if osp.exists(cur_save_path):
+                            os.remove(cur_save_path)
+                        cv2.imwrite(cur_save_path, cur_crop_img)
+            continue
+
+
+        # 裁切 spikes
+        # since 00000000.dat is abandoned
+        spike_idx_offset = 10
+        # 1. central crop
+        spikes = spikes[:, 32:-32, 128:-128]
+        # 2. crop
+        for spk_idx in range(11, 28+1):
+            crop_spike = spikes[spk_idx*10-spike_idx_offset : spk_idx*10-spike_idx_offset+10]
+            
+            for sub_scene_idx in range(4):
+                cur_scene = '{:s}_{:d}'.format(scene, sub_scene_idx)
+                if sub_scene_idx == 0:
+                    cur_crop_spike = crop_spike[:, :384, :512]
+                elif sub_scene_idx == 1:
+                    cur_crop_spike = crop_spike[:, -384:, :512]
+                elif sub_scene_idx == 2:
+                    cur_crop_spike = crop_spike[:, :384, -512:]
+                elif sub_scene_idx == 3:
+                    cur_crop_spike = crop_spike[:, -384:, -512:]
+
+                cur_save_root = osp.join(args.output_path, 'val_small',
+                                "eta_{:.2f}_gamma_{:d}_alpha_{:.1f}".format(args.eta, args.gamma, args.alpha), 
+                                cur_scene,
+                                'spikes')
+
+                os.makedirs(cur_save_root, exist_ok=True)
+                cur_save_path = osp.join(cur_save_root,'{:08d}.dat'.format(spk_idx))
+                if osp.exists(cur_save_path):
+                    os.remove(cur_save_path)
+                SpikeToRaw(SpikeSeq=cur_crop_spike, save_path=cur_save_path)
+
+
+        # crop dsft
+        dsft_idx_offset = 10 + 100
+        # 1. central crop
+        dsft = dsft[:, 32:-32, 128:-128]
+        # 2. crop
+        for dsft_idx in range(11, 28+1):
+            crop_dsft = dsft[dsft_idx*10-dsft_idx_offset : dsft_idx*10-dsft_idx_offset+10]
+
+            for sub_scene_idx in range(4):
+                cur_scene = '{:s}_{:d}'.format(scene, sub_scene_idx)
+                if sub_scene_idx == 0:
+                    cur_crop_dsft = crop_dsft[:, :384, :512]
+                elif sub_scene_idx == 1:
+                    cur_crop_dsft = crop_dsft[:, -384:, :512]
+                elif sub_scene_idx == 2:
+                    cur_crop_dsft = crop_dsft[:, :384, -512:]
+                elif sub_scene_idx == 3:
+                    cur_crop_dsft = crop_dsft[:, -384:, -512:]
+
+                cur_save_root = osp.join(args.output_path, 'val_small',
+                                "eta_{:.2f}_gamma_{:d}_alpha_{:.1f}".format(args.eta, args.gamma, args.alpha), 
+                                cur_scene,
+                                'dsft')
+                os.makedirs(cur_save_root, exist_ok=True)
+                cur_save_path = osp.join(cur_save_root, '{:08d}.h5'.format(dsft_idx))
+                if osp.exists(cur_save_path):
+                    os.remove(cur_save_path)
+                f = h5py.File(cur_save_path, 'w')
+                f['dsft'] = cur_crop_dsft
+                f.close()
+

spikezoo/archs/bsf/prepare_data/crop_train.sh

@@ -0,0 +1,4 @@
+python3 crop_dataset_train.py -c $1 --eta 1.00 &&
+python3 crop_dataset_train.py -c $1 --eta 0.75 &&
+python3 crop_dataset_train.py -c $1 --eta 0.50 &&
+python3 crop_dataset_train.py --crop_image

spikezoo/archs/bsf/prepare_data/crop_val.sh

@@ -0,0 +1,4 @@
+python3 crop_dataset_val.py -c $1 --eta 1.00 &&
+python3 crop_dataset_val.py -c $1 --eta 0.75 &&
+python3 crop_dataset_val.py -c $1 --eta 0.50 &&
+python3 crop_dataset_val.py --crop_image

spikezoo/archs/bsf/prepare_data/io_utils.py

@@ -0,0 +1,64 @@
+import numpy as np
+import os
+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 dat_to_spmat(dat_path,  size):
+    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
+
+
+## Save Raw dat files
+def SpikeToRaw(SpikeSeq, save_path):
+    """
+        SpikeSeq: Numpy array (sfn x h x w)
+        save_path: full saving path (string)
+    """
+    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 save_to_h5(SpikeMatrix, h5path, name):
+    f = h5py.File(h5path, 'w')
+    f[name] = SpikeMatrix
+    f.close()

spikezoo/archs/bsf/requirements.txt

@@ -0,0 +1,9 @@
+torch
+torchvision
+opencv-python
+numpy
+tensorboardX
+scikit-image
+lpips
+tqdm
+h5py