Upload 28 files

.gitattributes

@@ -33,3 +33,15 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+snnTracker/0.dat filter=lfs diff=lfs merge=lfs -text
+snnTracker/1.dat filter=lfs diff=lfs merge=lfs -text
+snnTracker/2.dat filter=lfs diff=lfs merge=lfs -text
+snnTracker/3.dat filter=lfs diff=lfs merge=lfs -text
+snnTracker/4.dat filter=lfs diff=lfs merge=lfs -text
+snnTracker/5.dat filter=lfs diff=lfs merge=lfs -text
+snnTracker/6.dat filter=lfs diff=lfs merge=lfs -text
+snnTracker/7.dat filter=lfs diff=lfs merge=lfs -text
+snnTracker/8.dat filter=lfs diff=lfs merge=lfs -text
+snnTracker/9.dat filter=lfs diff=lfs merge=lfs -text
+snnTracker/driving_0_snntracker.avi filter=lfs diff=lfs merge=lfs -text
+snnTracker/driving_0_tfi.avi filter=lfs diff=lfs merge=lfs -text

snnTracker/0.dat

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1d325c5ecb53e64a6429c753df5e0c3b46deaf3bce97f445499c9dfb64b3db4d
+size 5000000

snnTracker/1.dat

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ba866d8427b0baa97a0c09bab6055edceaac2dfff6fcba0580f21973628e5d1a
+size 5000000

snnTracker/2.dat

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8a00b47dc8e9e84f902f159cd1f0ab387440091b470e4da9503e18e3bae4d354
+size 5000000

snnTracker/3.dat

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:47a782e5eea659f9ddd8ad0e20b688e1e69ed9bf589f99989586397d9e2bd825
+size 5000000

snnTracker/4.dat

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:02b795f9ca1f3affc59d59e280051bfd50f24c3d9c1c214ab8a680c9fdc2175a
+size 5000000

snnTracker/5.dat

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:68ebdd547fc9c36401a3033701f725d9552de1ae235ea49f9bf04a1afad4a331
+size 5000000

snnTracker/6.dat

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e5d0349fab1c1da0d3b3559c8bf30795f2488ec7caa182126525c5eb49f52d10
+size 5000000

snnTracker/7.dat

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a36642ec9cc8ef292cb1223d8cf4c263dde4203da83f0e55179e3570a46754b3
+size 5000000

snnTracker/8.dat

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cea9cb59395d668eab1c18b28cca3821e63f6ed2f097e5a120e67253ee78080a
+size 5000000

snnTracker/9.dat

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:50ef6ee3758465fba69b5c453e76cae4a7c805805e646445e8877cc886849ce6
+size 5000000

snnTracker/datasets/0/config.yaml

@@ -0,0 +1,13 @@
+
+# 基本配置
+spike_h: 250  # 脉冲数据高度
+spike_w: 400  # 脉冲数据宽度
+is_labeled: false  # 是否有标注数据
+
+# 数据标识符
+data_field_identifier: ''  # 数据文件标识符
+label_field_identifier: ''  # 标注文件标识符
+
+# 标注数据配置
+labeled_data_type: 'tracking'  # 标注数据类型
+labeled_data_suffix: 'txt'  # 标注文件后缀

snnTracker/driving_0_snntracker.avi

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f738eb42502b754cadf26539f749fe0281ba023cb139f32dff7af2a0a0cadf99
+size 578258

snnTracker/driving_0_tfi.avi

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:62d6381f1f002c2f68b0e935d69db5ca2545ceb68566c6e6706f613cba6d2af4
+size 9748482

snnTracker/path.py

@@ -0,0 +1,52 @@
+# -*- coding: utf-8 -*- 
+# @Time : 2023/7/16 20:19 
+# @Author : Yajing Zheng
+# @Email: [email protected]
+# @File : path.py
+# here put the import lib
+import os
+
+
+def seek_file(search_dirs, filename):
+    search_dir_split = split_path_into_pieces(search_dirs)
+    dir_level = len(search_dir_split)
+
+    for i_dir in range(0, dir_level):
+        if i_dir > 0:
+            search_dir_split.pop(-1)
+        # search_dir = os.path.join(str(search_dir_split[0:-i_dir]))
+        search_dir = os.path.join(*search_dir_split)
+        for root, dirs, files in os.walk(search_dir):
+            if filename in files:
+                print('{0}/{1}'.format(root, filename))
+                filepath = os.path.join(root, filename)
+                return filepath
+
+
+def split_path_into_pieces(path: str):
+    pieces = []
+    if path[-1] == '/':
+        path = path[0:-1]
+
+    while True:
+        splits = os.path.split(path)
+        if splits[0] == '':
+            pieces.insert(0, splits[-1])
+            break
+        if splits[-1] == '':
+            pieces.insert(0, splits[0])
+            break
+        pieces.insert(0, splits[-1])
+        path = splits[0]
+
+    return pieces
+
+def replace_identifier(path: list, src: str, dst: str):
+    new_path = []
+    for piece in path:
+        added_piece = piece
+        if piece == src:
+            added_piece = dst
+        new_path.append(added_piece)
+
+    return new_path

snnTracker/spkData/load_dat.py

@@ -0,0 +1,203 @@
+# -*- coding: utf-8 -*- 
+# @Time : 2023/7/16 20:13 
+# @Author : Yajing Zheng
+# @Email: [email protected]
+# @File : load_dat.py
+import os, sys
+import warnings
+import glob
+import yaml
+import numpy as np
+import path
+
+# key-value for generate data loader according to the type of label data
+LABEL_DATA_TYPE = {
+    'raw': 0,
+    'reconstruction': 1,
+    'optical_flow': 2,
+    'mono_depth_estimation': 3.1,
+    'stero_depth_estimation': 3.2,
+    'detection': 4,
+    'tracking': 5,
+    'recognition': 6
+}
+
+
+# generate parameters dictionary according to labeled or not
+def data_parameter_dict(data_filename, label_type):
+    filename = path.split_path_into_pieces(data_filename)
+
+    if os.path.isabs(data_filename):
+        file_root = data_filename
+        if os.path.isdir(file_root):
+            search_root = file_root
+        else:
+            search_root = '\\'.join(filename[0:-1])
+        config_filename = path.seek_file(search_root, 'config.yaml')
+    else:
+        file_root = os.path.join('', 'datasets', *filename)
+        config_filename = os.path.join('', 'datasets', filename[0], 'config.yaml')
+
+    try:
+        with open(config_filename, 'r', encoding='utf-8') as fin:
+            configs = yaml.load(fin, Loader=yaml.FullLoader)
+    except TypeError as err:
+        print("Cannot find config file" + str(err))
+        raise err
+
+    except KeyError as exception:
+        print('ERROR! Task name does not exist')
+        print('Task name must be in %s' % LABEL_DATA_TYPE.keys())
+        raise exception
+
+    is_labeled = configs.get('is_labeled')
+
+    paraDict = {'spike_h': configs.get('spike_h'), 'spike_w': configs.get('spike_w')}
+    paraDict['filelist'] = None
+
+    if is_labeled:
+        paraDict['labeled_data_type'] = configs.get('labeled_data_type')
+        paraDict['labeled_data_suffix'] = configs.get('labeled_data_suffix')
+        paraDict['label_root_list'] = None
+
+        if os.path.isdir(file_root):
+            filelist = sorted(glob.glob(file_root + '/*.dat'), key=os.path.getmtime)
+            filepath = filelist[0]
+
+            labelname = path.replace_identifier(filename, configs.get('data_field_identifier', ''),
+                                                configs.get('label_field_identifier', ''))
+            label_root_list = os.path.join('', 'datasets', *labelname)
+            paraDict['labeled_data_dir'] = sorted(glob.glob(label_root_list + '/*.' + paraDict['labeled_data_suffix']),
+                                                  key=os.path.getmtime)
+
+            paraDict['filelist'] = filelist
+            paraDict['label_root_list'] = label_root_list
+        else:
+            filepath = glob.glob(file_root)[0]
+            rawname = filename[-1].replace('.dat', '')
+            filename.pop(-1)
+            filename.append(rawname)
+            labelname = path.replace_identifier(filename, configs.get('data_field_identifier', ''),
+                                                configs.get('label_field_identifier', ''))
+            label_root = os.path.join('', 'datasets', *labelname)
+            paraDict['labeled_data_dir'] = glob.glob(label_root + '.' + paraDict['labeled_data_suffix'])[0]
+    else:
+        filepath = file_root
+
+    paraDict['filepath'] = filepath
+
+    return paraDict
+
+
+class SpikeStream:
+    def __init__(self, **kwargs):
+
+        self.SpikeMatrix = None
+        self.filename = kwargs.get('filepath')
+        if os.path.splitext(self.filename)[-1][1:] != 'dat':
+            self.filename = self.filename + '.dat'
+        self.spike_w = kwargs.get('spike_w')
+        self.spike_h = kwargs.get('spike_h')
+        if 'print_dat_detail' not in kwargs:
+            self.print_dat_detail = True
+        else:
+            self.print_dat_detail = kwargs.get('print_dat_detail')
+
+    def get_spike_matrix(self, flipud=True, with_head=False):
+
+        file_reader = open(self.filename, 'rb')
+        video_seq = file_reader.read()
+        video_seq = np.frombuffer(video_seq, 'b')
+
+        video_seq = np.array(video_seq).astype(np.byte)
+        if self.print_dat_detail:
+            print(video_seq)
+        if with_head:
+            decode_width = 416
+        else:
+            decode_width = self.spike_w
+        # img_size = self.spike_height * self.spike_width
+        img_size = self.spike_h * decode_width
+        img_num = len(video_seq) // (img_size // 8)
+
+        if self.print_dat_detail:
+            print('loading total spikes from dat file -- spatial resolution: %d x %d, total timestamp: %d' %
+                  (decode_width, self.spike_h, img_num))
+
+        # SpikeMatrix = np.zeros([img_num, self.spike_h, self.spike_width], np.byte)
+
+        pix_id = np.arange(0, img_num * self.spike_h * decode_width)
+        pix_id = np.reshape(pix_id, (img_num, self.spike_h, decode_width))
+        comparator = np.left_shift(1, np.mod(pix_id, 8))
+        byte_id = pix_id // 8
+
+        data = video_seq[byte_id]
+        result = np.bitwise_and(data, comparator)
+        tmp_matrix = (result == comparator)
+
+        # if with head, delete them
+        if with_head:
+            delete_indx = np.arange(400, 416)
+            tmp_matrix = np.delete(tmp_matrix, delete_indx, 2)
+
+        if flipud:
+            self.SpikeMatrix = tmp_matrix[:, ::-1, :]
+        else:
+            self.SpikeMatrix = tmp_matrix
+
+        file_reader.close()
+        self.SpikeMatrix = self.SpikeMatrix.astype(np.byte)
+        return self.SpikeMatrix
+
+        # return spikes with specified length and begin index
+
+    def get_block_spikes(self, begin_idx, block_len=1, flipud=True, with_head=False):
+
+        file_reader = open(self.filename, 'rb')
+        video_seq = file_reader.read()
+        video_seq = np.frombuffer(video_seq, 'b')
+
+        video_seq = np.array(video_seq).astype(np.uint8)
+
+        if with_head:
+            decode_width = 416
+        else:
+            decode_width = self.spike_w
+        # img_size = self.spike_height * self.spike_width
+        img_size = self.spike_h * decode_width
+        img_num = len(video_seq) // (img_size // 8)
+
+        end_idx = begin_idx + block_len
+        if end_idx > img_num:
+            warnings.warn("block_len exceeding upper limit! Zeros will be padded in the end. ", ResourceWarning)
+            end_idx = img_num
+
+        if self.print_dat_detail:
+            print(
+                'loading total spikes from dat file -- spatial resolution: %d x %d, begin index: %d total timestamp: %d' %
+                (decode_width, self.spike_h, begin_idx, block_len))
+
+        pix_id = np.arange(0, block_len * self.spike_h * decode_width)
+        pix_id = np.reshape(pix_id, (block_len, self.spike_h, decode_width))
+        comparator = np.left_shift(1, np.mod(pix_id, 8))
+        byte_id = pix_id // 8
+        id_start = begin_idx * img_size // 8
+        id_end = id_start + block_len * img_size // 8
+        data = video_seq[id_start:id_end]
+        data_frame = data[byte_id]
+        result = np.bitwise_and(data_frame, comparator)
+        tmp_matrix = (result == comparator)
+
+        # if with head, delete them
+        if with_head:
+            delete_indx = np.arange(400, 416)
+            tmp_matrix = np.delete(tmp_matrix, delete_indx, 2)
+
+        if flipud:
+            self.SpikeMatrix = tmp_matrix[:, ::-1, :]
+        else:
+            self.SpikeMatrix = tmp_matrix
+
+        file_reader.close()
+        self.SpikeMatrix = self.SpikeMatrix.astype(np.byte)
+        return self.SpikeMatrix

snnTracker/spkProc/detection/attention_select.py

@@ -0,0 +1,132 @@
+import skimage.morphology as smor
+from skimage.feature import peak_local_max
+from skimage.morphology import erosion
+from skimage.measure import label, regionprops
+import numpy as np
+import torch
+from torchvision.transforms import Resize
+
+
+# obtain 2D gaussian filter
+def get_kernel(filter_size, sigma):
+
+    assert (filter_size + 1) % 2 == 0, '2D filter size must be odd number!'
+    g = np.zeros((filter_size, filter_size), dtype=np.float32)
+    half_width = int((filter_size - 1) / 2)
+    # center location
+
+    xc = (filter_size + 1) / 2
+    yc = (filter_size + 1) / 2
+    for i in range(-half_width, half_width + 1, 1):
+        for j in range(-half_width, half_width + 1, 1):
+            x = int(xc + i)
+            y = int(yc + j)
+            g[y - 1, x - 1] = np.exp(- (i ** 2 + j ** 2) / 2 / sigma / sigma)
+
+    g = (g - g.min()) / (g.max() - g.min())
+    return g
+
+
+# detect moving connected regions
+class SaccadeInput:
+
+    def __init__(self, spike_h, spike_w, box_size, device, attentionThr=None, extend_edge=None):
+
+        self.spike_h = spike_h
+        self.spike_w = spike_w
+        self.device = device
+
+        self.U = torch.zeros(self.spike_h, self.spike_w, dtype=torch.float32)
+        self.tau_u = 0.5
+        self.global_inih = 0.01
+        self.box_width = box_size  # attention box width
+        self.Jxx_size = self.box_width * 2 + 1
+        self.Jxx = torch.nn.Conv2d(in_channels=1, out_channels=1, kernel_size=(self.Jxx_size, self.Jxx_size),
+                                   padding=(self.box_width, self.box_width), bias=False)
+
+        tmp_filter = get_kernel(self.Jxx_size, round(self.box_width / 2) + 1)
+        tmp_filter = tmp_filter.reshape((1, 1, self.Jxx_size, self.Jxx_size))
+        self.Jxx.weight.data = torch.from_numpy(tmp_filter)
+        self.resizer = Resize((self.Jxx_size, self.Jxx_size))
+
+        self.U = self.U.to(self.device)
+        self.Jxx = self.Jxx.to(self.device)
+
+        if attentionThr is not None:
+            self.attentionThr = attentionThr
+        else:
+            self.attentionThr = 40
+        if extend_edge is not None:
+            self.extend_edge = extend_edge
+        else:
+            self.extend_edge = 7
+            # self.extend_edge = 1
+        self.peak_width = int(self.extend_edge)
+
+    def update_dnf(self, spike):
+        inputSpk = torch.reshape(spike, (1, 1, self.spike_h, self.spike_w)).float()
+
+        maxU = torch.relu(self.U)
+        squareU = torch.square(maxU)
+        r = squareU / (1 + self.global_inih * torch.sum(squareU))
+        conv_fired = self.Jxx(inputSpk)
+        conv_fired = torch.squeeze(conv_fired).to(self.device)
+        du = conv_fired - self.U
+
+        r = torch.reshape(r, (1, 1, self.spike_h, self.spike_w))
+        conv_r = self.Jxx(r)
+        conv_r = torch.squeeze(conv_r).to(self.device)
+        du = conv_r + du
+        self.U += (du * self.tau_u).detach()
+
+        del inputSpk, maxU, squareU, r, conv_r, conv_fired, du
+
+    def get_attention_location(self, spikes):
+
+        tmpU = torch.relu(self.U - self.attentionThr)
+        tmpU = tmpU.cpu()
+        tmpU = tmpU.detach().numpy()
+        dilated_u = erosion(tmpU, smor.square(self.peak_width))
+        peak_cord = peak_local_max(dilated_u, min_distance=self.box_width)
+        num_max = len(peak_cord)
+        # print('detect %d attention location' % num_max)
+        dilated_u[dilated_u > 1] = 1
+        dilated_u[dilated_u < 1] = 0
+        region_labels = label(dilated_u)
+        regions = regionprops(region_labels)
+        num_box = len(regions)
+
+        attentionBox = torch.zeros((num_box, 4), dtype=torch.int)
+        attentionInput = torch.zeros(self.Jxx_size + 4, self.Jxx_size, num_box)
+
+        for region, iBox in zip(regions, range(num_box)):
+            minr, minc, maxr, maxc = region.bbox
+            beginX = minr - self.extend_edge >= 0 and minr - self.extend_edge or 0
+            beginY = minc - self.extend_edge >= 0 and minc - self.extend_edge or 0
+            endX = maxr + self.extend_edge < self.spike_h and maxr + self.extend_edge or self.spike_h - 1
+            endY = maxc + self.extend_edge < self.spike_w and maxc + self.extend_edge or self.spike_w - 1
+
+            attentionBox[iBox, :] = torch.tensor([beginX, beginY, endX, endY])
+            attentionI = torch.unsqueeze(spikes[beginX:endX + 1, beginY:endY + 1], dim=0)
+            attentionI = self.resizer.forward(attentionI)
+            fire_index = torch.where(attentionI > 0.9)
+            attentionI2 = torch.zeros_like(attentionI)
+            attentionI2[0, fire_index[1], fire_index[2]] = 1
+            attentionInput[:-4, :, iBox] = torch.squeeze(attentionI2).detach().clone()
+            tmp_spk = bin(beginX + 1)
+            tmp_spk = tmp_spk[2:].zfill(self.box_width)
+            attentionInput[-4, :-1, iBox] = torch.from_numpy(np.tile(np.array(list(tmp_spk), dtype=np.float32), (1, 2)))
+            tmp_spk = bin(beginY + 1)
+            tmp_spk = tmp_spk[2:].zfill(self.box_width)
+            attentionInput[-3, :-1, iBox] = torch.from_numpy(np.tile(np.array(list(tmp_spk), dtype=np.float32), (1, 2)))
+            tmp_spk = bin(endX + 1)
+            tmp_spk = tmp_spk[2:].zfill(self.box_width)
+            attentionInput[-2, :-1, iBox] = torch.from_numpy(np.tile(np.array(list(tmp_spk), dtype=np.float32), (1, 2)))
+            tmp_spk = bin(endY + 1)
+            tmp_spk = tmp_spk[2:].zfill(self.box_width)
+            attentionInput[-1, :-1, iBox] = torch.from_numpy(np.tile(np.array(list(tmp_spk), dtype=np.float32), (1, 2)))
+
+        attentionInput = attentionInput.to(self.device)
+        del tmpU, dilated_u, peak_cord
+
+        return attentionBox, attentionInput

snnTracker/spkProc/detection/motion_clustering.py

@@ -0,0 +1,101 @@
+# -*- coding: utf-8 -*- 
+# @Time : 2024/12/6 3:34 
+# @Author : Yajing Zheng
+# @Email: [email protected]
+# @File : motion_clustering.py
+from sklearn.cluster import DBSCAN, OPTICS, SpectralClustering
+from sklearn.metrics import pairwise_distances
+from sklearn.preprocessing import StandardScaler
+import numpy as np
+import scipy.ndimage.measurements as mnts
+import torch
+
+
+class detect_object:
+
+    def __init__(self, h, w):
+        self.h = h
+        self.w = w
+        params = {'quantile': .3,
+                  'eps': .4,
+                  'damping': .9,
+                  'preference': -200,
+                  'n_neighbors': 10,
+                  'min_samples': 50,
+                  'xi': 0.05,
+                  'min_cluster_size': 0.1,
+                  'n_cluster': 2}
+        self.dbscan = DBSCAN(eps=params['eps'], min_samples=params['min_samples'], metric='precomputed')
+        self.optics = OPTICS(min_samples=params['min_samples'], xi=params['xi'],
+                             min_cluster_size=params['min_cluster_size'], metric='precomputed')
+        self.spectral = SpectralClustering(n_clusters=params['n_cluster'], eigen_solver='arpack',
+                                           affinity='precomputed')
+
+    def get_object(self, motion_vector, max_motion=None):
+        # motion_vector = StandardScaler.transform(motion_vector)
+        if max_motion is None:
+            mv_idx = torch.where(torch.logical_or(motion_vector[:, :, 0] != 0, motion_vector[:, :, 1] != 0))
+            if len(mv_idx[0]) < 1:
+                return None, None
+            fire_idx = np.zeros((2, len(mv_idx[0])), dtype=np.int)
+            fire_idx[0, :] = mv_idx[0].cpu().numpy()
+            fire_idx[1, :] = mv_idx[1].cpu().numpy()
+
+        else:
+            max_motion = max_motion.cpu().numpy()
+            if max_motion.max() < 1:
+                return None, None
+            fire_idx = np.array(np.nonzero(max_motion))
+
+        motion_vector = motion_vector.cpu().numpy()
+
+        fire_idx = fire_idx.T
+        num_events = len(fire_idx)
+        fire_idx_ori = fire_idx
+        spatial_vector = StandardScaler().fit_transform(fire_idx)
+
+        motion_array = np.zeros((num_events, 2))
+        motion_array[:, 0] = motion_vector[fire_idx[:, 0], fire_idx[:, 1], 0]
+        motion_array[:, 1] = motion_vector[fire_idx[:, 0], fire_idx[:, 1], 1]
+        motion_array = StandardScaler().fit_transform(motion_array)
+
+        motion_dis = pairwise_distances(motion_array, metric='euclidean')
+        spatial_dis = pairwise_distances(spatial_vector, metric='euclidean')
+        total_dis = 0.5 * (motion_dis + spatial_dis)
+
+        self.dbscan.fit(total_dis)
+        # self.optics.fit(total_dis)
+        # self.spectral.fit(total_dis)
+        labels = self.dbscan.labels_.astype(np.int)
+        # labels_optics = self.optics.labels_.astype(np.int)
+        # labels_spetral = self.spectral.labels_.astype(np.int)
+
+        return labels, fire_idx_ori
+        # return labels_optics, fire_idx_ori
+        # return labels_spetral, fire_idx_ori
+
+    def detection_object_with_motion(self, fireID, clusterId):
+        L = np.zeros((self.h, self.w), dtype=np.int)
+        L[fireID[:, 0], fireID[:, 1]] = clusterId + 1
+
+        structure = np.array([
+            [1, 1, 1],
+            [1, 1, 1],
+            [1, 1, 1]
+        ])
+
+        bboxSlices = mnts.find_objects(L)
+        box_num = clusterId.max() + 1
+        bbox = np.zeros((box_num, 4))
+        for iBox in range(box_num):
+            tmpBox = np.array(bboxSlices[iBox])
+            begin_X = tmpBox[0].start
+            end_X = tmpBox[0].stop
+            begin_Y = tmpBox[1].start
+            end_Y = tmpBox[1].stop
+
+            bbox[iBox, :] = [begin_X, begin_Y, end_X, end_Y]
+
+        # pprint(bbox)
+        return bbox
+

snnTracker/spkProc/detection/stdp_clustering.py

@@ -0,0 +1,397 @@
+# from config import *
+import numpy as np
+import torch
+from collections import namedtuple
+
+detect_box = namedtuple('detect_box', ['zId', 'box', 'velocity'])
+tracks = namedtuple('tracks', ['id', 'color', 'bbox', 'predbox', 'visible', 'vel', 'age', 'unvisible_count'])
+trajectories = namedtuple('trajectories', ['id', 'x', 'y', 't', 'color'])
+tracks_bbox = namedtuple('tracks_bbox', ['id', 't', 'x', 'y', 'h', 'w'])
+
+
+class stdp_cluster():
+
+    def __init__(self, spike_h, spike_w, box_size, device):
+
+        self.spike_h = spike_h
+        self.spike_w = spike_w
+        self.box_size = box_size
+        self.K1 = 1
+        self.K2 = 5
+        # self.InputSize = (2 * box_size + 1)**2
+        self.InputSize = (2 * box_size + 1) * (2 * box_size + 5)
+        self.device = device
+        # self.InputSize = box_size * (box_size + 4)
+        # self.InputSize = box_size**2
+
+        # self.synaptic_weight = torch.ones(self.K1, self.InputSize, dtype=torch.float64) / self.K1
+        self.synaptic_weight = torch.rand(self.K2, self.K1, self.InputSize, dtype=torch.float64)
+        # self.synaptic_weight = self.synaptic_weight / torch.sum(self.synaptic_weight)
+        # self.synaptic_weight = torch.unsqueeze(self.synaptic_weight, dim=0)
+        # self.synaptic_weight = self.synaptic_weight.repeat(self.K2, 1, 1)
+        self.bias_weight = torch.ones(self.K2, 1, dtype=torch.float64) / self.K2
+
+        self.synaptic_weight = self.normalization_w(self.synaptic_weight)
+        self.bias_weight = self.normalization_w(self.bias_weight)
+        self.synaptic_weight = self.synaptic_weight.to(device)
+        self.bias_weight = self.bias_weight.to(device)
+
+        self.learning_rate = 0.001
+        self.iter_num = 5
+        self.stdp_coefficience = 1
+
+        self.w_up = 1
+        self.w_low = -8
+
+        self.w_up_tensor = torch.ones_like(self.synaptic_weight, dtype=torch.float64) * self.w_up
+        self.w_low_tensor = torch.ones_like(self.synaptic_weight, dtype=torch.float64) * self.w_low
+
+        self.tracks = []
+        self.trajectories = []
+        self.tracks_bbox = []
+        # self.seed_everything(5)
+
+        self.background_occ_fr = 10  # background oscillation rate　20 Hz
+        self.occ_fr = torch.Tensor([self.background_occ_fr / 20000.0])
+
+        for i_neuron in range(self.K2):
+            self.tracks.append(tracks(i_neuron, 255 * np.random.rand(1, 3), torch.zeros((1, 4), dtype=torch.float64),
+                                      torch.zeros((1, 4), dtype=torch.float64), 0,
+                                      torch.zeros((2,)), 0, 0))
+            self.trajectories.append(trajectories(i_neuron, [], [], [], self.tracks[i_neuron].color))
+            self.tracks_bbox.append(tracks_bbox(i_neuron, [], [], [], [], []))
+
+
+    def normalization_w(self, weight):
+
+        if len(weight.shape) == 2:
+            exp_w = torch.exp(weight)
+            if torch.sum(exp_w) == 0:
+                norm_w = weight.clone()
+            else:
+                exp_w = exp_w / torch.sum(exp_w)
+                norm_w = torch.log(exp_w)
+
+            del exp_w
+        else:
+            exp_w = torch.exp(weight)
+            w_norm = torch.sum(exp_w, dim=2)
+            w_norm = torch.unsqueeze(w_norm, dim=2)
+            w_norm = w_norm.tile(1, 1, self.InputSize)
+            norm_w = exp_w.detach().clone()
+
+            valid_index = torch.where(w_norm != 0)
+            norm_w[valid_index] = exp_w[valid_index] / w_norm[valid_index]
+            norm_w = torch.log(norm_w)
+
+            del exp_w, w_norm, valid_index
+
+        torch.cuda.empty_cache()
+        return norm_w
+
+    # winner-take-all
+    # @staticmethod
+    def wta(self, attention_spikes, synaptic_weight, bias_weight):
+
+        attention_spikes = attention_spikes.double()
+        synaptic_weight = torch.squeeze(synaptic_weight + abs(self.w_low))
+        # synaptic_weight = torch.squeeze(synaptic_weight)
+        intPre = torch.matmul(synaptic_weight, attention_spikes)
+        intPre = torch.squeeze(intPre)
+        intPre[intPre > 700] = 700
+        psp = intPre + torch.squeeze(bias_weight)
+
+        rate_norm = (torch.Tensor([1]) * 1.0).to(self.device)
+        sum_exp = torch.sum(torch.exp(psp))
+        fire_inhb = torch.log(sum_exp) - torch.log(rate_norm)
+
+        tmp_psp = torch.exp(psp - fire_inhb)
+        tmp_psp[torch.isinf(tmp_psp)] = 0
+        fire_index = torch.where(torch.rand(1).to(self.device) < tmp_psp)
+        Z_spike = torch.zeros(tmp_psp.shape)
+        Z_spike[fire_index] = 1
+
+        del intPre, psp, rate_norm, sum_exp, fire_inhb, tmp_psp, fire_index
+        torch.cuda.empty_cache()
+        return Z_spike
+
+    @staticmethod
+    def intersect(box_a, box_b):
+        """ We resize both tensors to [A,B,2] without new malloc:
+        [A,2] -> [A,1,2] -> [A,B,2]
+        [B,2] -> [1,B,2] -> [A,B,2]
+        Then we compute the area of intersect between box_a and box_b.
+        Args:
+          box_a: (tensor) bounding boxes, Shape: [A,4].
+          box_b: (tensor) bounding boxes, Shape: [B,4].
+        Return:
+          (tensor) intersection area, Shape: [A,B].
+        """
+        A = box_a.size(0)
+        B = box_b.size(0)
+        max_xy = torch.min(box_a[:, 2:].unsqueeze(1).expand(A, B, 2),
+                           box_b[:, 2:].unsqueeze(0).expand(A, B, 2))
+        min_xy = torch.max(box_a[:, :2].unsqueeze(1).expand(A, B, 2),
+                           box_b[:, :2].unsqueeze(0).expand(A, B, 2))
+        inter = torch.clamp((max_xy - min_xy), min=0)
+        return inter[:, :, 0] * inter[:, :, 1]
+
+    def jaccard(self, box_a, box_b):
+        """Compute the jaccard overlap of two sets of boxes.  The jaccard overlap
+        is simply the intersection over union of two boxes.  Here we operate on
+        ground truth boxes and default boxes.
+        E.g.:
+            A ∩ B / A ∪ B = A ∩ B / (area(A) + area(B) - A ∩ B)
+        Args:
+            box_a: (tensor) Ground truth bounding boxes, Shape: [A,4]
+            box_b: (tensor) Prior boxes from priorbox layers, Shape: [B,4]
+        Return:
+            jaccard overlap: (tensor) Shape: [A, B]
+        """
+        inter = self.intersect(box_a, box_b)
+        area_a = ((box_a[:, 2] - box_a[:, 0]) *
+                  (box_a[:, 3] - box_a[:, 1])).unsqueeze(1).expand_as(inter)  # [A,B]
+        area_b = ((box_b[:, 2] - box_b[:, 0]) *
+                  (box_b[:, 3] - box_b[:, 1])).unsqueeze(0).expand_as(inter)  # [A,B]
+        union = area_a + area_b - inter
+        return inter / union  # [A,B]
+
+    def update_weight(self, attention_input):
+
+        n_attention = attention_input.shape[2]
+        predict_fire = torch.zeros(n_attention, self.K2)
+        synaptic_weight = self.synaptic_weight.detach().clone()
+        bias_weight = self.bias_weight.detach().clone()
+        lr_weight = torch.zeros(self.K2).to(self.device)
+        has_fired = np.zeros((self.K2, 1))
+
+        for iPattern in range(n_attention):
+            detected = -1
+            input_spike = torch.reshape(attention_input[:, :, iPattern], (-1, 1))
+            # background_noise = (torch.rand(input_spike.shape) < self.occ_fr).to(device)
+            # input_spike = (torch.logical_or((input_spike).type(torch.bool), background_noise)).type(torch.float32)
+            # input_spike = torch.reshape(attention_input, (-1, n_attention))
+            confusion_flag = 0
+            for i in range(self.iter_num):
+                z_spike = self.wta(input_spike, synaptic_weight, bias_weight).to(self.device)
+                dw_bias = self.learning_rate * (z_spike * torch.squeeze(torch.exp(-bias_weight)) - 1)
+                # tmp_sum = torch.sum(dw_bias, dim=1)
+                bias_weight += torch.unsqueeze(dw_bias, dim=1).detach()
+
+                for iZ in range(self.K2):
+                    if z_spike[iZ] != 0 and has_fired[iZ] == 0 and (detected == -1 or iZ == detected):
+                        has_fired[iZ] = 1
+                        detected = iZ
+                        # fire_idx = torch.where(z_spike[iZ, :]!=0)
+                        tmpE = torch.exp(-synaptic_weight[iZ, :, :])
+                        dw = self.stdp_coefficience * tmpE * torch.transpose(input_spike, 0, 1) - 1
+                        lr_weight[iZ] += 1
+
+                        synaptic_weight[iZ, :, :] += ((1.0 / lr_weight[iZ]) * dw.to(self.device))
+                        # synaptic_weight[iZ, :, :] += self.learning_rate * dw.to(device)
+
+                        synaptic_weight = torch.where(synaptic_weight < self.w_up, synaptic_weight,
+                                                      self.w_up_tensor).detach()
+                        synaptic_weight = torch.where(synaptic_weight < self.w_low, self.w_low_tensor,
+                                                      synaptic_weight).detach()
+
+                        synaptic_weight = self.normalization_w(synaptic_weight)
+                        bias_weight = self.normalization_w(bias_weight)
+                        predict_fire[iPattern, iZ] = torch.Tensor([1])
+
+                # predict_fire[iPattern, :] = z_spike.detach()
+            # predict_fire = z_spike.detach()
+        # print(synaptic_weight.max())
+        # print(synaptic_weight.min())
+
+        del n_attention, lr_weight
+        torch.cuda.empty_cache()
+        return predict_fire, synaptic_weight, bias_weight
+
+    def seed_everything(self, seed=11):
+        np.random.seed(seed)
+        torch.manual_seed(seed)
+        torch.cuda.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+
+        return
+
+    @staticmethod
+    def detect_object(predict_fire, attention_box, motion_id, motion_vector, **kwargs):
+
+        spike_h = kwargs.get('spike_h')
+        spike_w = kwargs.get('spike_w')
+        device = kwargs.get('device')
+        nAttention = attention_box.shape[0]
+        boxId = torch.zeros(nAttention, 1)
+        predBox = torch.zeros((nAttention, 4), dtype=torch.int)
+        velocities = torch.zeros(nAttention, 2).to(device)
+        predict_box = []
+
+        for iPattern in range(nAttention):
+            z_spike = predict_fire[iPattern, :]
+            if torch.any(z_spike != 0):
+                if len(torch.where(z_spike != 0)[0]) > 1:
+                    print('check')
+
+                tmp_fired = torch.where(z_spike != 0)[0]
+                boxId[iPattern] = tmp_fired[0] + 1
+
+                x = attention_box[iPattern, 0]
+                y = attention_box[iPattern, 1]
+                end_x = attention_box[iPattern, 2]
+                end_y = attention_box[iPattern, 3]
+
+                tmp_motion = torch.zeros(spike_h, spike_w)
+                tmp_motion[x:end_x + 1, y:end_y + 1] = motion_id[x:end_x + 1, y:end_y + 1].clone()
+
+                motion_index2d = torch.where(tmp_motion != 0)
+                if len(motion_index2d[0]) == 0:
+                    continue
+
+                motion_num = len(motion_index2d[0])
+                block_veloctiy = torch.zeros(motion_num, 2).to(device)
+                block_veloctiy[:, 0] = motion_vector[motion_index2d[0], motion_index2d[1], 0].clone()
+                block_veloctiy[:, 1] = motion_vector[motion_index2d[0], motion_index2d[1], 1].clone()
+                tmp_veloctiy = torch.mean(block_veloctiy, dim=0)
+                velocities[iPattern, :] = tmp_veloctiy.data
+                predBox[iPattern, :] = attention_box[iPattern, :]
+
+                predict_box.append(detect_box(boxId[iPattern],
+                                              torch.unsqueeze(predBox[iPattern], dim=0),
+                                              velocities[iPattern]))
+            # else:
+            #     print('no tracking neuron fire..')
+
+        del boxId, predBox, velocities
+        torch.cuda.empty_cache()
+
+        return predict_box
+
+    def update_tracks(self, detect_objects, sw, bw, timestep):
+
+        objects_num = len(detect_objects)
+        id_check = torch.zeros(self.K2, 1)
+        AssignTrk = []
+
+        for iObject in range(objects_num):
+            tmp_object = detect_objects[iObject]
+            id = int(tmp_object.zId.detach().item())
+            box = tmp_object.box
+            velocity = tmp_object.velocity
+
+            if id_check[id - 1] != 0:
+                if id in AssignTrk:
+                    # print('id %d repeat' % (id-1))
+                    AssignTrk.remove(id)
+                continue
+            else:
+                id_check[id - 1] = 1
+
+            pred_box = self.tracks[id - 1].predbox
+            boxes_iou = self.jaccard(box, pred_box)
+            unvisible_count = self.tracks[id - 1].unvisible_count
+            if ~(self.tracks[id - 1].predbox[0, 3] != 0 and self.tracks[id - 1].age > 15
+                 and boxes_iou < 0.6):
+                self.tracks[id - 1] = self.tracks[id - 1]._replace(bbox=box)
+                beginX = box[0, 0]
+                beginY = box[0, 1]
+                endX = box[0, 2]
+                endY = box[0, 3]
+
+                beginX = beginX + velocity[0] >= 0 and (beginX + velocity[0]) or 0
+                beginY = beginY + velocity[1] >= 0 and (beginY + velocity[1]) or 0
+                # endX = beginX + self.box_size * 2 < self.spike_h and (beginX + self.box_size * 2) or (self.spike_h - 1)
+                # endY = beginY + self.box_size * 2 < self.spike_w and (beginY + self.box_size * 2) or (self.spike_w - 1)
+                endX = endX + velocity[0] < self.spike_h and (endX + velocity[0]) or (self.spike_h - 1)
+                endY = endY + velocity[1] < self.spike_w and (endY + velocity[1]) or (self.spike_w - 1)
+
+                tmp_box = torch.tensor([beginX, beginY, endX, endY])
+                tmp_box = torch.unsqueeze(tmp_box, dim=0)
+                self.tracks[id - 1] = self.tracks[id - 1]._replace(predbox=tmp_box)
+
+                self.tracks[id - 1] = self.tracks[id - 1]._replace(visible=1)
+                self.tracks[id - 1] = self.tracks[id - 1]._replace(vel=velocity)
+                self.tracks[id - 1] = self.tracks[id - 1]._replace(unvisible_count=0)
+                self.tracks[id - 1] = self.tracks[id - 1]._replace(age=self.tracks[id - 1].age + 1)
+
+                # update the trajectories
+                self.trajectories[id - 1].x.append((box[0, 0] + self.box_size).item())
+                self.trajectories[id - 1].y.append((box[0, 1] + self.box_size).item())
+                self.trajectories[id - 1].t.append(timestep)
+
+                # Check if beginX, beginY, endX, endY are int; otherwise, use .item()
+                self.tracks_bbox[id - 1].x.append(beginY if isinstance(beginY, int) else beginY.item())
+                self.tracks_bbox[id - 1].y.append(beginX if isinstance(beginX, int) else beginX.item())
+                self.tracks_bbox[id - 1].h.append(
+                    (endX - beginX) if isinstance(endX, int) and isinstance(beginX, int) else (endX - beginX).item())
+                self.tracks_bbox[id - 1].w.append(
+                    (endY - beginY) if isinstance(endY, int) and isinstance(beginY, int) else (endY - beginY).item())
+
+                self.tracks_bbox[id - 1].t.append(timestep)
+                AssignTrk.append(id)
+                # print('tracks %d velocity dx: %f dy: %f' % (id, velocity[0], velocity[1]))
+
+        all_id = list(range(1, self.K2 + 1, 1))
+        noAssign = [x for x in all_id if x not in AssignTrk]
+
+        noAssign_num = self.K2 - len(AssignTrk)
+
+        for iObject in range(noAssign_num):
+            id = noAssign[iObject]
+            unvisible_count = self.tracks[id - 1].unvisible_count
+            if unvisible_count > 5:
+                self.tracks[id - 1] = self.tracks[id - 1]._replace(age=0)
+                self.tracks[id - 1] = self.tracks[id - 1]._replace(visible=0)
+                # sw[id-1, :, :] = 1 / self.K1
+                # bw[id-1] = 1 / self.K2
+            else:
+                if self.tracks[id - 1].predbox[0, 2] != 0:
+                    self.tracks[id - 1] = self.tracks[id - 1]._replace(bbox=self.tracks[id - 1].predbox)
+                    beginX = self.tracks[id - 1].predbox[0, 0].item()
+                    beginY = self.tracks[id - 1].predbox[0, 1].item()
+                    endX = self.tracks[id - 1].predbox[0, 2].item()
+                    endY = self.tracks[id - 1].predbox[0, 3].item()
+
+                    beginX = beginX + self.tracks[id - 1].vel[0] >= 0 and (beginX + self.tracks[id - 1].vel[0]) or 0
+                    beginY = beginY + self.tracks[id - 1].vel[1] >= 0 and (beginY + self.tracks[id - 1].vel[1]) or 0
+                    # endX = beginX + self.box_size * 2 < self.spike_h and (beginX + self.box_size * 2) or (self.spike_h - 1)
+                    # endY = beginY + self.box_size * 2 < self.spike_w and (beginY + self.box_size * 2) or (self.spike_w - 1)
+                    endX = endX + self.tracks[id - 1].vel[0] < self.spike_h and (endX + self.tracks[id - 1].vel[0]) or (
+                                self.spike_h - 1)
+                    endY = endY + self.tracks[id - 1].vel[1] < self.spike_w and (endY + self.tracks[id - 1].vel[1]) or (
+                                self.spike_w - 1)
+
+                    pred_box = torch.tensor([beginX, beginY, endX, endY])
+                    pred_box = torch.unsqueeze(pred_box, dim=0)
+                    self.tracks[id - 1] = self.tracks[id - 1]._replace(predbox=pred_box)
+
+                    self.trajectories[id - 1].x.append(
+                        (beginX + self.box_size) if isinstance(beginX, int) else (beginX + self.box_size).item())
+                    self.trajectories[id - 1].y.append(
+                        (beginY + self.box_size) if isinstance(beginY, int) else (beginY + self.box_size).item())
+                    self.trajectories[id - 1].t.append(timestep)
+
+                    # Check if beginX, beginY, endX, endY are int; otherwise, use .item()
+                    self.tracks_bbox[id - 1].x.append(beginY if isinstance(beginY, int) else beginY.item())
+                    self.tracks_bbox[id - 1].y.append(beginX if isinstance(beginX, int) else beginX.item())
+                    self.tracks_bbox[id - 1].h.append(
+                        (endX - beginX) if isinstance(endX, int) and isinstance(beginX, int) else (
+                                endX - beginX).item())
+                    self.tracks_bbox[id - 1].w.append(
+                        (endY - beginY) if isinstance(endY, int) and isinstance(beginY, int) else (
+                                endY - beginY).item())
+                    self.tracks_bbox[id - 1].t.append(timestep)
+                    # print('predicting location of object %d the %d time' % (id, unvisible_count))
+
+                # print('tracks %d predictive velocity dx: %f, dy: %f' % (
+                # id, self.tracks[id - 1].vel[0], self.tracks[id - 1].vel[1]))
+                if ~(torch.all(self.synaptic_weight == 1)):
+                    sw[id - 1, :, :] = self.synaptic_weight[id - 1, :, :].detach().clone()
+                    bw[id - 1] = self.bias_weight[id - 1].detach().clone()
+                    # print('correct the weight')
+            self.tracks[id - 1] = self.tracks[id - 1]._replace(unvisible_count=self.tracks[id - 1].unvisible_count + 1)
+
+        return sw, bw

snnTracker/spkProc/filters/stp_filters_torch.py

@@ -0,0 +1,170 @@
+# -*- coding: utf-8 -*- 
+# @Time : 2021/11/19 16:25 
+# @Author : Yajing Zheng
+# @File : stp_filters_torch.py
+import copy
+
+import torch
+import numpy as np
+
+class STPFilter:
+
+    def __init__(self, spike_h, spike_w, device, diff_time=1, **STPargs):
+        self.spike_h = spike_h
+        self.spike_w = spike_w
+        self.device = device
+
+        # specify stp parameters
+        if STPargs.get('u0', None) is None:
+            self.u0 = 0.1
+            self.D = 0.02
+            self.F = 1.7
+            self.f = 0.11
+            self.time_unit = 2000
+
+        else:
+            self.u0 = STPargs.get('u0')
+            self.D = STPargs.get('D')
+            self.F = STPargs.get('F')
+            self.f = STPargs.get('f')
+            self.time_unit = STPargs.get('time_unit')
+
+        self.r0 = 1
+
+        self.diff_time = diff_time # duration of window for record past dynamics for calculating the differnece
+        self.R = torch.ones(self.spike_h, self.spike_w) * self.r0
+        self.u = torch.ones(self.spike_h, self.spike_w) * self.u0
+        self.r_old = torch.ones(self.diff_time, self.spike_h, self.spike_w) * self.r0
+
+        self.R = self.R.to(self.device)
+        self.u = self.u.to(self.device)
+        self.r_old = self.r_old.to(self.device)
+
+        # LIF detect layer parameters
+        self.detectVoltage = torch.zeros(self.spike_h, self.spike_w).to(self.device)
+        if STPargs.get('lifSize', None) is None:
+            lifSize = 3
+            paddingSize = 1
+        else:
+            lifSize = STPargs.get('lifSize')
+            paddingSize = int((lifSize - 1) / 2)
+
+        self.lifConv = torch.nn.Conv2d(in_channels=1, out_channels=1, kernel_size=(lifSize, lifSize),
+                                       padding=(paddingSize, paddingSize),
+                                       bias=False)
+        self.lifConv.weight.data = torch.ones(1, 1, lifSize, lifSize) * 3.0
+
+        self.lifConv = self.lifConv.to(self.device)
+        if STPargs.get('filterThr', None) is None:
+            self.filterThr = 0.1  # filter threshold
+            self.voltageMin = -8
+            self.lifThr = 2
+        else:
+            self.filterThr = STPargs.get('filterThr')
+            self.voltageMin = STPargs.get('voltageMin')
+            self.lifThr = STPargs.get('lifThr')
+
+        self.filter_spk = torch.zeros(self.spike_h, self.spike_w).to(self.device)
+        self.lif_spk = torch.zeros(self.spike_h, self.spike_w).to(self.device)
+        self.spikePrevMnt = torch.zeros([self.spike_h, self.spike_w], device=self.device)
+        self.stp_gradient = 0
+        self.adjusted_threshold = torch.zeros(self.spike_h, self.spike_w).to(self.device)
+
+    def update_dynamics(self, curT, spikes):
+
+        spikeCurMnt = self.spikePrevMnt.detach().clone()
+        spike_bool = spikes.bool()
+        spikeCurMnt[spike_bool] = curT + 1
+        dttimes = spikeCurMnt - self.spikePrevMnt
+        dttimes = dttimes / self.time_unit
+        exp_D = torch.exp((-dttimes[spike_bool] / self.D))
+        self.R[spike_bool] = 1 - (1 - self.R[spike_bool] * (1 - self.u[spike_bool])) * exp_D
+        exp_F = torch.exp((-dttimes[spike_bool] / self.F))
+        self.u[spike_bool] = self.u0 + (
+                self.u[spike_bool] + self.f * (1 - self.u[spike_bool]) - self.u0) * exp_F
+
+        tmp_diff = torch.abs(self.R - self.r_old[0])
+        # 根据梯度动态调整滤波器阈值
+        self.stp_gradient = (0.5 * self.stp_gradient + 0.5 * torch.div(tmp_diff, self.R))
+        gradient_sqrt = torch.from_numpy(np.sqrt(self.stp_gradient.cpu().numpy()) + 1).to(self.device)
+        self.adjusted_threshold = torch.div(self.filterThr, gradient_sqrt)
+
+        self.filter_spk[:] = 0
+        # self.filter_spk[spike_bool & (tmp_diff >= self.filterThr)] = 1
+        self.filter_spk[spike_bool & (tmp_diff >= self.adjusted_threshold)] = 1
+
+        if curT < self.diff_time:
+            self.r_old[curT] = self.R.detach().clone()
+        else:
+            self.r_old[0:-1] = self.r_old[1:].detach().clone()
+            self.r_old[-1] = self.R.detach().clone()
+        self.spikePrevMnt = spikeCurMnt.detach().clone()
+        del spikeCurMnt, dttimes, exp_D, exp_F, tmp_diff
+
+    def update_dynamic_offline(self, spikes, intervals):
+
+        isi_num = intervals.shape[0]
+        R = torch.ones(isi_num, self.spike_h, self.spike_w) * self.r0
+        u = torch.ones(isi_num, self.spike_h, self.spike_w) * self.u0
+        prev_isi = intervals[0, :, :]
+
+        for t in range(1, isi_num):
+            tmp_isi = intervals[t, :, :]
+            update_idx = (tmp_isi != prev_isi) & (spikes[t, :, :] == 1) | (tmp_isi == 1)
+            tmp_isi = torch.from_numpy(tmp_isi).to(self.device).float()
+
+            exp_D = torch.exp((-tmp_isi[update_idx] / self.D))
+            self.R[update_idx] = 1 - (1 - self.R[update_idx] * (1 - self.u[update_idx])) * exp_D
+            exp_F = torch.exp((-tmp_isi[update_idx] / self.F))
+            self.u[update_idx] = self.u0 + (
+                    self.u[update_idx] + self.f * (1 - self.u[update_idx]) - self.u0) * exp_F
+
+            tmp_r = self.R.detach().clone()
+            tmp_u = self.u.detach().clone()
+            R[t, :, :] = copy.deepcopy(tmp_r)
+            u[t, :, :] = copy.deepcopy(tmp_u)
+
+        return R, u
+
+    def local_connect(self, spikes):
+        inputSpk = torch.reshape(spikes, (1, 1, self.spike_h, self.spike_w)).float()
+        # tmp_fired = spikes != 0
+        self.detectVoltage[spikes == False] -= 1
+        tmpRes = self.lifConv(inputSpk)
+        tmpRes = torch.squeeze(tmpRes).to(self.device)
+        self.detectVoltage += tmpRes.data
+        self.detectVoltage[self.detectVoltage < self.voltageMin] = self.voltageMin
+
+        self.lif_spk[:] = 0
+        self.lif_spk[self.detectVoltage >= self.lifThr] = 1
+        self.detectVoltage[self.detectVoltage >= self.lifThr] *= 0.8
+        # self.detectVoltage[(self.detectVoltage < self.lifThr) & (self.detectVoltage > 0)] = 0
+
+        del inputSpk, tmpRes
+
+    def local_connect_offline(self, spikes):
+        timestamps = spikes.shape[0]
+        tmp_voltage = []
+        lif_spk = []
+
+        for iSpk in range(timestamps):
+            tmp_spikes = spikes[iSpk]
+            tmp_spk = torch.from_numpy(spikes[iSpk]).to(self.device)
+            inputSpk = torch.reshape(tmp_spk, (1, 1, self.spike_h, self.spike_w)).float()
+            # tmp_fired = spikes != 0
+            self.detectVoltage[tmp_spikes == 0] -= 1
+            tmpRes = self.lifConv(inputSpk)
+            tmpRes = torch.squeeze(tmpRes).to(self.device)
+            self.detectVoltage += tmpRes.data
+            self.detectVoltage[self.detectVoltage < self.voltageMin] = self.voltageMin
+
+            self.lif_spk[:] = 0
+            self.lif_spk[self.detectVoltage >= self.lifThr] = 1
+            # self.detectVoltage[(self.detectVoltage < self.lifThr) & (self.detectVoltage > 0)] = 0
+            self.detectVoltage[self.detectVoltage >= self.lifThr] *= 0.8
+            voltage = self.detectVoltage.cpu().detach().numpy()
+            tmp_voltage.append(copy.deepcopy(voltage))
+            lif_spk.append(self.lif_spk.cpu().detach().numpy())
+
+        del inputSpk, tmpRes
+        return tmp_voltage, lif_spk

snnTracker/spkProc/motion/motion_detection.py

@@ -0,0 +1,347 @@
+from utils import get_kernel, get_transform_matrix_new, visualize_images
+import torchgeometry as tgm
+import numpy as np
+import torch.nn.functional as F
+import torch
+
+class motion_estimation:
+
+    def __init__(self, dvs_h, dvs_w, device, logger):
+
+        self.dvs_h = dvs_h
+        self.dvs_w = dvs_w
+        self.device = device
+        self.logger = logger
+
+        # motion parameters
+        self.orientation = range(0, 180 - 1, int(180 / 4))
+        # eight moving direction
+        '''
+                self.ori = torch.Tensor(np.array([[-1, -1],
+                    [0, -1],
+                    [1, -1],
+                    [-1, 0],
+                    [1, 0],
+                    [-1, 1],
+                    [0, 1],
+                    [1, 1]], dtype=np.uint8)).to(self.device)
+        '''
+        # self.ori = np.array([[-1, -1],
+        #                      [0, -1],
+        #                      [1, -1],
+        #                      [1, 0],
+        #                      [-1, 0],
+        #                      [-1, 1],
+        #                      [0, 1],
+        #                      [1, 1]], dtype=np.int32)
+
+        self.ori = np.array([[1, 0],
+                             [1, 1],
+                             [0, 1],
+                             [-1, 1],
+                             [-1, 0],
+                             [-1, -1],
+                             [0, -1],
+                             [1, -1]], dtype=np.int32)
+        # self.ori = np.array(self.ori, dtype=np.int)
+        # self.speed = torch.Tensor(np.array([1, 2, 3, 4], np.uint8)).to(self.device)
+        self.speed = np.array([1, 2, 3, 4, 5, 6], dtype=np.int32)
+        # self.speed = np.array([1], dtype=np.int32)
+        # self.ori_x = torch.from_numpy(self.ori[:, 0]).to(self.device)
+        # self.ori_y = torch.from_numpy(self.ori[:, 1]).to(self.device)
+        self.ori_x = torch.from_numpy(np.expand_dims(self.ori[:, 0], axis=1)).to(self.device).float()
+        self.ori_y = torch.from_numpy(np.expand_dims(self.ori[:, 1], axis=1)).to(self.device).float()
+
+        self.warp_matrix = get_transform_matrix_new(self.ori, self.speed, self.dvs_w, self.dvs_h, self.device)
+        self.track_pre = torch.zeros(self.dvs_h, self.dvs_w)
+
+        self.num_ori = len(self.ori)
+        self.num_speed = len(self.speed)
+        self.motion_pattern_num = self.num_ori * self.num_speed
+        self.motion_weight = torch.ones(self.motion_pattern_num, 1, self.dvs_h, self.dvs_w) / self.motion_pattern_num
+        self.tracking_threshold = 1
+
+        # self.local_pool_size = 21
+        self.local_pool_size = 11
+        # self.local_pool_size = 5
+        padding_width = int((self.local_pool_size - 1) / 2)
+        self.pool_kernel = torch.nn.Conv2d(in_channels=1, out_channels=1,
+                                           kernel_size=(self.local_pool_size, self.local_pool_size),
+                                           padding=(padding_width, padding_width), bias=False)
+        self.pool_kernel.weight.data = torch.ones(1, 1, self.local_pool_size, self.local_pool_size)
+
+        self.gaussian_kernel = torch.nn.Conv2d(in_channels=1, out_channels=1,
+                                               kernel_size=(self.local_pool_size, self.local_pool_size),
+                                               padding=(padding_width, padding_width), bias=False)
+        tmp_filter = get_kernel(self.local_pool_size, round(self.local_pool_size / 4))
+        tmp_filter = tmp_filter.reshape((1, 1, self.local_pool_size, self.local_pool_size))
+        self.gaussian_kernel.weight.data = torch.from_numpy(tmp_filter).float()
+
+        # local wta inhibition size
+        # inh_size = 15
+        self.inh_size = 25
+        # inh_size = 11
+        self.padding_width = int((self.inh_size - 1) / 2)
+        self.inhb_kernel = torch.nn.Conv2d(in_channels=1, out_channels=1,
+                                           kernel_size=(self.inh_size, self.inh_size),
+                                           padding=(self.padding_width, self.padding_width), bias=False)
+        self.inhb_kernel.weight.data = torch.ones(1, 1, self.inh_size, self.inh_size)
+        self.inhb_threshold = 5
+
+        self.track_pre = self.track_pre.to(self.device)
+        self.motion_weight = self.motion_weight.to(self.device)
+        self.pool_kernel = self.pool_kernel.to(self.device)
+        self.gaussian_kernel = self.gaussian_kernel.to(self.device)
+        self.inhb_kernel = self.inhb_kernel.to(self.device)
+
+        # cc_motion = [[0, 33, 238],
+        #              [79, 0, 255],
+        #               [229, 0, 237],
+        #               [188, 0, 26],
+        #               [191, 198, 0],
+        #               [129, 241, 0],
+        #               [0, 205, 106],
+        #               [0, 205, 198]]
+
+        cc_motion = [[0, 255, 255],
+                     [205, 95, 85],
+                     [11, 134, 184],
+                     [255, 255, 0],
+                     [154, 250, 0],
+                     [147, 20, 255],
+                     [240, 32, 160],
+                     [48, 48, 255]]
+
+        cc_motion = np.transpose(np.array(cc_motion, dtype=np.float32))
+        self.cc_motion = torch.from_numpy(cc_motion / 255)
+        self.cc_motion = self.cc_motion.to(self.device)
+        self.learning_rate = 0.1
+
+        '''
+        self.dw_ltp = torch.zeros(self.motion_pattern_num, 1, self.dvs_h, self.dvs_w)
+        self.dw_ltd = torch.zeros(self.motion_pattern_num, 1, self.dvs_h, self.dvs_w)
+        self.dw_ltp = self.dw_ltp.to(self.device)
+        self.dw_ltd = self.dw_ltd.to(self.device)       
+        '''
+
+    def stdp_tracking(self, spikes):
+        track_post = torch.reshape(spikes, (1, 1, self.dvs_h, self.dvs_w))
+        tmp_pool = self.pool_kernel(track_post)
+        tmp_pool = tmp_pool.repeat(self.motion_pattern_num, 1, 1, 1)
+
+        predict_fired = torch.zeros(self.motion_pattern_num, 1, self.dvs_h, self.dvs_w).to(self.device)
+        fire_idx = torch.where(spikes != 0)
+
+        for i_ori in range(self.num_ori):
+            for i_speed in range(self.num_speed):
+                i_motion = i_ori * self.num_speed + i_speed
+                x = fire_idx[0] + self.ori[i_ori, 0] * self.speed[i_speed]
+                y = fire_idx[1] + self.ori[i_ori, 1] * self.speed[i_speed]
+                invalid_idx = torch.logical_or(torch.logical_or(x > self.dvs_h - 1, x < 0),
+                                               torch.logical_or(y > self.dvs_w - 1, y < 0))
+                x[invalid_idx] = 0
+                y[invalid_idx] = 0
+                predict_fired[i_motion, 0, x, y] = 1
+
+        # track_post = track_post.repeat(self.motion_pattern_num, 1, 1, 1)
+        # grid = F.affine_grid(self.warp_matrix, track_post.shape)
+        # predict_fired = F.grid_sample(track_post, grid, padding_mode='zeros', align_corners=True)
+        # # invalid_index = torch.where(torch.logical_and(predict_fired != 0, predict_fired != 1))
+        # invalid_index = torch.where(predict_fired>0)
+        # predict_fired[invalid_index] = 1
+        # predict_fired[torch.where(predict_fired<1)] = 0
+
+        track_pre_exp = torch.unsqueeze(self.track_pre, 0).repeat(self.motion_pattern_num, 1, 1)
+        track_pre_exp = torch.unsqueeze(track_pre_exp, 1)
+
+        # STDP update the motion weight
+        dw_ltd = torch.zeros(self.motion_pattern_num, 1, self.dvs_h, self.dvs_w).to(self.device)
+        dw_ltp = torch.zeros(self.motion_pattern_num, 1, self.dvs_h, self.dvs_w).to(self.device)
+
+        tmp_bool = torch.eq(predict_fired, track_pre_exp)
+        index = torch.where(torch.logical_and(tmp_bool, track_pre_exp == 1))
+        if len(index[0]) != 0:
+            dw_ltp[index] = 1
+
+        index = torch.where(torch.logical_and(~tmp_bool, predict_fired == 1))
+        if len(index[0]) != 0:
+            dw_ltd[index] = 2
+
+        dw_ltp = self.pool_kernel(dw_ltp)
+        dw_ltd = self.pool_kernel(dw_ltd)
+
+        # dw_ltp = self.gaussian_kernel(dw_ltp)
+        # dw_ltd = self.gaussian_kernel(dw_ltd)
+
+        # dw = dw_ltp - dw_ltd
+        # dw = self.gaussian_kernel(dw_ltp - dw_ltd)
+        # tmp_pool[torch.where(tmp_pool == 0)] = 1
+        dw = torch.div((dw_ltp - dw_ltd), tmp_pool)
+        # dw = dw / tmp_pool
+        # dw = dw_ltp - dw_ltd
+        self.motion_weight += self.learning_rate * dw.detach().clone()
+
+        max_weight, _ = torch.max(self.motion_weight, dim=0)
+        min_weight, _ = torch.min(self.motion_weight, dim=0)
+
+        for iMotion in range(self.motion_pattern_num):
+            tmp_weight = self.motion_weight[iMotion, :, :, :].detach()
+            tmp_weight = (tmp_weight - min_weight) / (max_weight - min_weight)
+            self.motion_weight[iMotion, :, :, :] = tmp_weight.detach()
+
+        # self.motion_weight.data = F.normalize(self.motion_weight, p=2, dim=0)
+        self.motion_weight[torch.isnan(self.motion_weight)] = 0
+        # self.motion_weight[torch.isinf(self.motion_weight)] = 0
+        self.track_pre = spikes.detach().clone()
+
+        del track_post, tmp_pool, predict_fired, track_pre_exp, tmp_bool, dw
+        del tmp_weight, max_weight, min_weight, spikes
+        del dw_ltd, dw_ltp
+        torch.cuda.empty_cache()
+
+    def local_wta(self, spikes, timestamp, visualize=False):
+        input_spike = torch.reshape(spikes, (1, 1, self.dvs_h, self.dvs_w))
+
+        motion_vector_layer1 = torch.zeros(self.dvs_h, self.dvs_w, 2, dtype=torch.float32).to(self.device)
+        max_w, max_wid = torch.max(self.motion_weight, dim=0)
+        max_wid = torch.squeeze(max_wid)
+        speedId = (max_wid % self.num_speed).detach()
+        oriId = (torch.floor(max_wid / self.num_speed)).detach()
+
+        tmp_weight = self.motion_weight.permute(2, 3, 1, 0)
+        # change the dimension of matrix from (ori_num, speed_num, height, width) to (h,w, speed_num, ori_num)
+        tmp_weight = torch.reshape(tmp_weight, [self.dvs_h, self.dvs_w, self.num_ori, self.num_speed])
+        tmp_weight = tmp_weight.permute(0, 1, 3, 2)
+        tmp_weight_x = torch.matmul(tmp_weight, self.ori_x)
+        tmp_weight_y = torch.matmul(tmp_weight, self.ori_y)
+        # tmp_weight_x = torch.reshape(torch.mm(tmp_weight, self.ori_x), [self.dvs_h, self.dvs_w, self.num_speed])
+        # tmp_weight_y = torch.reshape(torch.mm(tmp_weight, self.ori_y), [self.dvs_h, self.dvs_w, self.num_speed])
+
+        max_w = torch.squeeze(max_w)
+        fired_spk_index2d = torch.where(torch.logical_and(spikes != 0, max_w > 0))
+
+        # speedId = speedId[fired_spk_index2d].cpu().numpy()
+        # oriId = oriId[fired_spk_index2d].int().cpu().numpy()
+        #
+        # dx = -1 * self.ori[oriId, 1] * self.speed[speedId]
+        # dy = -1 * self.ori[oriId, 0] * self.speed[speedId]
+
+        tmp_weight_x = torch.mean(tmp_weight_x, dim=2)
+        tmp_weight_y = torch.mean(tmp_weight_y, dim=2)
+        tmp_weight_x = torch.squeeze(tmp_weight_x)
+        tmp_weight_y = torch.squeeze(tmp_weight_y)
+
+        dx = tmp_weight_x[fired_spk_index2d]
+        dy = tmp_weight_y[fired_spk_index2d]
+
+        motion_vector_layer1[fired_spk_index2d[0], fired_spk_index2d[1], 0] = dx
+        motion_vector_layer1[fired_spk_index2d[0], fired_spk_index2d[1], 1] = dy
+        dy_numpy = dy.cpu().numpy()
+        dx_numpy = dx.cpu().numpy()
+
+        # motion_vector_layer1[fired_spk_index2d[0], fired_spk_index2d[1], 0] = torch.from_numpy(dx.astype('float32')).to(self.device)
+        # motion_vector_layer1[fired_spk_index2d[0], fired_spk_index2d[1], 1] = torch.from_numpy(dy.astype('float32')).to(self.device)
+        # dx_numpy = dx
+        # dy_numpy = dy
+
+        rotAng = np.arctan2(-dy_numpy, dx_numpy) * 180 / np.pi + 180
+        rotAng[np.where(rotAng == 360)] = 0
+        tmp_motion = np.floor(rotAng / (360 / 8))
+        # tmp_motion[np.where(tmp_motion == 8)] = 0
+        track_voltage = torch.zeros(self.num_ori, self.dvs_h, self.dvs_w)
+        track_voltage[tmp_motion, fired_spk_index2d[0], fired_spk_index2d[1]] = 1
+
+        track_voltage = torch.unsqueeze(track_voltage, 1)
+        track_voltage = track_voltage.to(self.device)
+        track_voltage = torch.squeeze(self.inhb_kernel(track_voltage))
+        max_v, max_vid = torch.max(track_voltage, dim=0)
+
+        fired_layer2_index = torch.where(
+            torch.logical_and(max_v >= self.inhb_threshold, torch.logical_and(spikes != 0, max_w > 0)))
+        max_motion = torch.zeros(self.dvs_h, self.dvs_w, dtype=torch.int64)
+        max_motion_layer1 = torch.zeros(self.dvs_h, self.dvs_w, dtype=torch.int64)
+        max_motion_layer1 = max_motion_layer1.to(self.device)
+
+        motion_vector_max = torch.zeros(self.dvs_h, self.dvs_w, 2, dtype=torch.float32)
+        max_motion = max_motion.to(self.device)
+        motion_vector_max = motion_vector_max.to(self.device)
+
+        max_motion[fired_layer2_index] = max_vid[fired_layer2_index].detach() + 1
+        motion_tensor = torch.from_numpy(tmp_motion + 1).to(self.device)
+        max_motion_layer1[fired_spk_index2d[0], fired_spk_index2d[1]] = motion_tensor.long()
+        max_motion_layer1[max_motion == 0] = 0
+
+        # 1. find the difference between m1 and mc motion
+        tmp_vid = max_vid[fired_layer2_index].cpu().detach().numpy()
+        if len(tmp_vid) != 0:
+            motion_vector_max[fired_layer2_index] = motion_vector_layer1[fired_layer2_index].detach()
+            loser_pattern_index = torch.where(torch.logical_and(max_motion != 0, max_motion_layer1 != max_motion))
+            fired2_index_x = loser_pattern_index[0]
+            fired2_index_y = loser_pattern_index[1]
+            voltage_block = max_v[None, None, :, :]
+            voltage_block = F.pad(voltage_block, (self.padding_width, self.padding_width, self.padding_width, self.padding_width),
+                                  mode='constant', value=0)
+            voltage_block = F.unfold(voltage_block, (self.inh_size, self.inh_size))
+            voltage_block = voltage_block.reshape([1, self.inh_size*self.inh_size, self.dvs_h, self.dvs_w])
+            offset_pattern = torch.argmax(voltage_block, dim=1)
+            offset_pattern = torch.squeeze(offset_pattern)
+            offset_pattern_loser = offset_pattern[fired2_index_x, fired2_index_y]
+            offset_x = offset_pattern_loser / self.inh_size - self.padding_width
+            offset_y = torch.fmod(offset_pattern_loser, self.inh_size) - self.padding_width
+            offset_x = offset_x.int()
+            offset_y = offset_y.int()
+            motion_vector_max[fired2_index_x, fired2_index_y, :] = motion_vector_max[fired2_index_x + offset_x,
+                                                                                     fired2_index_y + offset_y, :]
+
+        # for i_vector in range(len(fired2_index_x)):
+        #
+        #     if fired2_index_x[i_vector] - self.padding_width < 0:
+        #         x_begin = 0
+        #     else:
+        #         x_begin = fired2_index_x[i_vector] - self.padding_width
+        #
+        #     if fired2_index_x[i_vector] + self.padding_width >= self.dvs_h:
+        #         x_end = self.dvs_h
+        #     else:
+        #         x_end = fired2_index_x[i_vector] + self.padding_width
+        #
+        #     if fired2_index_y[i_vector] - self.padding_width < 0:
+        #         y_begin = 0
+        #     else:
+        #         y_begin = fired2_index_y[i_vector] - self.padding_width
+        #
+        #     if fired2_index_y[i_vector] + self.padding_width >= self.dvs_w:
+        #         y_end = self.dvs_w
+        #     else:
+        #         y_end = fired2_index_y[i_vector] + self.padding_width
+        #
+        #     winner_motion = max_motion[fired2_index_x[i_vector], fired2_index_y[i_vector]]
+        #     motion_id_block = max_motion_layer1[x_begin:x_end, y_begin:y_end]
+        #     motion_block = motion_vector_max[x_begin:x_end, y_begin:y_end, :]
+        #     motion_voltage_block = max_v[x_begin:x_end, y_begin:y_end]
+        #     winner_id = torch.where(motion_id_block == winner_motion)
+        #     if len(winner_id[0]) > 0:
+        #         winner_mv = motion_block[winner_id[0], winner_id[1], :]
+        #         winner_voltage = motion_voltage_block[winner_id[0], winner_id[1]]
+        #         motion_vector_max[fired2_index_x[i_vector], fired2_index_y[i_vector], :] = \
+        #             winner_mv[torch.argmax(winner_voltage), :]
+
+        # 2. replace the loser motion pattern
+
+        if visualize is True:
+            Image_layer1 = torch.zeros(3, self.dvs_h, self.dvs_w).to(self.device)
+            Image_layer1[:, fired_spk_index2d[0], fired_spk_index2d[1]] = self.cc_motion[:, tmp_motion]
+
+            Image_layer2 = torch.zeros(3, self.dvs_h, self.dvs_w).to(self.device)
+            Image_layer2[:, fired_layer2_index[0], fired_layer2_index[1]] = self.cc_motion[:, tmp_vid]
+
+            self.logger.add_image('motion_estimation/M1 estimation', Image_layer1, timestamp)
+            self.logger.add_image('motion_estimation/MC estimation', Image_layer2, timestamp)
+
+        # track_voltage.to(self.device_cpu)
+
+        del input_spike, fired_spk_index2d, fired_layer2_index
+        del track_voltage, dx, dy
+        torch.cuda.empty_cache()
+
+        return max_motion, motion_vector_max, motion_vector_layer1

snnTracker/spkProc/tracking/snn_tracker.py

@@ -0,0 +1,227 @@
+# -*- coding: utf-8 -*- 
+# @Time : 2023/7/16 20:23 
+# @Author : Yajing Zheng
+# @Email: [email protected]
+# @File : snn_tracker.py
+import os, sys
+sys.path.append('../..')
+import time
+import numpy as np
+import torch
+
+from spkProc.filters.stp_filters_torch import STPFilter
+# from filters import stpFilter
+from spkProc.detection.attention_select import SaccadeInput
+from spkProc.motion.motion_detection import motion_estimation
+from spkProc.detection.stdp_clustering import stdp_cluster
+from utils import NumpyEncoder
+from collections import namedtuple
+import json
+import cv2
+from tqdm import tqdm
+
+trajectories = namedtuple('trajectories', ['id', 'x', 'y', 't', 'color'])
+
+class SNNTracker:
+
+    def __init__(self, spike_h, spike_w, device, attention_size=20, diff_time=1, **STPargs):
+        self.spike_h = spike_h
+        self.spike_w = spike_w
+        self.device = device
+
+        # self.stp_filter = STPFilter(spike_h, spike_w, device)
+        if STPargs is not None:
+            self.stp_filter = STPFilter(spike_h, spike_w, device, diff_time, **STPargs)
+        else:
+            self.stp_filter = STPFilter(spike_h, spike_w, device, diff_time)
+        # self.stp_filter = stpFilter()
+        self.attention_size = attention_size
+        self.object_detection = SaccadeInput(spike_h, spike_w, box_size=self.attention_size, device=device)
+        from tensorboardX import SummaryWriter
+        logger = SummaryWriter(log_dir='data/log_pkuvidar')
+        self.motion_estimator = motion_estimation(spike_h, spike_w, device, logger=logger)
+        # gpu_tracker.track()  # run function between the code line where uses GPU
+
+        self.object_cluster = stdp_cluster(spike_h, spike_w, box_size=self.attention_size, device=device)
+
+        # self.timestamps = spikes.shape[0]
+        # self.filterd_spikes = np.zeros([self.timestamps, self.spike_h, self.spike_w], np.uint8)
+        self.calibration_time = 150
+        self.timestamps = 0
+        self.trajectories = {}
+        self.filterd_spikes = []
+
+    def calibrate_motion(self, spikes, calibration_time=None):
+
+        if calibration_time is None:
+            calibration_time = self.calibration_time
+        else:
+            self.calibration_time = calibration_time
+
+        print('begin calibrate..')
+        for t in range(calibration_time):
+            input_spk = torch.from_numpy(spikes[t, :, :]).to(self.device)
+            self.stp_filter.update_dynamics(t, input_spk)
+            self.timestamps += 1
+
+    def get_results(self, spikes, res_filepath, mov_writer=None, save_video=False):
+
+        result_file = open(res_filepath, 'a+')
+
+        timestamps = spikes.shape[0]
+        total_time = 0
+        predict_kwargs = {'spike_h': self.spike_h, 'spike_w': self.spike_w, 'device': self.device}
+
+        for t in tqdm(range(timestamps), desc=f'Saving tracking results to {str(result_file)}'):
+            try:
+                input_spk = torch.from_numpy(spikes[t, :, :]).to(self.device)
+                self.stp_filter.update_dynamics(self.timestamps, input_spk)
+
+                self.stp_filter.local_connect(self.stp_filter.filter_spk)
+                # self.filterd_spikes[t, :, :] = self.stp_filter.lif_spk.cpu().detach().numpy()
+
+                self.object_detection.update_dnf(self.stp_filter.lif_spk)
+                attentionBox, attentionInput = self.object_detection.get_attention_location(self.stp_filter.lif_spk)
+                # attentionInput = attentionInput.to(self.device)
+                num_box = attentionBox.shape[0]
+                self.motion_estimator.stdp_tracking(self.stp_filter.lif_spk)
+
+                motion_id, motion_vector, _ = self.motion_estimator.local_wta(self.stp_filter.lif_spk, self.timestamps, visualize=True)
+                # gpu_tracker.track()  # run function between the code line where uses GPU
+
+                predict_fire, sw, bw = self.object_cluster.update_weight(attentionInput)
+
+                predict_object = self.object_cluster.detect_object(predict_fire, attentionBox, motion_id, motion_vector, **predict_kwargs)
+
+                # visualize_weights(sw, 'before update tracks', t)
+
+                sw, bw = self.object_cluster.update_tracks(predict_object, sw, bw, self.timestamps)
+
+                self.object_cluster.synaptic_weight = sw.detach().clone()
+                self.object_cluster.bias_weight = bw.detach().clone()
+
+                dets = torch.zeros((num_box, 6), dtype=torch.int)
+                for i_box, bbox in enumerate(attentionBox):
+                    dets[i_box, :] = torch.tensor([bbox[0], bbox[1], bbox[2], bbox[3], 1, 1])
+
+                track_ids = []
+                if save_video:
+                    track_frame = self.stp_filter.lif_spk.cpu().numpy()
+                    track_frame = (track_frame * 255).astype(np.uint8)
+                    # track_frame = np.transpose(track_frame, (1, 2, 0))
+
+                    # track_frame = np.tile(track_frame, (3, 1, 1))
+                    # track_frame = np.squeeze(track_frame)
+                    track_frame = cv2.cvtColor(track_frame, cv2.COLOR_GRAY2BGR)
+
+                    for i_box in range(attentionBox.shape[0]):
+                        tmp_box = attentionBox[i_box, :]
+                        cv2.rectangle(track_frame, (int(tmp_box[1]), int(tmp_box[0])), (int(tmp_box[3]), int(tmp_box[2])),
+                                           (int(0), int(0), int(255)), 2)
+
+                for i_box in range(self.object_cluster.K2):
+                    if self.object_cluster.tracks[i_box].visible == 1:
+                        tmp_box = self.object_cluster.tracks[i_box].bbox.numpy()
+                        pred_box = self.object_cluster.tracks[i_box].predbox.numpy()
+                        id = self.object_cluster.tracks[i_box].id
+                        color = self.object_cluster.tracks[i_box].color
+
+                        # update the trajectories
+                        mid_y = (tmp_box[0, 0] + tmp_box[0, 2]) / 2  # height
+                        mid_x = (tmp_box[0, 1] + tmp_box[0, 3]) / 2  # width
+                        box_w = int(tmp_box[0, 3] - tmp_box[0, 1])
+                        box_h = int(tmp_box[0,2] - tmp_box[0, 0])
+                        print('%d,%d,%.2f,%.2f,%.2f,%.2f,1,-1,-1,-1' % (
+                            self.timestamps, id, tmp_box[0, 1], tmp_box[0, 0], box_w, box_h), file=result_file)
+
+                        if id not in self.trajectories:
+                            self.trajectories[id] = trajectories(int(id), [], [], [], 255 * np.random.rand(1, 3))
+                            self.trajectories[id].x.append(mid_x)
+                            self.trajectories[id].y.append(mid_y)
+                            self.trajectories[id].t.append(self.timestamps)
+
+                        else:
+                            self.trajectories[id].x.append(mid_x)
+                            self.trajectories[id].y.append(mid_y)
+                            self.trajectories[id].t.append(self.timestamps)
+                            # the detection results
+
+                        if save_video:
+                            cv2.rectangle(track_frame, (int(tmp_box[0, 1]), int(tmp_box[0, 0])),
+                                          (int(tmp_box[0, 3]), int(tmp_box[0, 2])),
+                                          (int(color[0, 0]), int(color[0, 1]), int(color[0, 2])), 2)
+
+                            # # the predicted results
+                            # cv2.rectangle(track_frame, (int(pred_box[0, 1]), int(pred_box[0, 0])),
+                            #               (int(pred_box[0, 3]), int(pred_box[0, 2])), (int(0), int(0), int(255)), 2)
+
+                            # the label box
+                            cv2.rectangle(track_frame, (int(tmp_box[0, 1]), int(tmp_box[0, 0] - 35)),
+                                          (int(tmp_box[0, 1] + 60), int(tmp_box[0, 0])),
+                                          (int(color[0, 0]), int(color[0, 1]), int(color[0, 2])), -1)
+                            if self.object_cluster.tracks[i_box].unvisible_count > 0:
+                                show_text = 'predict' + str(id)
+                            else:
+                                show_text = 'object' + str(id)
+                            cv2.putText(track_frame, show_text, (int(tmp_box[0, 1]), int(tmp_box[0, 0] - 10)),
+                                        cv2.FONT_HERSHEY_SIMPLEX,
+                                        1, (255, 255, 255), 2)
+
+                if save_video:
+                    cv2.putText(track_frame, str(int(self.timestamps)),
+                                (10, 70), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 255), 2)
+                    mov_writer.write(track_frame)
+                self.timestamps += 1
+
+            except RuntimeError as exception:
+                if "out of memory" in str(exception):
+                    print('WARNING: out of memory')
+                    if hasattr(torch.cuda, 'empty_cache'):
+                        torch.cuda.empty_cache()
+                    else:
+                        raise exception
+
+        print('Total tracking took: %.3f seconds for %d timestamps spikes' %
+              (total_time, self.timestamps - self.calibration_time))
+
+        # if save_video:
+        #     mov_writer.release()
+        #     cv2.destroyAllWindows()
+
+        result_file.close()
+
+    def save_trajectory(self, results_dir, data_name):
+        trajectories_filename = os.path.join(results_dir, data_name + '_py.json')
+        mat_trajectories_filename = 'results/' + data_name + '.json'
+        track_box_filename = 'results/' + data_name + '_bbox.json'
+
+        if os.path.exists(trajectories_filename):
+            os.remove(trajectories_filename)
+
+        if os.path.exists(mat_trajectories_filename):
+            os.remove(mat_trajectories_filename)
+
+        if os.path.exists(track_box_filename):
+            os.remove(track_box_filename)
+
+        for i_traj in range(self.object_cluster.K2):
+            tmp_traj = self.object_cluster.trajectories[i_traj]
+            tmp_bbox = self.object_cluster.tracks_bbox[i_traj]
+
+            traj_json_string = json.dumps(tmp_traj._asdict(), cls=NumpyEncoder)
+            bbox_json_string = json.dumps(tmp_bbox._asdict(), cls=NumpyEncoder)
+
+            with open(mat_trajectories_filename, 'a+') as f:
+                f.write(traj_json_string)
+
+            with open(track_box_filename, 'a+') as f:
+                f.write(bbox_json_string)
+
+        num_len = len(self.trajectories)
+        for i_traj in self.trajectories:
+            traj_json_string = json.dumps(self.trajectories[i_traj]._asdict(), cls=NumpyEncoder)
+
+            with open(trajectories_filename, 'a+') as f:
+                f.write(traj_json_string)
+
+                f.write('\n')

snnTracker/test_motion_detection.py

@@ -0,0 +1,201 @@
+import numpy as np
+import torch
+from spkProc.tracking.snn_tracker import SNNTracker
+import matplotlib.pyplot as plt
+import cv2
+
+
+def load_vidar_dat(filename, frame_cnt=None, width=640, height=480, reverse_spike=True):
+    '''
+    output: <class 'numpy.ndarray'> (frame_cnt, height, width) {0，1} float32
+    '''
+    array = np.fromfile(filename, dtype=np.uint8)
+
+    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_ = frame_.reshape((height, width), order='C')
+        if reverse_spike:
+            frame_ = np.flipud(frame_)
+        spikes.append(frame_)
+
+    return np.array(spikes).astype(np.float32)
+
+
+
+def detect_motion(spikes, calibration_frames=200, device=None):
+    """
+    使用SNN进行运动目标检测
+    Args:
+        spikes: shape为[frames, height, width]的脉冲数据
+        calibration_frames: 用于校准的帧数
+        device: 运行设备（CPU/GPU）
+    Returns:
+        motion_mask: 第calibration_frames帧的运动目标掩码
+    """
+    if device is None:
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    
+    spike_h, spike_w = spikes.shape[1:]
+    
+    # 初始化SNN跟踪器
+    spike_tracker = SNNTracker(spike_h, spike_w, device, attention_size=15)
+    
+    # 使用前calibration_frames帧进行校准
+    calibration_spikes = spikes[:calibration_frames]
+    spike_tracker.calibrate_motion(calibration_spikes, calibration_frames)
+    
+    # 获取第calibration_frames帧的运动检测结果
+    target_frame = spikes[calibration_frames]
+    target_frame = torch.from_numpy(target_frame).to(device)
+    # target_frame = target_frame.reshape(1, 1, spike_h, spike_w)
+    
+    # 获取运动检测结果
+    motion_id, motion_vector, _ = spike_tracker.motion_estimator.local_wta(target_frame, calibration_frames)
+    
+    # 生成运动掩码
+    motion_mask = (motion_id > 0).cpu().numpy()
+    
+    return motion_mask
+
+def spikes_to_tfi(spk_seq):
+    n, h, w = spk_seq.shape
+    last_index = np.zeros((1, h, w))
+    cur_index = np.zeros((1, h, w))
+    c_frames = np.zeros_like(spk_seq).astype(np.float64)
+    for i in range(n - 1):
+        last_index = cur_index
+        cur_index = spk_seq[i+1,:,:] * (i + 1) + (1 - spk_seq[i+1,:,:]) * last_index
+        c_frames[i,:,:] = cur_index - last_index
+    last_frame = c_frames[n-1:,:]
+    last_frame[last_frame==0] = n
+    c_frames[n-1,:,:] = last_frame
+    last_interval = n * np.ones((1, h, w))
+    for i in range(n - 2, -1, -1):
+        last_interval = spk_seq[i+1,:,:] * c_frames[i,:,:] + (1 - spk_seq[i+1,:,:]) * last_interval
+        tmp_frame = np.expand_dims(c_frames[i,:,:], 0)
+        tmp_frame[tmp_frame==0] = last_interval[tmp_frame==0]
+        c_frames[i] = tmp_frame
+    return 1.0 / c_frames
+
+def detect_object(spikes, calibration_frames=200, device=None):
+    """
+    使用SNN进行目标检测
+    Args:
+        spikes: shape为[frames, height, width]的脉冲数据
+        calibration_frames: 用于校准的帧数
+        device: 运行设备（CPU/GPU）
+    Returns:
+        object_mask: 第calibration_frames帧的目标掩码
+    """
+    if device is None:
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    
+    spike_h, spike_w = spikes.shape[1:]
+    
+    # 初始化SNN跟踪器
+    spike_tracker = SNNTracker(spike_h, spike_w, device, attention_size=15)
+    spike_tracker.object_cluster.K2 = 4
+
+    # 使用前calibration_frames帧进行校准
+    calibration_spikes = spikes[:calibration_frames]
+    spike_tracker.calibrate_motion(calibration_spikes, calibration_frames)
+    
+    # 获取第calibration_frames帧的目标检测结果
+    target_frame = spikes[calibration_frames: calibration_frames + 200]
+    print(target_frame.shape)
+    # target_frame = target_frame.reshape(1, 1, spike_h, spike_w)
+    
+    # 获取目标检测结果
+    save_filename = "testtest.avi"
+    mov = cv2.VideoWriter(save_filename, cv2.VideoWriter_fourcc(*'MJPG'), 30, (400, 250))
+    spike_tracker.get_results(target_frame, save_filename, mov, save_video=True)
+    
+    mov.release()
+    cv2.destroyAllWindows()
+    return 0
+
+
+if __name__ == "__main__":
+    height = 250
+    width = 400  
+    spikes = load_vidar_dat("0.dat", width=width, height=height)
+    for n in range(1,10):
+        tmp_spikes = load_vidar_dat(f"{n}.dat", width=width, height=height)
+        spikes = np.concatenate((spikes, tmp_spikes), axis=0)
+    print(spikes.shape)
+
+    spikes = spikes[::10]
+
+    motion_mask = detect_object(spikes, calibration_frames=200)
+
+    
+
+    tfi = spikes_to_tfi(spikes)
+    # 保存重建的视频
+    save_recon_filename = "tfi.avi"
+    recon_mov = cv2.VideoWriter(save_recon_filename, cv2.VideoWriter_fourcc(*'MJPG'), 30, (width, height))
+    
+    for frame in tfi:
+        frame_norm = (frame * 255).astype(np.uint8)
+        frame_rgb = cv2.cvtColor(frame_norm, cv2.COLOR_GRAY2BGR)
+        recon_mov.write(frame_rgb)
+    
+    recon_mov.release()
+
+
+
+    # 检测运动目标
+    # motion_mask = detect_motion(spikes, calibration_frames=200)
+    # print(f"Motion mask shape: {motion_mask.shape}")
+    # print(f"Number of motion pixels: {motion_mask.sum()}") 
+
+    # 可视化运动目标检测结果
+    # plt.figure(figsize=(10, 5))
+    # plt.subplot(1, 2, 1)
+    # plt.imshow(spikes[200], cmap='gray')
+    # plt.title("Input frame")
+    # plt.axis('off')
+    # plt.subplot(1, 2, 2)
+    # plt.imshow(motion_mask, cmap='gray')
+    # plt.title("Motion mask")
+    # plt.axis('off')
+    # plt.show()
+
+
+    # 计算原始脉冲图和运动掩码之间的差异
+    # spike_frame = spikes[200]  # 获取第200帧脉冲图
+    
+    # # 计算差异指标
+    # pixel_diff = np.logical_xor(spike_frame > 0, motion_mask).sum()
+    # total_pixels = height * width
+    # diff_ratio = pixel_diff / total_pixels
+    
+    # print("\n运动检测结果分析:")
+    # print(f"原始脉冲图中的活跃像素数: {(spike_frame > 0).sum()}")
+    # print(f"运动掩码中的运动像素数: {motion_mask.sum()}")
+    # print(f"不一致的像素数: {pixel_diff}")
+    # print(f"像素差异比例: {diff_ratio:.2%}")
+    
+    # # 可视化差异
+    # plt.figure(figsize=(10, 5))
+    # plt.subplot(1, 2, 1)
+    # plt.imshow(np.logical_xor(spike_frame > 0, motion_mask), cmap='gray')
+    # plt.title("Difference map (white indicates inconsistency)")
+    # plt.axis('off')
+    
+    # plt.subplot(1, 2, 2)
+    # plt.imshow(spike_frame > 0, cmap='gray', alpha=0.5)
+    # plt.imshow(motion_mask, cmap='Reds', alpha=0.5)
+    # plt.title("Overlay (Red: Motion mask, Gray: Original spikes)")
+    # plt.axis('off')
+    # plt.show()

snnTracker/test_snntracker copy.py

@@ -0,0 +1,77 @@
+# -*- coding: utf-8 -*-
+# @Time : 2024/12/05 20:17
+# @Author : Yajing Zheng
+# @Email: [email protected]
+# @File : test_snntracker.py
+import os, sys
+sys.path.append("..")
+import path
+
+import numpy as np
+from spkData.load_dat import data_parameter_dict, SpikeStream
+from pprint import pprint
+import torch
+from spkProc.tracking.snn_tracker import SNNTracker
+from utils import vis_trajectory
+from visualization.get_video import obtain_mot_video
+import cv2
+# from tracking_mot import TrackingMetrics
+
+from visualization.get_video import obtain_detection_video
+
+# change the path to where you put the datasets
+test_scene = "0"
+# data_filename = 'motVidarReal2020/rotTrans'
+data_filename = test_scene
+label_type = 'tracking'
+para_dict = data_parameter_dict(data_filename, label_type)
+pprint(para_dict)
+vidarSpikes = SpikeStream(**para_dict)
+
+# block_len = 2000
+# spikes = vidarSpikes.get_block_spikes(begin_idx=0, block_len=block_len)
+spikes = vidarSpikes.get_spike_matrix()
+pprint(spikes.shape)
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+calibration_time = 150
+filename = path.split_path_into_pieces(data_filename)
+result_filename = filename[-1] + '_snn.txt'
+if not os.path.exists('results'):
+    os.makedirs('results')
+tracking_file = os.path.join('results', result_filename)
+if os.path.exists(tracking_file):
+    os.remove(tracking_file)
+
+# stp_params = {'filterThr': 0.12,  # filter threshold
+#               'voltageMin': -10,
+#               'lifThr': 3}
+spike_tracker = SNNTracker(para_dict.get('spike_h'), para_dict.get('spike_w'), device, attention_size=15)
+spike_tracker.object_cluster.K2 = 4
+# total_spikes = spikes
+
+# using stp filter to filter out static spikes
+spike_tracker.calibrate_motion(spikes, calibration_time)
+# start tracking
+track_videoName = tracking_file.replace('txt', 'avi')
+mov = cv2.VideoWriter(track_videoName, cv2.VideoWriter_fourcc(*'MJPG'), 30, (para_dict.get('spike_w'), para_dict.get('spike_h')))
+spike_tracker.get_results(spikes[calibration_time:], tracking_file, mov, save_video=True)
+
+data_name = test_scene
+trajectories_filename = os.path.join('results', data_name + '_py.json')
+visTraj_filename = os.path.join('results', data_name + '.png')
+
+spike_tracker.save_trajectory('results', data_name)
+vis_trajectory(trajectories_filename, visTraj_filename, **para_dict)
+# measure the multi-object tracking performance
+# metrics = TrackingMetrics(tracking_file, **para_dict)
+# metrics.get_results()
+#
+# block_len = total_spikes.shape[0]
+mov.release()
+cv2.destroyAllWindows()
+# # visualize the tracking results to a video
+# video_filename = os.path.join('results', filename[-1] + '_mot.avi')
+# obtain_mot_video(spike_tracker.filterd_spikes, video_filename, tracking_file, **para_dict)
+# obtain_detection_video(total_spikes, video_filename, tracking_file, evaluate_seq_len=evaluate_seq_len, **para_dict)

snnTracker/test_snntracker.py

@@ -0,0 +1,78 @@
+# -*- coding: utf-8 -*-
+# @Time : 2024/12/05 20:17
+# @Author : Yajing Zheng
+# @Email: [email protected]
+# @File : test_snntracker.py
+import os, sys
+sys.path.append("..")
+import path
+
+import numpy as np
+from spkData.load_dat import data_parameter_dict, SpikeStream
+from pprint import pprint
+import torch
+from spkProc.tracking.snn_tracker import SNNTracker
+from utils import vis_trajectory
+from visualization.get_video import obtain_mot_video
+import cv2
+from tracking_mot import TrackingMetrics
+
+from visualization.get_video import obtain_detection_video
+
+# change the path to where you put the datasets
+test_scene = ['spike59', 'rotTrans', 'cplCam', 'cpl1', 'badminton', 'ball']
+# data_filename = 'motVidarReal2020/rotTrans'
+scene_idx = 2
+data_filename = 'motVidarReal2020/' + test_scene[scene_idx]
+label_type = 'tracking'
+para_dict = data_parameter_dict(data_filename, label_type)
+pprint(para_dict)
+vidarSpikes = SpikeStream(**para_dict)
+
+# block_len = 2000
+# spikes = vidarSpikes.get_block_spikes(begin_idx=0, block_len=block_len)
+spikes = vidarSpikes.get_spike_matrix()
+pprint(spikes.shape)
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+calibration_time = 150
+filename = path.split_path_into_pieces(data_filename)
+result_filename = filename[-1] + '_snn.txt'
+if not os.path.exists('results'):
+    os.makedirs('results')
+tracking_file = os.path.join('results', result_filename)
+if os.path.exists(tracking_file):
+    os.remove(tracking_file)
+
+# stp_params = {'filterThr': 0.12,  # filter threshold
+#               'voltageMin': -10,
+#               'lifThr': 3}
+spike_tracker = SNNTracker(para_dict.get('spike_h'), para_dict.get('spike_w'), device, attention_size=15)
+spike_tracker.object_cluster.K2 = 4
+# total_spikes = spikes
+
+# using stp filter to filter out static spikes
+spike_tracker.calibrate_motion(spikes, calibration_time)
+# start tracking
+track_videoName = tracking_file.replace('txt', 'avi')
+mov = cv2.VideoWriter(track_videoName, cv2.VideoWriter_fourcc(*'MJPG'), 30, (para_dict.get('spike_w'), para_dict.get('spike_h')))
+spike_tracker.get_results(spikes[calibration_time:], tracking_file, mov, save_video=True)
+
+data_name = test_scene[scene_idx]
+trajectories_filename = os.path.join('results', data_name + '_py.json')
+visTraj_filename = os.path.join('results', data_name + '.png')
+
+spike_tracker.save_trajectory('results', data_name)
+vis_trajectory(trajectories_filename, visTraj_filename, **para_dict)
+# measure the multi-object tracking performance
+# metrics = TrackingMetrics(tracking_file, **para_dict)
+# metrics.get_results()
+#
+# block_len = total_spikes.shape[0]
+mov.release()
+cv2.destroyAllWindows()
+# # visualize the tracking results to a video
+# video_filename = os.path.join('results', filename[-1] + '_mot.avi')
+# obtain_mot_video(spike_tracker.filterd_spikes, video_filename, tracking_file, **para_dict)
+# obtain_detection_video(total_spikes, video_filename, tracking_file, evaluate_seq_len=evaluate_seq_len, **para_dict)

snnTracker/utils.py

@@ -0,0 +1,207 @@
+import sys
+import numpy as np
+import torch
+import threading
+import cv2
+import json
+
+# import matplotlib
+# matplotlib.use('TkAgg')
+import matplotlib.pyplot as plt
+from mpl_toolkits.mplot3d import Axes3D
+from matplotlib.pyplot import MultipleLocator
+
+
+class dataReader(threading.Thread):
+    def __init__(self, file_reader, device, q, is_dat=True, is_npy=False, filedir=None):
+        super(dataReader, self).__init__()
+        self.file_reader = file_reader
+        self.device = device
+        self.q = q
+        self.is_dat = is_dat
+        self.is_npy = is_npy
+        self.filedir = filedir
+        self.stream = torch.cuda.Stream()
+
+    def run(self):
+        with torch.cuda.stream(self.stream):
+            for t in range(tnum):
+                if self.is_dat:
+                    ibuffer = self.file_reader.read(int(ivs_w * ivs_h / 8))
+                    a = bin(int.from_bytes(ibuffer, byteorder=sys.byteorder))
+                    a = a[2:].zfill(ivs_w * ivs_h)
+
+                    a = list(a)
+                    a = np.array(a, dtype=np.byte)
+                    a = np.reshape(a, [ivs_h, ivs_w])
+                    if ivs_h == 600:
+                        a = np.flip(a, 0)
+                    if ivs_h == 250:
+                        a = np.flip(a, 1)
+                    input_spk = torch.from_numpy(a != 0).to(device)
+                elif self.is_npy:
+                    npy_filename = self.filedir + str(t + 442) + '.npy'
+                    tmp_data = np.load(npy_filename)
+                    superResolution_rate = tmp_data.shape[2]
+                    for i_data in range(superResolution_rate):
+                        tmp_spk = tmp_data[:, :, i_data]
+                        input_spk = torch.from_numpy(tmp_spk).to(device)
+                        self.q.put(input_spk)
+
+                else:
+                    # img_filename = self.filedir + str(t + 4200) + '.png'
+                    img_filename = self.filedir + 'spike_' + str(t + 1) + '.png'
+                    # print('reading %d frames' % (t+1))
+                    # print('reading %d frames' % (t+5000))
+                    a = cv2.imread(img_filename)
+                    a = cv2.cvtColor(a, cv2.COLOR_BGR2GRAY)
+                    a = a / 255
+                    a = np.array(a, dtype=np.byte)
+                    input_spk = torch.from_numpy(a != 0).to(device)
+
+                self.q.put(input_spk)
+
+
+# obtain 2D gaussian filter
+def get_kernel(filter_size, sigma):
+    assert (filter_size + 1) % 2 == 0, '2D filter size must be odd number!'
+    g = np.zeros((filter_size, filter_size), dtype=np.float32)
+    half_width = int((filter_size - 1) / 2)
+    # center location
+
+    xc = (filter_size + 1) / 2
+    yc = (filter_size + 1) / 2
+    for i in range(-half_width, half_width + 1, 1):
+        for j in range(-half_width, half_width + 1, 1):
+            x = int(xc + i)
+            y = int(yc + j)
+            g[y - 1, x - 1] = np.exp(- (i ** 2 + j ** 2) / 2 / sigma / sigma)
+
+    g = (g - g.min()) / (g.max() - g.min())
+    return g
+
+
+def get_transform_matrix(ori, speed):
+    ori_num = len(ori)
+    speed_num = len(speed)
+    transform_matrix = torch.zeros(ori_num * speed_num, 2, 3)
+    cnt = 0
+    for iOri in range(ori_num):
+        for iSpeed in range(speed_num):
+            transform_matrix[cnt, 0, 0] = 1
+            transform_matrix[cnt, 1, 1] = 1
+
+            transform_matrix[cnt, 0, 2] = - float(ori[iOri, 1] * speed[iSpeed] / ivs_w)
+            transform_matrix[cnt, 1, 2] = - float(ori[iOri, 0] * speed[iSpeed] / ivs_h)
+
+            cnt += 1
+
+    transform_matrix = transform_matrix.to(device)
+    return transform_matrix
+
+
+def get_transform_matrix_new(ori, speed, dvs_w, dvs_h, device):
+    ori_num = len(ori)
+    speed_num = len(speed)
+    transform_matrix = torch.zeros(ori_num * speed_num, 2, 3)
+    cnt = 0
+    for iOri in range(ori_num):
+        for iSpeed in range(speed_num):
+            transform_matrix[cnt, 0, 0] = 1
+            transform_matrix[cnt, 1, 1] = 1
+
+            transform_matrix[cnt, 0, 2] = - float(ori[iOri, 1] * speed[iSpeed] / dvs_w)
+            transform_matrix[cnt, 1, 2] = - float(ori[iOri, 0] * speed[iSpeed] / dvs_h)
+
+            cnt += 1
+
+    transform_matrix = transform_matrix.to(device)
+    return transform_matrix
+
+
+# monitor the inference process
+def visualize_img(gray_img, tag, curT):
+    gray_img = gray_img.float32()
+    img = torch.unsqueeze(gray_img, 0)
+    logger.add_image(tag, img, global_step=curT)
+
+
+def visualize_images(images, tag, curT):
+    if images.shape[0] < 1:
+        return
+    images = torch.squeeze(images)
+    img_num = images.shape[-1]
+    for iImg in range(img_num):
+        tmp_img = images[:, :, iImg]
+        tmp_img = torch.squeeze(tmp_img)
+        tmp_img = torch.unsqueeze(tmp_img, 0)
+        logger.add_image(tag + str(iImg), tmp_img, global_step=curT)
+
+
+def visualize_weights(weights, tag, curT):
+    if weights.shape[0] < 1:
+        return
+    weights = torch.squeeze(weights)
+    weights_num = weights.shape[0]
+    input_size = weights.shape[1]
+    stim_size = int(np.sqrt(input_size))
+    for iw in range(weights_num):
+        tmp_w = weights[iw, :]
+        tmp_w = torch.squeeze(tmp_w)
+        tmp_w = (tmp_w - torch.min(tmp_w)) / (torch.max(tmp_w) - torch.min(tmp_w))
+        tmp_w = torch.reshape(tmp_w, (stim_size, stim_size))
+        tmp_w = torch.unsqueeze(tmp_w, 0)
+        logger.add_image(tag + str(iw), tmp_w, global_step=curT)
+
+
+class NumpyEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, np.ndarray):
+            return obj.tolist()
+        return json.JSONEncoder.default(self, obj)
+
+
+def vis_trajectory(json_file, filename, **dataDict):
+    spike_h = dataDict.get('spike_h')
+    spike_w = dataDict.get('spike_w')
+    traj_dict = []
+    with open(json_file, 'r') as f:
+        for line in f.readlines():
+            traj_dict.append(json.loads(line))
+
+    num_traj = len(traj_dict)
+
+    fig = plt.figure(figsize=[10, 6])
+    ax = fig.add_subplot(111, projection='3d')
+    min_t = 1000
+    max_t = 0
+
+    for tmp_traj in traj_dict:
+        tmp_t = np.array(tmp_traj['t'])
+        if np.min(tmp_t) < min_t:
+            min_t = np.min(tmp_t)
+        if np.max(tmp_t) > max_t:
+            max_t = np.max(tmp_t)
+
+        tmp_x = spike_w - np.array(tmp_traj['x'])
+        tmp_y = np.array(tmp_traj['y'])
+        tmp_color = np.array(tmp_traj['color']) / 255.
+        ax.plot(tmp_t, tmp_x, tmp_y, color=tmp_color, linewidth=2, label='traj ' + str(tmp_traj['id']))
+
+    ax.legend(loc='best', bbox_to_anchor=(0.7, 0., 0.4, 0.8))
+    zoom = [2.2, 0.8, 0.5, 1]
+    ax.get_proj = lambda: np.dot(Axes3D.get_proj(ax), np.diag([zoom[0], zoom[1], zoom[2], zoom[3]]))
+    ax.set_xlim(min_t, max_t)
+    ax.set_ylim(0, spike_w)
+    ax.set_zlim(0, spike_h)
+
+    ax.set_xlabel('time', fontsize=15)
+    ax.set_ylabel('width', fontsize=15)
+    ax.set_zlabel('height', fontsize=15)
+
+    ax.view_init(elev=16, azim=135)
+    ax.yaxis.set_major_locator(MultipleLocator(100))
+    fig.subplots_adjust(top=1., bottom=0., left=0.2, right=1.)
+    # fig.tight_layout()
+    plt.show()
+    plt.savefig(filename, dpi=500, transparent=True)

snnTracker/visualization/get_image.py

@@ -0,0 +1,121 @@
+# -*- coding: utf-8 -*- 
+# @Time : 2023/8/20 16:06 
+# @Author : Yajing Zheng
+# @Email: [email protected]
+# @File : get_image.py
+import numpy as np
+import matplotlib
+matplotlib.use('Agg')
+import matplotlib.pyplot as plt
+import json
+from mpl_toolkits.mplot3d import Axes3D
+from matplotlib.pyplot import MultipleLocator
+from matplotlib.patches import Rectangle
+import torch
+import copy
+
+def get_spike_raster(data):
+    num_neuron, timesteps = data.shape
+    colors = [f'C{i}' for i in range(num_neuron)]
+    # set different line properties for each set of positions
+    # note that some overlap
+    lineoffsets1 = np.array(range(1, num_neuron*2+1, 2))
+    linelengthts1 = np.ones((num_neuron, )) * 1.5
+
+    plt.figure(figsize=(8, 6))
+    plt.eventplot(data, colors=colors, lineoffsets=lineoffsets1, linelengths=linelengthts1)
+    return plt.gcf()
+
+
+def get_heatmap_handle(data, marker=None, bounding_box=None):
+
+    if torch.is_tensor(data):
+        data = copy.deepcopy(data.cpu().detach().numpy())
+
+    fig, ax = plt.subplots(figsize=(8, 6))
+    h, w = data.shape
+    if marker is not None:
+        num_points = marker.shape[1]
+        colors = [f'C{i}' for i in range(num_points)]
+        for i_point in range(num_points):
+            ax.plot(marker[1, i_point], h-marker[0, i_point], 'o', color=colors[i_point], markersize=10)
+            ax.annotate('P{}'.format(i_point), (marker[1, i_point], h-marker[0, i_point]))
+
+    if bounding_box is not None:
+        for i_box, bbox in enumerate(bounding_box):
+            ax.add_patch(Rectangle((bbox[1], bbox[0]), bbox[3]-bbox[1], bbox[2] - bbox[0],
+                                   edgecolor='red', facecolor='none', lw=2))
+
+    ax.imshow(data, cmap='Blues', interpolation='nearest')
+
+    # plt.colorbar()
+    plt.axis('off')  # 可选，关闭坐标轴
+    plt.title('Heatmap')
+
+    return plt.gcf()
+
+
+def get_histogram_handle(data, marker=None, bounding_box=None):
+
+    if torch.is_tensor(data):
+        data = copy.deepcopy(data.cpu().detach().numpy())
+
+    fig, ax = plt.subplots(figsize=(8, 6))
+    h, w = data.shape
+    ax.hist(data.reshape((-1, 1)), bins=20)
+
+    # plt.colorbar()
+    # plt.axis('off')  # 可选，关闭坐标轴
+    plt.title('Heatmap')
+
+    return plt.gcf()
+def vis_trajectory(box_file, json_file, filename, **dataDict):
+
+    spike_h = dataDict.get('spike_h')
+    spike_w = dataDict.get('spike_w')
+    traj_dict = []
+    with open(json_file, 'r') as f:
+        for line in f.readlines():
+            traj_dict.append(json.loads(line))
+
+    box_file = open(box_file, 'r')
+    result_lines = box_file.readlines()
+    num_traj = len(traj_dict)
+
+    fig = plt.figure(figsize=[10, 6])
+    ax = fig.add_subplot(111, projection='3d')
+    min_t = 1000
+    max_t = 0
+
+    for tmp_traj in traj_dict:
+        tmp_t = np.array(tmp_traj['t'])
+        if np.min(tmp_t) < min_t:
+            min_t = np.min(tmp_t)
+        if np.max(tmp_t) > max_t:
+            max_t = np.max(tmp_t)
+
+        tmp_x = spike_w - np.array(tmp_traj['x'])
+        tmp_y = np.array(tmp_traj['y'])
+        tmp_color = np.array(tmp_traj['color']) / 255.
+        ax.plot(tmp_t, tmp_x, tmp_y, color=tmp_color, linewidth=2, label='traj ' + str(tmp_traj['id']))
+
+    ax.legend(loc='best', bbox_to_anchor=(0.7, 0., 0.4, 0.8))
+    zoom = [2.2, 0.8, 0.5, 1]
+    ax.get_proj = lambda: np.dot(Axes3D.get_proj(ax), np.diag([zoom[0], zoom[1], zoom[2], zoom[3]]))
+    ax.set_xlim(min_t, max_t)
+    ax.set_ylim(0, spike_w)
+    ax.set_zlim(0, spike_h)
+
+    ax.set_xlabel('time', fontsize=15)
+    ax.set_ylabel('width', fontsize=15)
+    ax.set_zlabel('height', fontsize=15)
+
+    ax.view_init(elev=16, azim=135)
+    # ax.view_init(elev=2, azim=27)
+    ax.yaxis.set_major_locator(MultipleLocator(100))
+    fig.subplots_adjust(top=1., bottom=0., left=0.2, right=1.)
+    # fig.tight_layout()
+    # plt.savefig(filename, dpi=500, transparent=True)
+    # filename = filename.replace('png', 'eps')
+    # plt.savefig(filename, dpi=500, transparent=True)
+    plt.show()

snnTracker/visualization/get_video.py

@@ -0,0 +1,245 @@
+# -*- coding: utf-8 -*- 
+# @Time : 2022/6/12 15:21 
+# @Author : Yajing Zheng
+# @File : visualize.py
+import cv2
+import numpy as np
+import matplotlib
+matplotlib.use('Agg')
+import matplotlib.pyplot as plt
+
+import matplotlib.cm as cm
+import matplotlib.animation as animation
+
+def obtain_spike_video(spikes, video_filename, **dataDict):
+    spike_h = dataDict.get('spike_h')
+    spike_w = dataDict.get('spike_w')
+    timestamps = spikes.shape[0]
+
+    mov = cv2.VideoWriter(video_filename, cv2.VideoWriter_fourcc(*'MJPG'), 30, (spike_w, spike_h))
+
+    for iSpk in range(timestamps):
+        tmpSpk = spikes[iSpk, :, :] * 255
+        tmpSpk = cv2.cvtColor(tmpSpk.astype(np.uint8), cv2.COLOR_GRAY2BGR)
+        mov.write(tmpSpk)
+
+    mov.release()
+
+
+def obtain_reconstruction_video(images, video_filename, **dataDict):
+    spike_h = dataDict.get('spike_h')
+    spike_w = dataDict.get('spike_w')
+
+    img_num = images.shape[0]
+    mov = cv2.VideoWriter(video_filename, cv2.VideoWriter_fourcc(*'MJPG'), 30, (spike_w, spike_h))
+    for iImg in range(img_num):
+        tmp_img = images[iImg, :, :]
+        tmp_img = cv2.cvtColor(tmp_img, cv2.COLOR_GRAY2BGR)
+        mov.write(tmp_img)
+
+    mov.release()
+
+
+def obtain_mot_video(spikes, video_filename, res_filepath, **dataDict):
+    spike_h = dataDict.get('spike_h')
+    spike_w = dataDict.get('spike_w')
+
+    gt_file = dataDict.get('labeled_data_dir')
+    gt_boxes = {}
+    if gt_file is not None:
+        gt_f = open(gt_file, 'r')
+        gt_lines = gt_f.readlines()
+        for line in gt_lines:
+            gt_term = line.split(',')
+            time_step = gt_term[0]
+            box_id = gt_term[1]
+            x = float(gt_term[2])
+            y = float(gt_term[3])
+            w = float(gt_term[4])
+            h = float(gt_term[5])
+
+            if str(time_step) not in gt_boxes:
+                gt_boxes[str(time_step)] = []
+            bbox = [box_id, x, y, w, h]
+            gt_boxes[str(time_step)].append(bbox)
+
+        gt_f.close()
+
+    result_file = res_filepath
+    test_boxes = {}
+    result_f = open(result_file, 'r')
+    result_lines = result_f.readlines()
+    color_dict = {}
+
+    for line in result_lines:
+        res_box = line.split(',')
+        time_step = res_box[0]
+        track_id = res_box[1]
+        if track_id not in color_dict.keys():
+            colors = (np.random.rand(1, 3) * 255).astype(np.uint8)
+            color_dict[track_id] = np.squeeze(colors)
+
+        x = float(res_box[2])
+        y = float(res_box[3])
+        w = float(res_box[4])
+        h = float(res_box[5])
+
+        if str(time_step) not in test_boxes:
+            test_boxes[str(time_step)] = []
+
+        test_box = [track_id, x, y, w, h]
+        test_boxes[str(time_step)].append(test_box)
+
+    result_f.close()
+
+    mov = cv2.VideoWriter(video_filename, cv2.VideoWriter_fourcc(*'MJPG'), 30, (spike_w, spike_h))
+
+    timestamps = spikes.shape[0]
+    for t in range(151, timestamps):
+    # for t in range(160, 1000):
+        tmp_ivs = spikes[t, :, :] * 255
+        tmp_ivs = cv2.cvtColor(tmp_ivs.astype(np.uint8), cv2.COLOR_GRAY2BGR)
+
+        if len(gt_boxes) > 0:
+            if str(t) in gt_boxes:
+                gts = gt_boxes[str(t)]
+                gt_num = len(gts)
+                for i in range(gt_num):
+                    box = gts[i]
+                    box_id = box[0]
+                    cv2.rectangle(tmp_ivs, (int(box[2]), int(box[1])),
+                                  (int(box[2] + box[4]), int(box[1] + box[3])),
+                                  (int(255), int(255), int(255)), 2)
+
+        if str(t) in test_boxes:
+            test = test_boxes[str(t)]
+            test_num = len(test)
+            for i in range(test_num):
+                box = test[i]
+                box_id = box[0]
+                colors = color_dict[box_id]
+                cv2.rectangle(tmp_ivs, (int(box[2]), int(box[1])),
+                              (int(box[2] + box[4]), int(box[1] + box[3])),
+                              (int(colors[0]), int(colors[1]), int(colors[2])), 2)
+
+        mov.write(tmp_ivs)
+
+    mov.release()
+
+
+def obtain_detection_video(spikes, video_filename, res_filepath, evaluate_seq_len, begin_idx=0, **dataDict):
+    spike_h = dataDict.get('spike_h')
+    spike_w = dataDict.get('spike_w')
+
+    gt_file = dataDict.get('labeled_data_dir')
+    gt_boxes = {}
+    if gt_file is not None:
+        start_idx = begin_idx
+        end_idx = begin_idx + evaluate_seq_len
+        for seq_no in range(start_idx, end_idx):
+            gt_filename = gt_file[seq_no]
+            gt_f = open(gt_filename, 'r')
+
+            gt_lines = gt_f.readlines()
+            for line in gt_lines:
+                tmp_box = line.split(',')
+
+                x = float(tmp_box[0])
+                y = float(tmp_box[1])
+                w = float(tmp_box[2])
+                h = float(tmp_box[3])
+                box_id = int(0)
+
+                if str(seq_no) not in gt_boxes:
+                    gt_boxes[str(seq_no)] = []
+                bbox = [box_id, x, y, w, h]
+                gt_boxes[str(seq_no)].append(bbox)
+
+            gt_f.close()
+
+    result_file = res_filepath
+    test_boxes = {}
+    result_f = open(result_file, 'r')
+    result_lines = result_f.readlines()
+    color_dict = {}
+
+    for line in result_lines:
+        res_box = line.split(',')
+        time_step = res_box[0]
+        track_id = res_box[1]
+        if track_id not in color_dict.keys():
+            colors = (np.random.rand(1, 3) * 255).astype(np.uint8)
+            color_dict[track_id] = np.squeeze(colors)
+
+        x = float(res_box[2])
+        y = float(res_box[3])
+        w = float(res_box[4])
+        h = float(res_box[5])
+
+        if str(time_step) not in test_boxes:
+            test_boxes[str(time_step)] = []
+
+        test_box = [track_id, x, y, w, h]
+        test_boxes[str(time_step)].append(test_box)
+
+    result_f.close()
+
+    mov = cv2.VideoWriter(video_filename, cv2.VideoWriter_fourcc(*'MJPG'), 30, (spike_w, spike_h))
+
+    block_len = spikes.shape[0]
+    # gt_intv = int(block_len/evaluate_seq_len)
+    gt_intv = 400
+
+    # for t in range(150, block_len):
+    for i_gt in range(start_idx+1, end_idx):
+        t = i_gt * gt_intv + int(gt_intv/2)
+        tmp_ivs = spikes[t, :, :] * 255
+        tmp_ivs = cv2.cvtColor(tmp_ivs.astype(np.uint8), cv2.COLOR_GRAY2BGR)
+
+        if len(gt_boxes) > 0:
+            gts = gt_boxes[str(i_gt)]
+            gt_num = len(gts)
+            for i in range(gt_num):
+                box = gts[i]
+                cv2.rectangle(tmp_ivs, (int(spike_w - box[1]), int(box[2])),
+                              (int(spike_w - box[1] - box[3]), int(box[2] + box[4])),
+                              (int(255), int(255), int(255)), 2)
+
+        if str(t) in test_boxes:
+            test = test_boxes[str(t)]
+            test_num = len(test)
+            for i in range(test_num):
+                box = test[i]
+                box_id = box[0]
+                colors = color_dict[box_id]
+                cv2.rectangle(tmp_ivs, (int(box[2]), int(box[1])),
+                              (int(box[2] + box[4]), int(box[1] + box[3])),
+                              (int(colors[0]), int(colors[1]), int(colors[2])), 2)
+
+        mov.write(tmp_ivs)
+
+    mov.release()
+
+def get_heatVideo(results, video_filename):
+    results = np.array(results)
+    frame_num = results.shape[0]
+    frames = []
+
+    fig = plt.figure()
+    for i in range(frame_num):
+        tmp_res = results[i]
+        # frames.append([plt.imshow(tmp_res, cmap=cm.Greys_r, animated=True)])
+        frames.append([plt.imshow(tmp_res, cmap=cm.Blues, animated=True)])
+
+    ani = animation.ArtistAnimation(fig, frames, interval=50, blit=True,
+                                    repeat_delay=1000)
+
+    # change the path to where you save ffmpeg
+    plt.rcParams['animation.ffmpeg_path'] = 'F:\\ffmpeg-N-99818-g993429cfb4-win64-gpl-shared-vulkan\\bin\\ffmpeg.exe'
+    FFwrite = animation.FFMpegWriter(fps=30, extra_args=['-vcodec', 'libx264'])
+    ani.save(video_filename, writer=FFwrite)
+    plt.show()
+
+
+
+

snnTracker/visualization/optical_flow_visualization.py

@@ -0,0 +1,272 @@
+# -*- coding: utf-8 -*- 
+# @Time : 2022/7/21
+# @Author : Rui Zhao
+# @File : optical_flow_visualization.py
+
+import torch
+import numpy as np
+import cv2
+import math
+
+
+#################### Interface ####################
+def flow_visualization(flow, mode='normal', use_cv2=True):
+    if mode == 'normal':
+        flow_vis = flow_to_image(flow_uv=flow, convert_to_bgr=use_cv2)
+    elif mode == 'scflow':
+        flow_vis = flow_to_img_scflow(flow_uv=flow)
+        if not use_cv2:
+            flow_vis = cv2.cvtColor(flow_vis, cv2.COLOR_BGR2RGB)
+    elif mode == 'evflow':
+        flow_vis = flow_viz_np(flow_x=flow[:,:,0], flow_y=flow[:,:,1])
+    
+    return flow_vis
+
+
+def vis_color_map(use_cv2=True):
+    u = np.linspace(-100, 99, 200)
+    v = np.linspace(-100, 99, 200)
+    xx, yy = np.meshgrid(u, v)
+    flow = np.concatenate((xx[:,:,None], yy[:,:,None]), axis=2)
+    map_normal = flow_visualization(flow=flow, mode='normal', use_cv2=use_cv2)
+    map_scflow = flow_visualization(flow=flow, mode='scflow', use_cv2=use_cv2)
+    map_evflow = flow_visualization(flow=flow, mode='evflow', use_cv2=use_cv2)
+    return [map_normal, map_scflow, map_evflow]
+
+
+def make_colorwheel():
+    """
+    Generates a color wheel for optical flow visualization as presented in:
+        Baker et al. "A Database and Evaluation Methodology for Optical Flow" (ICCV, 2007)
+        URL: http://vision.middlebury.edu/flow/flowEval-iccv07.pdf
+
+    Code follows the original C++ source code of Daniel Scharstein.
+    Code follows the the Matlab source code of Deqing Sun.
+
+    Returns:
+        np.ndarray: Color wheel
+    """
+
+    RY = 15
+    YG = 6
+    GC = 4
+    CB = 11
+    BM = 13
+    MR = 6
+
+    ncols = RY + YG + GC + CB + BM + MR
+    colorwheel = np.zeros((ncols, 3))
+    col = 0
+
+    # RY
+    colorwheel[0:RY, 0] = 255
+    colorwheel[0:RY, 1] = np.floor(255*np.arange(0,RY)/RY)
+    col = col+RY
+    # YG
+    colorwheel[col:col+YG, 0] = 255 - np.floor(255*np.arange(0,YG)/YG)
+    colorwheel[col:col+YG, 1] = 255
+    col = col+YG
+    # GC
+    colorwheel[col:col+GC, 1] = 255
+    colorwheel[col:col+GC, 2] = np.floor(255*np.arange(0,GC)/GC)
+    col = col+GC
+    # CB
+    colorwheel[col:col+CB, 1] = 255 - np.floor(255*np.arange(CB)/CB)
+    colorwheel[col:col+CB, 2] = 255
+    col = col+CB
+    # BM
+    colorwheel[col:col+BM, 2] = 255
+    colorwheel[col:col+BM, 0] = np.floor(255*np.arange(0,BM)/BM)
+    col = col+BM
+    # MR
+    colorwheel[col:col+MR, 2] = 255 - np.floor(255*np.arange(MR)/MR)
+    colorwheel[col:col+MR, 0] = 255
+    return colorwheel
+
+
+#################### Normal Version ####################
+"""
+From https://github.com/princeton-vl/RAFT/blob/master/core/utils/flow_viz.py
+"""
+def flow_uv_to_colors(u, v, convert_to_bgr=False):
+    """
+    Applies the flow color wheel to (possibly clipped) flow components u and v.
+    According to the C++ source code of Daniel Scharstein
+    According to the Matlab source code of Deqing Sun
+    Args:
+        u (np.ndarray): Input horizontal flow of shape [H,W]
+        v (np.ndarray): Input vertical flow of shape [H,W]
+        convert_to_bgr (bool, optional): Convert output image to BGR. Defaults to False.
+    Returns:
+        np.ndarray: Flow visualization image of shape [H,W,3]
+    """
+    flow_image = np.zeros((u.shape[0], u.shape[1], 3), np.uint8)
+    colorwheel = make_colorwheel()  # shape [55x3]
+    ncols = colorwheel.shape[0]
+    rad = np.sqrt(np.square(u) + np.square(v))
+    a = np.arctan2(-v, -u)/np.pi
+    fk = (a+1) / 2*(ncols-1)
+    k0 = np.floor(fk).astype(np.int32)
+    k1 = k0 + 1
+    k1[k1 == ncols] = 0
+    f = fk - k0
+    for i in range(colorwheel.shape[1]):
+        tmp = colorwheel[:,i]
+        col0 = tmp[k0] / 255.0
+        col1 = tmp[k1] / 255.0
+        col = (1-f)*col0 + f*col1
+        idx = (rad <= 1)
+        col[idx]  = 1 - rad[idx] * (1-col[idx])
+        col[~idx] = col[~idx] * 0.75   # out of range
+        # Note the 2-i => BGR instead of RGB
+        ch_idx = 2-i if convert_to_bgr else i
+        flow_image[:,:,ch_idx] = np.floor(255 * col)
+    return flow_image
+
+
+def flow_to_image(flow_uv, clip_flow=None, convert_to_bgr=False):
+    """
+    Expects a two dimensional flow image of shape.
+    Args:
+        flow_uv (np.ndarray): Flow UV image of shape [H,W,2]
+        clip_flow (float, optional): Clip maximum of flow values. Defaults to None.
+        convert_to_bgr (bool, optional): Convert output image to BGR. Defaults to False.
+    Returns:
+        np.ndarray: Flow visualization image of shape [H,W,3]
+    """
+    assert flow_uv.ndim == 3, 'input flow must have three dimensions'
+    assert flow_uv.shape[2] == 2, 'input flow must have shape [H,W,2]'
+    if clip_flow is not None:
+        flow_uv = np.clip(flow_uv, 0, clip_flow)
+    u = flow_uv[:,:,0]
+    v = flow_uv[:,:,1]
+    rad = np.sqrt(np.square(u) + np.square(v))
+    rad_max = np.max(rad)
+    epsilon = 1e-5
+    u = u / (rad_max + epsilon)
+    v = v / (rad_max + epsilon)
+    return flow_uv_to_colors(u, v, convert_to_bgr)
+
+
+#################### SCFlow Version ####################
+def flow_uv_to_colors_scflow(u, v, convert_to_bgr=False):
+    """
+    Applies the flow color wheel to (possibly clipped) flow components u and v.
+
+    According to the C++ source code of Daniel Scharstein
+    According to the Matlab source code of Deqing Sun
+
+    Args:
+        u (np.ndarray): Input horizontal flow of shape [H,W]
+        v (np.ndarray): Input vertical flow of shape [H,W]
+        convert_to_bgr (bool, optional): Convert output image to BGR. Defaults to False.
+
+    Returns:
+        np.ndarray: Flow visualization image of shape [H,W,3]
+    """
+    flow_image = np.zeros((u.shape[0], u.shape[1], 3), np.uint8)
+    colorwheel = make_colorwheel()  # shape [55x3]
+    ncols = colorwheel.shape[0]
+    rad = np.sqrt(np.square(u) + np.square(v))
+    a = np.arctan2(-v, u)/np.pi
+    fk = (a+1) / 2*(ncols-1)
+    k0 = np.floor(fk).astype(np.int32)
+    k1 = k0 + 1
+    k1[k1 == ncols] = 0
+    f = fk - k0
+    for i in range(colorwheel.shape[1]):
+        tmp = colorwheel[:,i]
+        col0 = tmp[k0] / 255.0
+        col1 = tmp[k1] / 255.0
+        col = (1-f)*col0 + f*col1
+        idx = (rad <= 1)
+        col[idx]  = 1 - rad[idx] * (1-col[idx])
+        col[~idx] = col[~idx] * 0.75   # out of range
+        # Note the 2-i => BGR instead of RGB
+        ch_idx = 2-i if convert_to_bgr else i
+        flow_image[:,:,ch_idx] = np.floor(255 * col)
+    return flow_image
+
+
+def flow_to_img_scflow(flow_uv, clip_flow=None):
+    """
+    Expects a two dimensional flow image of shape.
+
+    Args:
+        flow_uv (np.ndarray): Flow UV image of shape [H,W,2]
+        clip_flow (float, optional): Clip maximum of flow values. Defaults to None.
+
+    Returns:
+        np.ndarray: Flow visualization image of shape [H,W,3]
+    """
+    convert_to_bgr = False
+    assert flow_uv.ndim == 3, 'input flow must have three dimensions'
+    assert flow_uv.shape[2] == 2, 'input flow must have shape [H,W,2]'
+    if clip_flow is not None:
+        flow_uv = np.clip(flow_uv, 0, clip_flow)
+    u = flow_uv[:,:,0]
+    v = flow_uv[:,:,1]
+    rad = np.sqrt(np.square(u) + np.square(v))
+    rad_max = np.max(rad)
+    epsilon = 1e-5
+    u = u / (rad_max + epsilon)
+    v = v / (rad_max + epsilon)
+    return flow_uv_to_colors_scflow(u, v, convert_to_bgr)
+
+
+#################### EVFlow Version ####################
+"""
+From https://github.com/chan8972/Spike-FlowNet/blob/master/vis_utils.py
+"""
+
+"""
+Generates an RGB image where each point corresponds to flow in that direction from the center,
+as visualized by flow_viz_tf.
+Output: color_wheel_rgb: [1, width, height, 3]
+"""
+def draw_color_wheel_np(width, height):
+    color_wheel_x = np.linspace(-width / 2.,width / 2.,width)
+    color_wheel_y = np.linspace(-height / 2.,height / 2.,height)
+    color_wheel_X, color_wheel_Y = np.meshgrid(color_wheel_x, color_wheel_y)
+    color_wheel_rgb = flow_viz_np(color_wheel_X, color_wheel_Y)
+    return color_wheel_rgb
+
+
+"""
+Visualizes optical flow in HSV space using TensorFlow, with orientation as H, magnitude as V.
+Returned as RGB.
+Input: flow: [batch_size, width, height, 2]
+Output: flow_rgb: [batch_size, width, height, 3]
+"""
+def flow_viz_np(flow_x, flow_y):
+    import cv2
+    flows = np.stack((flow_x, flow_y), axis=2)
+    mag = np.linalg.norm(flows, axis=2)
+
+    ang = np.arctan2(flow_y, flow_x)
+    ang += np.pi
+    ang *= 180. / np.pi / 2.
+    ang = ang.astype(np.uint8)
+    hsv = np.zeros([flow_x.shape[0], flow_x.shape[1], 3], dtype=np.uint8)
+    hsv[:, :, 0] = ang
+    hsv[:, :, 1] = 255
+    hsv[:, :, 2] = cv2.normalize(mag, None, 0, 255, cv2.NORM_MINMAX)
+    flow_rgb = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)
+    return flow_rgb
+
+
+
+#################### Visualization tools when training SCFlow ####################
+def outflow_img(flow_list, vis_path, name_prefix='flow', max_batch=4):
+    flow = flow_list[0]
+    batch_size, c, h, w = flow.shape
+
+    for batch in range(batch_size):
+        if batch > max_batch:
+            break
+        flow_current = flow[batch,:,:,:].permute(1,2,0).detach().cpu().numpy()
+        flow_img = flow_visualization(flow_current, mode='scflow', use_cv2=True)
+
+        cv2.imwrite(vis_path + '/{:s}_batch_id={:02d}.png'.format(name_prefix, batch), flow_img)
+    
+    return