Upload 96 files

.gitattributes

@@ -33,3 +33,20 @@ 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
+assets/11_180_380.gif filter=lfs diff=lfs merge=lfs -text
+assets/3_143_243.gif filter=lfs diff=lfs merge=lfs -text
+assets/indoor_fruit_410_510.gif filter=lfs diff=lfs merge=lfs -text
+assets/indoor_fruit_guobao2_5_155.gif filter=lfs diff=lfs merge=lfs -text
+assets/indoor_hand_move_dynamic2_300_400.gif filter=lfs diff=lfs merge=lfs -text
+assets/outdoor_day2-1_317_417.gif filter=lfs diff=lfs merge=lfs -text
+assets/peanuts_light_160_386.gif filter=lfs diff=lfs merge=lfs -text
+assets/peanuts_running_2360_2460.gif filter=lfs diff=lfs merge=lfs -text
+assets/teaser.png filter=lfs diff=lfs merge=lfs -text
+assets/toulan_360_460.gif filter=lfs diff=lfs merge=lfs -text
+gt_tracks/boxes_rotation_198_278_gt_track.mp4 filter=lfs diff=lfs merge=lfs -text
+gt_tracks/boxes_translation_330_410_gt_track.mp4 filter=lfs diff=lfs merge=lfs -text
+gt_tracks/peanuts_light_160_386_gt_track_old.mp4 filter=lfs diff=lfs merge=lfs -text
+gt_tracks/peanuts_light_160_386.gt_track.mp4 filter=lfs diff=lfs merge=lfs -text
+gt_tracks/peanuts_running_2360_2460_gt_track.mp4 filter=lfs diff=lfs merge=lfs -text
+gt_tracks/rocket_earth_light_338_438_gt_track.mp4 filter=lfs diff=lfs merge=lfs -text
+gt_tracks/ziggy_in_the_arena_1350_1650_gt_track.mp4 filter=lfs diff=lfs merge=lfs -text

LFE_TAP/datasets/EC_dataset.py

@@ -0,0 +1,111 @@
+import os
+import torch
+import imageio
+import re
+import numpy as np
+from LFE_TAP.utils.event.utils import read_input
+from LFE_TAP.utils.dataset_utils import FrameEventData_test
+
+class EC_dataset(torch.utils.data.Dataset):
+    def __init__(self, data_root, dt=0.0200, representation="time_surfaces_v2_5", event_template_type = "sobel"):
+        self.data_root = data_root
+        self.dt = dt
+        self.representation = representation
+        self.event_template = event_template_type
+        self.seq_names = [
+            fname
+            for fname in os.listdir(data_root)
+            if os.path.isdir(os.path.join(data_root, fname))
+        ]
+        print("found %d unique seqences in %s" % (len(self.seq_names), self.data_root))
+
+    def __getitem__(self, index):
+        gotit =False
+        event_dir_path = os.path.join(str(self.data_root), self.seq_names[index], "events", f"{self.dt:.4f}", self.representation)
+        rgb_path = os.path.join(str(self.data_root), self.seq_names[index], "images_corrected")
+        track_point_path = os.path.join(str(self.data_root), self.seq_names[index], "track.gt.txt")
+
+        img_paths = sorted(os.listdir(rgb_path))
+        # img_paths = img_paths[::16]
+        event_paths = sorted(os.listdir(event_dir_path))
+        rgb_imgs = []
+        rgb_ifnew = []
+        rgb_times = []
+        rgb_imgs_plus = []
+        rgb_ind = 0
+        event_imgs = []
+        event_time = []
+        for i, img_path in enumerate(img_paths):
+            try:
+                rgb_imgs.append(imageio.v2.imread(os.path.join(rgb_path, img_path)))
+                rgb_times.append(int(re.match(r"\d+", img_path).group()))
+            except Exception as e:
+                print(f"error reading image at path:{img_path}")
+                print(f"error mrssage:{str(e)}")
+                gotit = False
+
+        for i, event_path in enumerate(event_paths):
+            try:
+                event_imgs.append(read_input(os.path.join(event_dir_path, event_path), self.representation))
+                rgb_time = rgb_times[rgb_ind] if rgb_ind < len(rgb_times) else float('inf')
+                event_time.append(int(re.match(r"\d+", event_path).group()))
+                if int(re.match(r"\d+", event_path).group()) >= rgb_time:
+                    # event_time.append(rgb_times[rgb_ind])
+                    rgb_imgs_plus.append(rgb_imgs[rgb_ind])
+                    rgb_ind += 1
+                    rgb_ifnew.append(1)
+                else:
+                    # event_time.append(int(re.match(r"\d+", event_path).group())-5000)
+                    rgb_imgs_plus.append(rgb_imgs[rgb_ind-1])
+                    rgb_ifnew.append(0)
+            except Exception as e:
+                print(f"error reading event at path:{event_path}")
+                print(f"error mrssage:{str(e)}")
+                gotit = False
+
+        # rgb_imgs = np.stack(rgb_imgs_plus)
+        # try:
+        #     event_imgs = np.stack(event_imgs)
+        # except Exception as e:
+        #     print(f"error reading at path:{str(self.data_root)}")
+        #     gotit = False
+        #     return [], gotit
+        rgb_imgs = np.stack(rgb_imgs_plus).transpose(0, 3, 1, 2)
+        event_imgs = np.stack(event_imgs).transpose(0, 3, 1, 2)
+        traj_data = np.genfromtxt(track_point_path, delimiter=" ")       # id, t, x, y
+        track_num, track_ind = np.unique(traj_data[:, 0], return_index=True)   
+        query_points = traj_data[track_ind, 2:]
+
+        T, H, W, C = event_imgs.shape
+        
+        segs = np.array(event_time)
+        rgb_timestamp = np.array(rgb_times)
+        img_ifnew = np.array(rgb_ifnew)
+        query_points = torch.from_numpy(query_points).float()
+        query_points = torch.cat([torch.zeros_like(query_points[:, :1]), query_points], dim=1)
+        gotit = True
+
+        sample = FrameEventData_test(
+            rgb_imgs,
+            event_imgs,
+            segs,
+            traj_data,
+            seq_name=self.seq_names[index],
+            query_points=query_points,
+            img_ifnew=img_ifnew,
+            rgb_timestamp=rgb_timestamp,
+        )
+
+        return sample, gotit
+    
+    def get_a_seq(self, seq_name):
+        for i in range(len(self.seq_names)):
+            if self.seq_names[i] == seq_name:
+                sample, gotit = self.__getitem__(i)
+                return sample, gotit
+        
+        print("WARNNING: did not find the sequence", seq_name)
+        return [], False
+
+    def __len__(self):
+        return len(self.seq_names)

LFE_TAP/datasets/EDS_dataset.py

@@ -0,0 +1,115 @@
+import os
+import torch
+import imageio
+import re
+import numpy as np
+import torch.nn.functional as F
+from LFE_TAP.utils.event.utils import read_input
+from LFE_TAP.utils.dataset_utils import FrameEventData, FrameEventData_test
+
+class EDS_dataset(torch.utils.data.Dataset):
+    def __init__(self, data_root, dt=0.0050, representation="time_surfaces_v2_5"):
+        self.data_root = data_root
+        self.dt = dt
+        self.representation = representation
+        self.seq_names = [
+            fname
+            for fname in os.listdir(data_root)
+            if os.path.isdir(os.path.join(data_root, fname))
+        ]
+        print("found %d unique seqences in %s" % (len(self.seq_names), self.data_root))
+
+    def __getitem__(self, index):
+        gotit =False
+        event_dir_path = os.path.join(str(self.data_root), self.seq_names[index], "events", f"{self.dt:.4f}", self.representation)
+        rgb_path = os.path.join(str(self.data_root), self.seq_names[index], "images_corrected")
+        track_point_path = os.path.join(f"gt_tracks/{self.seq_names[index]}.gt.txt")
+
+        img_paths = sorted(os.listdir(rgb_path))
+        # img_paths = img_paths[::8]
+        event_paths = sorted([f for f in os.listdir(event_dir_path) if f.endswith('.h5')])
+        rgb_imgs = []
+        rgb_ifnew = []
+        rgb_times = []
+        rgb_imgs_plus = []
+        rgb_ind = 0
+        event_imgs = []
+        event_time = []
+        time_emmbed = []
+        for i, img_path in enumerate(img_paths):
+            try:
+                # gray
+                img = imageio.v2.imread(os.path.join(rgb_path, img_path)).reshape(480, 640, 1)
+                img = img.repeat(3, axis=2)
+                # rgb
+                # img = imageio.v2.imread(os.path.join(rgb_path, img_path))
+                
+                rgb_imgs.append(img)
+                rgb_times.append(int(re.match(r"\d+", img_path).group()))
+            except Exception as e:
+                print(f"error reading image at path:{img_path}")
+                print(f"error mrssage:{str(e)}")
+                gotit = False
+        for i, event_path in enumerate(event_paths):
+            try:
+                event_imgs.append(read_input(os.path.join(event_dir_path, event_path), self.representation))
+                rgb_time = rgb_times[rgb_ind] if rgb_ind < len(rgb_times) else float('inf')
+                if int(event_path.split('.')[0]) >= rgb_time:
+                    # event_imgs.append(read_input(os.path.join(event_dir_path, event_path), self.representation))
+                    event_time.append(rgb_times[rgb_ind])
+                    rgb_imgs_plus.append(rgb_imgs[rgb_ind])
+                    time_emmbed.append(round((int(event_path.split('.')[0]) - rgb_times[rgb_ind])*1e-4, 1))
+                    rgb_ind += 1
+                    rgb_ifnew.append(1)
+                else:
+                    # continue
+                    event_time.append(int(event_path.split('.')[0]))
+                    rgb_imgs_plus.append(rgb_imgs[max(rgb_ind-1,0)])
+                    time_emmbed.append(round((int(event_path.split('.')[0]) - rgb_times[max(rgb_ind-1,0)])*1e-4, 1))
+                    rgb_ifnew.append(0)
+            except Exception as e:
+                print(f"error reading event at path:{event_path}")
+                print(f"error mrssage:{str(e)}")
+                gotit = False
+
+        rgb_imgs = np.stack(rgb_imgs_plus).transpose(0, 3, 1, 2)
+        event_imgs = np.stack(event_imgs)
+        if self.representation != "event_stack":
+            event_imgs = event_imgs.transpose(0, 3, 1, 2)
+        traj_data = np.genfromtxt(track_point_path, delimiter=" ")       # id, t, x, y
+        track_num, track_ind = np.unique(traj_data[:, 0], return_index=True)   
+        query_points = traj_data[track_ind, 2:]
+        
+        T, C, _, _ = event_imgs.shape
+
+        segs = np.array(event_time)
+        rgb_timestamp = np.array(rgb_times)
+        img_ifnew = np.array(rgb_ifnew)
+        query_points = torch.from_numpy(query_points).float()
+        query_points = torch.cat([torch.zeros_like(query_points[:, :1]), query_points], dim=1)
+        gotit = True
+
+        sample = FrameEventData_test(
+            rgb_imgs,
+            event_imgs,
+            segs,
+            traj_data,
+            seq_name=self.seq_names[index],
+            query_points=query_points,
+            img_ifnew=img_ifnew,
+            rgb_timestamp=rgb_timestamp,
+        )
+
+        return sample, gotit
+    
+    def get_a_seq(self, seq_name):
+        for i in range(len(self.seq_names)):
+            if self.seq_names[i] == seq_name:
+                sample, gotit = self.__getitem__(i)
+                return sample, gotit
+        
+        print("WARNNING: did not find the sequence", seq_name)
+        return [], False
+
+    def __len__(self):
+        return len(self.seq_names)

LFE_TAP/datasets/TAPFormer_dataset.py

@@ -0,0 +1,151 @@
+import os
+import torch
+import imageio
+import re
+import numpy as np
+from LFE_TAP.utils.event.utils import read_input
+from LFE_TAP.utils.dataset_utils import FrameEventData, FrameEventData_test
+
+class TAPFormer_dataset(torch.utils.data.Dataset):
+    def __init__(self, data_root, dt=0.020, representation="time_surfaces_v2_5", fix_num=None, with_gt=True):
+        self.data_root = data_root
+        self.dt = dt
+        self.representation = representation
+        self.fix_num = fix_num
+        self.with_gt = with_gt
+        self.seq_names = [
+            fname
+            for fname in os.listdir(data_root)
+            if os.path.isdir(os.path.join(data_root, fname))
+        ]
+        print("found %d unique seqences in %s" % (len(self.seq_names), self.data_root))
+
+    def __getitem__(self, index):
+        gotit =False
+        if self.fix_num is not None:
+            event_dir_path = os.path.join(str(self.data_root), self.seq_names[index], "events", self.representation, f"fix_num_{self.fix_num}")
+        else:
+            event_dir_path = os.path.join(str(self.data_root), self.seq_names[index], "events", self.representation, f"{self.dt:.4f}")
+        rgb_path = os.path.join(str(self.data_root), self.seq_names[index], "images_corrected")
+        track_point_path = os.path.join(str(self.data_root), self.seq_names[index], "annotations.npy")
+        img_time_full = np.stack(np.loadtxt(os.path.join(str(self.data_root), self.seq_names[index], "image_timestamps.txt")))
+
+        img_paths = sorted(os.listdir(rgb_path), key=lambda x: int(re.search(r'\d+', x).group()))
+        event_paths = sorted(os.listdir(event_dir_path), key=lambda x: int(re.search(r'\d+', x).group()))
+        rgb_imgs = []
+        rgb_ifnew = []
+        rgb_ifnew_full = []
+        rgb_times = []
+        rgb_imgs_plus = []
+        rgb_ind = 0
+        rgb_ind_full = 0
+        event_imgs = []
+        event_time = []
+        time_emmbed = []
+        for i, img_path in enumerate(img_paths):
+            try:
+                img = imageio.v2.imread(os.path.join(rgb_path, img_path))
+                if len(img.shape) == 2:
+                    img = img[:,:,np.newaxis]
+                    rgb_imgs.append(img.repeat(3, axis=2))
+                else:
+                    rgb_imgs.append(img)
+                rgb_times.append(int(re.match(r"\d+", img_path).group()))
+            except Exception as e:
+                print(f"error reading image at path:{img_path}")
+                print(f"error mrssage:{str(e)}")
+                gotit = False
+        # rgb_imgs_plus.append(rgb_imgs[0])
+        # event_time.append(rgb_times[0])
+        # event_imgs.append(read_input(os.path.join(self.data_root, self.seq_names[index], "events", "template", self.event_template, str(event_time[0])+".h5"), self.event_template))
+        # rgb_ifnew.append(1)
+        # rgb_ind += 1
+        for i, event_path in enumerate(event_paths):
+            try:
+                if int(event_path.split('.')[0]) < rgb_times[0]:
+                    continue
+                event_imgs.append(read_input(os.path.join(event_dir_path, event_path), self.representation))
+                # event_time.append(int(event_path.split('.')[0]))
+                rgb_time = rgb_times[min(rgb_ind, len(rgb_times)-1)]   
+                if int(event_path.split('.')[0]) >= rgb_time:
+                    rgb_imgs_plus.append(rgb_imgs[min(rgb_ind, len(rgb_times)-1)])
+                    event_time.append(rgb_times[min(rgb_ind, len(rgb_times)-1)])
+                    # time_emmbed.append(round((int(event_path.split('.')[0]) - rgb_times[min(rgb_ind, len(rgb_times)-1)])*1e-4, 1))
+                    rgb_ind += 1
+                    rgb_ifnew.append(1)
+                else:
+                    rgb_imgs_plus.append(rgb_imgs[max(rgb_ind-1,0)])
+                    event_time.append(int(event_path.split('.')[0]))
+                    # time_emmbed.append(round((int(event_path.split('.')[0]) - rgb_times[max(rgb_ind-1,0)])*1e-4, 1))
+                    rgb_ifnew.append(0)
+                
+                    
+                rgb_time_full = img_time_full[rgb_ind_full]
+                if int(event_path.split('.')[0]) >= rgb_time_full:
+                    rgb_ind_full += 1
+                    rgb_ifnew_full.append(1)
+                else:
+                    rgb_ifnew_full.append(0)
+                    
+            except Exception as e:
+                print(f"error reading event at path:{event_path}")
+                print(f"error mrssage:{str(e)}")
+                gotit = False
+
+        rgb_imgs = np.stack(rgb_imgs_plus).transpose(0, 3, 1, 2)
+        event_imgs = np.asarray(event_imgs).transpose(0, 3, 1, 2)
+        # event_imgs = np.stack(event_imgs).transpose(0, 3, 1, 2)
+        if self.with_gt:
+            annot_dict = np.load(track_point_path, allow_pickle=True).item()
+            traj_2d = annot_dict["coords"]              # N, T, 2
+            visibility = annot_dict["visibility"]       # N, T， 1
+        
+            traj_data = np.transpose(traj_2d, (1, 0, 2))                      # N, T, 2 -> T, N, 2
+            visibility = np.transpose(np.logical_not(np.squeeze(visibility)), (1, 0))   # N, T -> T, N
+            query_points = traj_data[0]
+
+        T, H, W, C = event_imgs.shape
+        
+        segs = np.array(event_time)
+        rgb_timestamp = np.array(rgb_times)
+        img_ifnew = np.array(rgb_ifnew)
+        img_ifnew_full = np.array(rgb_ifnew_full)
+        rgb_times = np.array(rgb_times) * 1e-6
+        rgb_time_full = np.array(img_time_full) * 1e-6
+        if self.with_gt:
+            traj_data = np.concatenate([rgb_time_full[:, np.newaxis, np.newaxis].repeat(len(query_points), axis=1), traj_data], axis=2)
+            
+            query_points = torch.from_numpy(query_points).float()
+            query_points = torch.cat([torch.zeros_like(query_points[:, :1]), query_points], dim=1)
+        else:
+            query_points = None
+            traj_data = None
+            visibility = None
+        gotit = True
+
+        sample = FrameEventData_test(
+            rgb_imgs,
+            event_imgs,
+            segs,
+            traj_data,
+            visibility=visibility,
+            seq_name=self.seq_names[index],
+            query_points=query_points,
+            img_ifnew=img_ifnew,
+            img_ifnew_full=img_ifnew_full,
+            rgb_timestamp=rgb_timestamp,
+        )
+
+        return sample, gotit
+    
+    def get_a_seq(self, seq_name):
+        for i in range(len(self.seq_names)):
+            if self.seq_names[i] == seq_name:
+                sample, gotit = self.__getitem__(i)
+                return sample, gotit
+        
+        print("WARNNING: did not find the sequence", seq_name)
+        return [], False
+
+    def __len__(self):
+        return len(self.seq_names)

LFE_TAP/datasets/kubric_movif_dataset.py

@@ -0,0 +1,778 @@
+import os
+import torch
+import imageio
+import cv2
+from PIL import Image
+import numpy as np
+from torchvision.transforms import ColorJitter, GaussianBlur
+from LFE_TAP.utils.dataset_utils import FrameEventData
+from LFE_TAP.utils.event.utils import *
+
+
+class FETAPDataset(torch.utils.data.Dataset):
+    def __init__(self, root_dir, crop_size=(512, 512), seq_len=24, traj_per_sample=512, use_augs=False, **kwargs,):
+        super(FETAPDataset, self).__init__()
+        np.random.seed(0)
+        torch.manual_seed(0)
+        self.root_dir = Path(root_dir)
+        self.crop_size = crop_size
+        self.seq_len = seq_len
+        self.traj_per_sample = traj_per_sample
+        self.use_augs = use_augs
+        
+        # photometric augmentation for rgb images
+        self.photo_aug = ColorJitter(
+            brightness=0.2, contrast=0.2, saturation=0.2, hue=0.25 / 3.14
+        )
+        self.blur_aug = GaussianBlur(11, sigma=(0.1, 2.0))
+        
+        self.blur_aug_prob = 0.25
+        self.color_aug_prob = 0.25
+        
+        # photometric augmentation for event images
+
+        
+        # occlusion augmentation
+        self.eraser_aug_prob = 0.5
+        self.eraser_bounds = [2, 100]
+        self.eraser_max = 10
+
+        # occlusion augmentation
+        self.replace_aug_prob = 0.5
+        self.replace_bounds = [2, 100]
+        self.replace_max = 10
+
+        # spatial augmentations
+        self.pad_bounds = [0, 100]
+        self.crop_size = crop_size
+        self.resize_lim = [0.25, 2.0]  # sample resizes from here
+        self.resize_delta = 0.2
+        self.max_crop_offset = 50
+
+        self.do_flip = True
+        self.h_flip_prob = 0.5
+        self.v_flip_prob = 0.5
+
+    def getitem_helper(self, index):
+        return NotImplementedError
+
+    def __getitem__(self, index):
+        gotit = False
+
+        sample, gotit = self.getitem_helper(index)
+        if not gotit:
+            print("warning: sampling failed")
+            # fake sample, so we can still collate
+            sample = FrameEventData(
+                video=torch.zeros(
+                    (self.seq_len, 3, self.crop_size[0], self.crop_size[1])
+                ),
+                events=torch.zeros((self.seq_len, 10, self.crop_size[0], self.crop_size[1])),
+                segmentation=torch.zeros(
+                    (self.seq_len, 1, self.crop_size[0], self.crop_size[1])
+                ),
+                trajectory=torch.zeros((self.seq_len, self.traj_per_sample, 2)),
+                visibility=torch.zeros((self.seq_len, self.traj_per_sample)),
+                valid=torch.zeros((self.seq_len, self.traj_per_sample)),
+            )
+
+        return sample, gotit
+        
+    def add_photometric_augs(self, rgbs, events, trajs, visibles, eraser=True, replace=True):
+        T, N, _ = trajs.shape
+
+        S = len(rgbs)
+        H, W = rgbs[0].shape[:2]
+        assert S == T
+        
+        # 事件图像增加椒盐噪声
+
+        if eraser:
+            ############ eraser transform (per image after the first) ############
+            rgbs = [rgb.astype(np.float32) for rgb in rgbs]
+            events = [event.astype(np.float32) for event in events]
+            for i in range(1, S):
+                if np.random.rand() < self.eraser_aug_prob:
+                    for _ in range(
+                        np.random.randint(1, self.eraser_max + 1)
+                    ):  # number of times to occlude
+
+                        xc = np.random.randint(0, W)
+                        yc = np.random.randint(0, H)
+                        dx = np.random.randint(
+                            self.eraser_bounds[0], self.eraser_bounds[1]
+                        )
+                        dy = np.random.randint(
+                            self.eraser_bounds[0], self.eraser_bounds[1]
+                        )
+                        x0 = np.clip(xc - dx / 2, 0, W - 1).round().astype(np.int32)
+                        x1 = np.clip(xc + dx / 2, 0, W - 1).round().astype(np.int32)
+                        y0 = np.clip(yc - dy / 2, 0, H - 1).round().astype(np.int32)
+                        y1 = np.clip(yc + dy / 2, 0, H - 1).round().astype(np.int32)
+
+                        mean_color_rgb = np.mean(
+                            rgbs[i][y0:y1, x0:x1, :].reshape(-1, 3), axis=0
+                        )
+                        rgbs[i][y0:y1, x0:x1, :] = mean_color_rgb
+                        
+                        mean_value_event = np.mean(
+                            events[i][y0:y1, x0:x1, :].reshape(-1, 10), axis=0
+                        )
+                        events[i][y0:y1, x0:x1, :] = mean_value_event
+
+                        occ_inds = np.logical_and(
+                            np.logical_and(trajs[i, :, 0] >= x0, trajs[i, :, 0] < x1),
+                            np.logical_and(trajs[i, :, 1] >= y0, trajs[i, :, 1] < y1),
+                        )
+                        visibles[i, occ_inds] = 0
+            rgbs = [rgb.astype(np.uint8) for rgb in rgbs]
+            events = [event.astype(np.uint8) for event in events]
+
+        if replace:
+
+            rgbs_alt = [
+                np.array(self.blur_aug(Image.fromarray(rgb)), dtype=np.uint8)
+                for rgb in rgbs
+            ]
+            events_alt = [
+                np.array(self.blur_aug(torch.from_numpy(event)), dtype=np.uint8)
+                for event in events
+            ]
+
+            ############ replace transform (per image after the first) ############
+            rgbs = [rgb.astype(np.float32) for rgb in rgbs]
+            rgbs_alt = [rgb.astype(np.float32) for rgb in rgbs_alt]
+            events = [event.astype(np.float32) for event in events]
+            events_alt = [event.astype(np.float32) for event in events_alt]
+            for i in range(1, S):
+                if np.random.rand() < self.replace_aug_prob:
+                    for _ in range(
+                        np.random.randint(1, self.replace_max + 1)
+                    ):  # number of times to occlude
+                        xc = np.random.randint(0, W)
+                        yc = np.random.randint(0, H)
+                        dx = np.random.randint(
+                            self.replace_bounds[0], self.replace_bounds[1]
+                        )
+                        dy = np.random.randint(
+                            self.replace_bounds[0], self.replace_bounds[1]
+                        )
+                        x0 = np.clip(xc - dx / 2, 0, W - 1).round().astype(np.int32)
+                        x1 = np.clip(xc + dx / 2, 0, W - 1).round().astype(np.int32)
+                        y0 = np.clip(yc - dy / 2, 0, H - 1).round().astype(np.int32)
+                        y1 = np.clip(yc + dy / 2, 0, H - 1).round().astype(np.int32)
+
+                        wid = x1 - x0
+                        hei = y1 - y0
+                        y00 = np.random.randint(0, H - hei)
+                        x00 = np.random.randint(0, W - wid)
+                        fr = np.random.randint(0, S)
+                        rep_rgb = rgbs_alt[fr][y00 : y00 + hei, x00 : x00 + wid, :]
+                        rgbs[i][y0:y1, x0:x1, :] = rep_rgb
+                        rep_event = events_alt[fr][y00 : y00 + hei, x00 : x00 + wid, :]
+                        events[i][y0:y1, x0:x1, :] = rep_event
+
+                        occ_inds = np.logical_and(
+                            np.logical_and(trajs[i, :, 0] >= x0, trajs[i, :, 0] < x1),
+                            np.logical_and(trajs[i, :, 1] >= y0, trajs[i, :, 1] < y1),
+                        )
+                        visibles[i, occ_inds] = 0
+            rgbs = [rgb.astype(np.uint8) for rgb in rgbs]
+            events = [event.astype(np.uint8) for event in events]
+
+        ############ photometric augmentation ############
+        if np.random.rand() < self.color_aug_prob:
+            # random per-frame amount of aug
+            rgbs = [
+                np.array(self.photo_aug(Image.fromarray(rgb)), dtype=np.uint8)
+                for rgb in rgbs
+            ]
+
+        if np.random.rand() < self.blur_aug_prob:
+            # random per-frame amount of blur
+            rgbs = [
+                np.array(self.blur_aug(Image.fromarray(rgb)), dtype=np.uint8)
+                for rgb in rgbs
+            ]
+            events = [
+                np.array(self.blur_aug(torch.from_numpy(event)), dtype=np.uint8)
+                for event in events
+            ]
+
+        return rgbs, events, trajs, visibles
+
+    def add_spatial_augs(self, rgbs, events, trajs, visibles):
+        T, N, __ = trajs.shape
+
+        S = len(rgbs)
+        H, W = events[0].shape[:2]
+        assert S == T
+
+        rgbs = [rgb.astype(np.float32) for rgb in rgbs]
+        events = [event.astype(np.float32) for event in events]
+
+        ############ spatial transform ############
+
+        # padding
+        pad_x0 = np.random.randint(self.pad_bounds[0], self.pad_bounds[1])
+        pad_x1 = np.random.randint(self.pad_bounds[0], self.pad_bounds[1])
+        pad_y0 = np.random.randint(self.pad_bounds[0], self.pad_bounds[1])
+        pad_y1 = np.random.randint(self.pad_bounds[0], self.pad_bounds[1])
+
+        rgbs = [
+            np.pad(rgb, ((pad_y0, pad_y1), (pad_x0, pad_x1), (0, 0))) for rgb in rgbs
+        ]
+        events = [
+            np.pad(event, ((pad_y0, pad_y1), (pad_x0, pad_x1), (0, 0))) for event in events
+        ]
+        trajs[:, :, 0] += pad_x0
+        trajs[:, :, 1] += pad_y0
+        H, W = rgbs[0].shape[:2]
+
+        # scaling + stretching
+        scale = np.random.uniform(self.resize_lim[0], self.resize_lim[1])
+        scale_x = scale
+        scale_y = scale
+        H_new = H
+        W_new = W
+
+        scale_delta_x = 0.0
+        scale_delta_y = 0.0
+
+        rgbs_scaled = []
+        events_scaled = []
+        for s in range(S):
+            if s == 1:
+                scale_delta_x = np.random.uniform(-self.resize_delta, self.resize_delta)
+                scale_delta_y = np.random.uniform(-self.resize_delta, self.resize_delta)
+            elif s > 1:
+                scale_delta_x = (
+                    scale_delta_x * 0.8
+                    + np.random.uniform(-self.resize_delta, self.resize_delta) * 0.2
+                )
+                scale_delta_y = (
+                    scale_delta_y * 0.8
+                    + np.random.uniform(-self.resize_delta, self.resize_delta) * 0.2
+                )
+            scale_x = scale_x + scale_delta_x
+            scale_y = scale_y + scale_delta_y
+
+            # bring h/w closer
+            scale_xy = (scale_x + scale_y) * 0.5
+            scale_x = scale_x * 0.5 + scale_xy * 0.5
+            scale_y = scale_y * 0.5 + scale_xy * 0.5
+
+            # don't get too crazy
+            scale_x = np.clip(scale_x, 0.2, 2.0)
+            scale_y = np.clip(scale_y, 0.2, 2.0)
+
+            H_new = int(H * scale_y)
+            W_new = int(W * scale_x)
+
+            # make it at least slightly bigger than the crop area,
+            # so that the random cropping can add diversity
+            H_new = np.clip(H_new, self.crop_size[0] + 10, None)
+            W_new = np.clip(W_new, self.crop_size[1] + 10, None)
+            # recompute scale in case we clipped
+            scale_x = (W_new - 1) / float(W - 1)
+            scale_y = (H_new - 1) / float(H - 1)
+
+            rgbs_scaled.append(
+                cv2.resize(rgbs[s], (W_new, H_new), interpolation=cv2.INTER_LINEAR)
+            )
+            events_scaled.append(
+                cv2.resize(events[s], (W_new, H_new), interpolation=cv2.INTER_LINEAR)
+            )
+            trajs[s, :, 0] *= scale_x
+            trajs[s, :, 1] *= scale_y
+        rgbs = rgbs_scaled
+        events = events_scaled
+        ok_inds = visibles[0, :] > 0
+        vis_trajs = trajs[:, ok_inds]  # S,?,2
+
+        if vis_trajs.shape[1] > 0:
+            mid_x = np.mean(vis_trajs[0, :, 0])
+            mid_y = np.mean(vis_trajs[0, :, 1])
+        else:
+            mid_x = self.crop_size[0]
+            mid_y = self.crop_size[1]
+
+        x0 = int(mid_x - self.crop_size[1] // 2)
+        y0 = int(mid_y - self.crop_size[0] // 2)
+
+        offset_x = 0
+        offset_y = 0
+
+        for s in range(S):
+            # on each frame, shift a bit more
+            if s == 1:
+                offset_x = np.random.randint(
+                    -self.max_crop_offset, self.max_crop_offset
+                )
+                offset_y = np.random.randint(
+                    -self.max_crop_offset, self.max_crop_offset
+                )
+            elif s > 1:
+                offset_x = int(
+                    offset_x * 0.8
+                    + np.random.randint(-self.max_crop_offset, self.max_crop_offset + 1)
+                    * 0.2
+                )
+                offset_y = int(
+                    offset_y * 0.8
+                    + np.random.randint(-self.max_crop_offset, self.max_crop_offset + 1)
+                    * 0.2
+                )
+            x0 = x0 + offset_x
+            y0 = y0 + offset_y
+
+            H_new, W_new = rgbs[s].shape[:2]
+            if H_new == self.crop_size[0]:
+                y0 = 0
+            else:
+                y0 = min(max(0, y0), H_new - self.crop_size[0] - 1)
+
+            if W_new == self.crop_size[1]:
+                x0 = 0
+            else:
+                x0 = min(max(0, x0), W_new - self.crop_size[1] - 1)
+
+            rgbs[s] = rgbs[s][y0 : y0 + self.crop_size[0], x0 : x0 + self.crop_size[1]]
+            events[s] = events[s][y0 : y0 + self.crop_size[0], x0 : x0 + self.crop_size[1]]
+            trajs[s, :, 0] -= x0
+            trajs[s, :, 1] -= y0
+
+        H_new = self.crop_size[0]
+        W_new = self.crop_size[1]
+
+        # flip
+        h_flipped = False
+        v_flipped = False
+        if self.do_flip:
+            # h flip
+            if np.random.rand() < self.h_flip_prob:
+                h_flipped = True
+                rgbs = [rgb[:, ::-1] for rgb in rgbs]
+                events = [event[:, ::-1] for event in events]
+            # v flip
+            if np.random.rand() < self.v_flip_prob:
+                v_flipped = True
+                rgbs = [rgb[::-1] for rgb in rgbs]
+                events = [event[::-1] for event in events]
+        if h_flipped:
+            trajs[:, :, 0] = W_new - trajs[:, :, 0]
+        if v_flipped:
+            trajs[:, :, 1] = H_new - trajs[:, :, 1]
+
+        return rgbs, events, trajs
+
+    def crop(self, rgbs, events, trajs, clear_rgbs=None):
+        T, N, _ = trajs.shape
+
+        S = len(events)
+        H, W = events[0].shape[:2]
+        assert S == T
+
+        ############ spatial transform ############
+
+        H_new = H
+        W_new = W
+
+        # simple random crop
+        y0 = 0 if self.crop_size[0] >= H_new else (H_new - self.crop_size[0]) // 2
+        # np.random.randint(0,
+        x0 = 0 if self.crop_size[1] >= W_new else np.random.randint(0, W_new - self.crop_size[1])
+        rgbs = [
+            rgb[y0 : y0 + self.crop_size[0], x0 : x0 + self.crop_size[1]]
+            for rgb in rgbs
+        ]
+        events = [
+            event[y0: y0 + self.crop_size[0], x0: x0 + self.crop_size[1]]
+            for event in events
+        ]
+
+        trajs[:, :, 0] -= x0
+        trajs[:, :, 1] -= y0
+        if clear_rgbs is not None:
+            clear_rgbs = [
+                clear_rgb[y0 : y0 + self.crop_size[0], x0 : x0 + self.crop_size[1]]
+                for clear_rgb in clear_rgbs
+            ]
+            return rgbs, events, trajs, clear_rgbs
+
+        return rgbs, events, trajs
+    
+    
+class KubricMovifDataset(FETAPDataset):
+    def __init__(
+        self,
+        root_dir,
+        representation="time_surfaces_v2_5",
+        event_template="sobel",
+        crop_size=(384, 512),
+        seq_len=24,
+        traj_per_sample=512,
+        sample_vis_1st_frame=False,
+        choose_long_point=False,
+        use_augs=False,
+    ):
+        super(KubricMovifDataset, self).__init__(
+            root_dir=root_dir,
+            representation=representation,
+            event_template=event_template,
+            crop_size=crop_size,
+            seq_len=seq_len,
+            traj_per_sample=traj_per_sample,
+            sample_vis_1st_frame=sample_vis_1st_frame,
+            choose_long_point=choose_long_point,
+            use_augs=use_augs,
+        )
+        self.representation = representation
+        self.event_template = event_template
+        self.sample_vis_1st_frame = sample_vis_1st_frame
+        self.choose_long_point = choose_long_point
+        self.pad_bounds = [0, 25]
+        self.resize_lim = [0.75, 1.25]
+        self.resize_delta = 0.05
+        self.max_crop_offset = 15
+        self.seq_names = [
+            fname
+            for fname in os.listdir(self.root_dir) if os.path.isdir(os.path.join(self.root_dir, fname))
+        ]
+        print("found %d unique seqences in %s" % (len(self.seq_names), self.root_dir))
+        
+    def getitem_helper(self, index):
+        gotit = True
+        seq_name = self.seq_names[index]
+        
+        npy_path = os.path.join(self.root_dir, seq_name, seq_name + ".npy")
+        rgb_dir_path = os.path.join(self.root_dir, seq_name, "frames")
+        event_dir_path = os.path.join(self.root_dir, seq_name, "events", self.representation)
+        
+        rgb_files = sorted(os.listdir(rgb_dir_path))
+        event_files = sorted(os.listdir(event_dir_path))
+        rgb_imgs = []
+        img_ifnew = []
+        event_imgs = []
+        event_imgs.append(read_input(os.path.join(self.root_dir, seq_name, "events", "template", self.event_template, "000.h5"), self.event_template))
+        
+        for i, img_path in enumerate(rgb_files):
+            try:
+                if i % 3 == 0:
+                    rgb_imgs.append(imageio.v2.imread(os.path.join(rgb_dir_path, img_path)))
+                    img_ifnew.append(1)
+                else:
+                    rgb_imgs.append(rgb_imgs[-1])
+                    img_ifnew.append(0)
+            except Exception as e:
+                print(f"error reading image at path:{img_path}_{rgb_dir_path}")
+                print(f"error mrssage:{str(e)}")
+                gotit = False
+                return [], gotit
+            
+        for i, event_path in enumerate(event_files):
+            try:
+                event_imgs.append(read_input(os.path.join(event_dir_path, event_path), self.representation))
+            except Exception as e:
+                print(f"error reading event at path:{event_path}_{event_dir_path}")
+                print(f"error mrssage:{str(e)}")
+                gotit = False
+                return [], gotit
+            
+        rgbs = np.stack(rgb_imgs)
+        events = np.stack(event_imgs)
+        annot_dict = np.load(npy_path, allow_pickle=True).item()
+        traj_2d = annot_dict["coords"]
+        visibility = annot_dict["visibility"]
+        
+        assert self.seq_len == len(rgbs)
+        
+        traj_2d = np.transpose(traj_2d, (1, 0, 2))                      # N, T, 2 -> T, N, 2
+        visibility = np.transpose(np.logical_not(visibility), (1, 0))   # N, T -> T, N
+        if self.use_augs:  
+            # rgbs, events, traj_2d, visibility = self.add_photometric_augs(rgbs, events, traj_2d, visibility)      
+            rgbs, events, traj_2d = self.add_spatial_augs(rgbs, events, traj_2d, visibility)
+        else:
+            rgbs, events, traj_2d = self.crop(rgbs, events, traj_2d)
+            
+        visibility[traj_2d[:, :, 0] > self.crop_size[1] - 1] = False
+        visibility[traj_2d[:, :, 0] < 0] = False
+        visibility[traj_2d[:, :, 1] > self.crop_size[0] - 1] = False
+        visibility[traj_2d[:, :, 1] < 0] = False
+        
+        visibility = torch.from_numpy(visibility)
+        traj_2d = torch.from_numpy(traj_2d)
+        
+        crop_tensor = torch.tensor(self.crop_size).flip(0)[None, None] / 2.0
+        close_pts_inds = torch.all(
+            torch.linalg.vector_norm(traj_2d[..., :2] - crop_tensor, dim=-1) < 1000.0,
+            dim=0,
+        )
+        traj_2d = traj_2d[:, close_pts_inds]
+        visibility = visibility[:, close_pts_inds]
+        
+        visibile_pts_first_frame_inds = (visibility[0]).nonzero(as_tuple=False)
+        
+        if self.sample_vis_1st_frame:
+            visiblile_pts_inds = visibile_pts_first_frame_inds
+        else:
+            visiblile_pts_mid_frame_inds = (visibility[self.seq_len // 2]).nonzero(as_tuple=False)
+            visiblile_pts_inds = torch.cat((visibile_pts_first_frame_inds, visiblile_pts_mid_frame_inds), dim=0)
+            
+        point_inds = torch.randperm(len(visiblile_pts_inds))[:self.traj_per_sample]
+        if len(point_inds) < self.traj_per_sample:
+            gotit = False
+        
+        if self.choose_long_point:
+            distance = np.linalg.norm(traj_2d[-1, visiblile_pts_inds, :] - traj_2d[0, visiblile_pts_inds, :], axis=-1)[:,0]
+            weight = distance / np.sum(distance)
+            point_inds = torch.tensor(np.random.choice(len(distance), size=self.traj_per_sample, p=weight))
+            visible_inds_sampled = visiblile_pts_inds[point_inds]
+            # visible_inds_sampled = visiblile_pts_inds[np.argsort(distance, axis=0)[-self.traj_per_sample:,0]]
+        else:
+            visible_inds_sampled = visiblile_pts_inds[point_inds]
+        
+        if len(visible_inds_sampled.shape) == 2:
+            visible_inds_sampled = visible_inds_sampled.squeeze(1)
+        trajs = traj_2d[:, visible_inds_sampled].float()
+        visibles = visibility[:, visible_inds_sampled]
+        valid = torch.ones((self.seq_len, self.traj_per_sample))
+        
+        rgbs = torch.from_numpy(np.stack(rgbs)).permute(0, 3, 1, 2).float()
+        events = torch.from_numpy(np.stack(events)).permute(0, 3, 1, 2).float()
+        seqs = torch.ones((self.seq_len, 1, self.crop_size[0], self.crop_size[1]))
+        img_ifnew = np.array(img_ifnew)
+        sample = FrameEventData(
+            video = rgbs,
+            events = events,
+            segmentation=seqs,
+            trajectory=trajs,
+            visibility=visibles,
+            valid=valid,
+            seq_name=seq_name,
+            img_ifnew=img_ifnew,
+        )
+        return sample, gotit
+    
+    def __len__(self):
+        return len(self.seq_names)
+    
+    
+class KubricMovifDataset_new(FETAPDataset):
+    def __init__(
+        self,
+        root_dir,
+        root_dir_fast_dataset=None,
+        representation="time_surfaces_v2_5",
+        crop_size=(384, 512),
+        seq_len=24,
+        traj_per_sample=256,
+        sample_vis_1st_frame=False,
+        choose_long_point=False,
+        use_augs=False,
+        if_test=False,
+    ):
+        super(KubricMovifDataset_new, self).__init__(
+            root_dir=root_dir,
+            root_dir_fast_dataset=root_dir_fast_dataset,
+            representation=representation,
+            crop_size=crop_size,
+            seq_len=seq_len,
+            traj_per_sample=traj_per_sample,
+            sample_vis_1st_frame=sample_vis_1st_frame,
+            choose_long_point=choose_long_point,
+            use_augs=use_augs,
+            if_test=if_test,
+        )
+        self.root_dir1 = os.path.join(self.root_dir, "kubric_ori_dataset1")
+        self.root_dir2 = os.path.join(self.root_dir, "kubric_ori_dataset2")
+        self.representation = representation
+        self.sample_vis_1st_frame = sample_vis_1st_frame
+        self.choose_long_point = choose_long_point
+        self.pad_bounds = [0, 25]
+        self.resize_lim = [0.75, 1.25]
+        self.resize_delta = 0.05
+        self.max_crop_offset = 15
+        self.if_test = if_test
+        if root_dir_fast_dataset is not None:
+            root_dir_fast_dataset = Path(root_dir_fast_dataset)
+            seq_names1 = [
+                os.path.join(self.root_dir1, fname)
+                for fname in os.listdir(self.root_dir1) if os.path.isdir(os.path.join(self.root_dir1, fname))
+            ]
+            seq_names2 = [
+                os.path.join(self.root_dir2, fname)
+                for fname in os.listdir(self.root_dir2) if os.path.isdir(os.path.join(self.root_dir2, fname))
+            ]
+            
+            seq_names_fast = [
+                os.path.join(root_dir_fast_dataset, fname)
+                for fname in os.listdir(root_dir_fast_dataset) if os.path.isdir(os.path.join(root_dir_fast_dataset, fname))
+            ]
+            self.seq_names = seq_names1 + seq_names2 + seq_names_fast
+        else:
+            self.seq_names = [
+                os.path.join(self.root_dir, fname)
+                for fname in os.listdir(self.root_dir) if os.path.isdir(os.path.join(self.root_dir, fname))
+            ]
+        print("found %d unique seqences in %s" % (len(self.seq_names), self.root_dir))
+        
+    def getitem_helper(self, index):
+        gotit = True
+        seq_name = self.seq_names[index]
+        
+        npy_path = os.path.join(seq_name, os.path.basename(seq_name) + ".npy")
+        if not os.path.exists(npy_path):
+            print(npy_path)
+        rgb_dir_path = os.path.join(seq_name, "blur_frames")
+        clear_rgb_dir_path = os.path.join(seq_name, "frames")
+        event_dir_path = os.path.join(seq_name, "events", self.representation)
+        
+        rgb_files = sorted(os.listdir(rgb_dir_path))
+        clear_rgb_files = sorted(os.listdir(clear_rgb_dir_path))
+        event_files = sorted(os.listdir(event_dir_path))
+        rgb_imgs = []
+        img_ifnew = []
+        clear_rgb_imgs = []
+        event_imgs = []
+        if not self.if_test:
+            random_id = np.random.randint(0, 95-24)
+            rgb_files = rgb_files[random_id: random_id+24]
+            event_files = event_files[random_id: random_id+24]
+            clear_rgb_files = clear_rgb_files[random_id: random_id+24]
+        else:
+            rgb_files = rgb_files[:-1]
+            clear_rgb_files = clear_rgb_files[:-1]
+        
+        next_insert = 0  # 记录下一个需要插入新图像的索引
+        for i, img_path in enumerate(rgb_files):
+            try:
+                if i == next_insert:
+                    # 读取新图像并添加到列表
+                    img = imageio.v2.imread(os.path.join(rgb_dir_path, img_path))
+                    rgb_imgs.append(img)
+                    img_ifnew.append(1)
+                    # 随机选择下一个间隔（3或4）
+                    if not self.if_test:
+                        next_insert += np.random.choice([3, 4])
+                    else:
+                        next_insert += 4
+                else:
+                    rgb_imgs.append(rgb_imgs[-1])
+                    img_ifnew.append(0)
+            except Exception as e:
+                print(f"error reading image at path:{img_path}_{rgb_dir_path}")
+                print(f"error mrssage:{str(e)}")
+                gotit = False
+                return [], gotit
+            
+        for i, img_path in enumerate(clear_rgb_files):
+            try:
+                clear_rgb_imgs.append(imageio.v2.imread(os.path.join(clear_rgb_dir_path, img_path)))
+            except Exception as e:
+                print(f"error reading image at path:{img_path}_{clear_rgb_dir_path}")
+                print(f"error mrssage:{str(e)}")
+                gotit = False
+                return [], gotit
+            
+        for i, event_path in enumerate(event_files):
+            try:
+                event_imgs.append(read_input(os.path.join(event_dir_path, event_path), self.representation))
+            except Exception as e:
+                print(f"error reading event at path:{event_path}_{event_dir_path}")
+                print(f"error mrssage:{str(e)}")
+                gotit = False
+                return [], gotit
+            
+        rgbs = np.stack(rgb_imgs)
+        clear_rgbs = np.stack(clear_rgb_imgs)
+        events = np.stack(event_imgs)
+        annot_dict = np.load(npy_path, allow_pickle=True).item()
+        if not self.if_test:
+            traj_2d = annot_dict["coords"][:,random_id: random_id+24]
+            visibility = annot_dict["visibility"][:,random_id: random_id+24]
+        else:
+            traj_2d = annot_dict["coords"][:,:-1]
+            visibility = annot_dict["visibility"][:,:-1]
+        
+        assert self.seq_len == len(rgbs) == len(clear_rgbs)
+        
+        traj_2d = np.transpose(traj_2d, (1, 0, 2))                      # N, T, 2 -> T, N, 2
+        visibility = np.transpose(np.logical_not(visibility), (1, 0))   # N, T -> T, N
+        if self.use_augs: 
+            print("new kubric dataset can't use augs!!!") 
+            # rgbs, events, traj_2d, visibility = self.add_photometric_augs(rgbs, events, traj_2d, visibility)      
+            rgbs, events, traj_2d = self.add_spatial_augs(rgbs, events, traj_2d, visibility)
+        else:
+            rgbs, events, traj_2d, clear_rgbs = self.crop(rgbs, events, traj_2d, clear_rgbs=clear_rgbs)
+            
+        visibility[traj_2d[:, :, 0] > self.crop_size[1] - 1] = False
+        visibility[traj_2d[:, :, 0] < 0] = False
+        visibility[traj_2d[:, :, 1] > self.crop_size[0] - 1] = False
+        visibility[traj_2d[:, :, 1] < 0] = False 
+        
+        visibility = torch.from_numpy(visibility)
+        traj_2d = torch.from_numpy(traj_2d)
+        
+        crop_tensor = torch.tensor(self.crop_size).flip(0)[None, None] / 2.0
+        close_pts_inds = torch.all(
+            torch.linalg.vector_norm(traj_2d[..., :2] - crop_tensor, dim=-1) < 1000.0,
+            dim=0,
+        )
+        traj_2d = traj_2d[:, close_pts_inds]
+        visibility = visibility[:, close_pts_inds]
+        
+        visibile_pts_first_frame_inds = (visibility[0]).nonzero(as_tuple=False)
+        
+        if self.sample_vis_1st_frame:
+            visiblile_pts_inds = visibile_pts_first_frame_inds
+        else:
+            visiblile_pts_mid_frame_inds = (visibility[self.seq_len // 2]).nonzero(as_tuple=False)
+            visibile_pts_last_frame_inds = (visibility[self.seq_len - 1]).nonzero(as_tuple=False)
+            visiblile_pts_inds = torch.cat((visibile_pts_first_frame_inds, visiblile_pts_mid_frame_inds, visibile_pts_last_frame_inds), dim=0)
+            
+        point_inds = torch.randperm(len(visiblile_pts_inds))[:self.traj_per_sample]
+        if len(point_inds) < self.traj_per_sample and not self.if_test:
+            print(seq_name, "get point num ", len(point_inds), "less than ", self.traj_per_sample, " random_id is", random_id)
+            gotit = False
+            # shutil.rmtree(seq_name)
+        
+        if self.choose_long_point:
+            distance = np.linalg.norm(traj_2d[-1, visiblile_pts_inds, :] - traj_2d[0, visiblile_pts_inds, :], axis=-1)[:,0]
+            weight = distance / np.sum(distance)
+            point_inds = torch.tensor(np.random.choice(len(distance), size=self.traj_per_sample, p=weight))
+            visible_inds_sampled = visiblile_pts_inds[point_inds]
+            # visible_inds_sampled = visiblile_pts_inds[np.argsort(distance, axis=0)[-self.traj_per_sample:,0]]
+        else:
+            visible_inds_sampled = visiblile_pts_inds[point_inds]
+        
+        if len(visible_inds_sampled.shape) == 2:
+            visible_inds_sampled = visible_inds_sampled.squeeze(1)
+        trajs = traj_2d[:, visible_inds_sampled].float()
+        visibles = visibility[:, visible_inds_sampled]
+        valid = torch.ones((self.seq_len, self.traj_per_sample))
+        
+        rgbs = torch.from_numpy(np.stack(rgbs)).permute(0, 3, 1, 2).float()
+        clear_rgbs = torch.from_numpy(np.stack(clear_rgbs)).permute(0, 3, 1, 2).float()
+        events = torch.from_numpy(np.stack(events)).float()
+        if "event_stack" not in self.representation:
+            events = events.permute(0, 3, 1, 2)
+        # events = torch.from_numpy(np.stack(events)).permute(0, 3, 1, 2).float()
+        seqs = torch.ones((self.seq_len, 1, self.crop_size[0], self.crop_size[1]))
+        img_ifnew = torch.Tensor(img_ifnew)
+        sample = FrameEventData(
+            video = rgbs,
+            events = events,
+            segmentation=seqs,
+            trajectory=trajs,
+            visibility=visibles,
+            valid=valid,
+            seq_name=seq_name,
+            img_ifnew=img_ifnew,
+            clear_video = clear_rgbs,
+        )
+        assert sample.img_ifnew is not None, f"发现空值样本"
+        return sample, gotit
+    
+    def __len__(self):
+        return len(self.seq_names)

LFE_TAP/evaluator/evaluation_pred.py

@@ -0,0 +1,184 @@
+import torch
+import numpy as np
+import torch.nn.functional as F
+from tqdm import tqdm
+from typing import Tuple
+
+from LFE_TAP.models.tapfe import TAPFormer
+from LFE_TAP.utils.model_utils import get_points_on_a_grid, get_sift_sampled_pts, get_uniformly_sampled_pts, normalize_voxels
+
+class EvaluationPredictor(torch.nn.Module):
+    def __init__(
+        self,
+        model: TAPFormer,
+        interp_shape: Tuple[int, int] = [384, 512],
+        grid_size: int = 5,
+        local_grid_size: int = 8,
+        single_point: bool = True,
+        sift_size: int = 0,
+        num_uniformly_sampled_pts: int = 0,
+        n_iters: int = 6,
+        local_extent: int = 50,
+        if_test: bool = False,
+    ) -> None:
+        super(EvaluationPredictor, self).__init__()
+        self.model = model
+        self.interp_shape = interp_shape
+        self.grid_size = grid_size
+        self.local_grid_size = local_grid_size
+        self.single_point = single_point
+        self.sift_size = sift_size
+        self.num_uniformly_sampled_pts = num_uniformly_sampled_pts
+        self.n_iters = n_iters
+        self.local_extent = local_extent
+        self.if_test = if_test
+        self.model.eval()
+    
+    def forward(self, video, events, queries, img_ifnew=None):
+        if queries is None and self.grid_size > 0:
+            grid_pts = get_points_on_a_grid(
+                self.grid_size, self.interp_shape, device="cuda"
+            )
+            queries = torch.cat(
+                [torch.ones_like(grid_pts[:, :, :1]) * 0, grid_pts],
+                dim=2,
+            )
+            self.grid_size = 0
+            
+        queries = queries.clone()
+        B, T, C_r, H, W = video.shape
+        C_e = events[0].shape[1]
+        B, N, D = queries.shape
+        device = queries.device
+        
+        assert D == 3
+        assert B == 1
+        
+        interp_shape = self.interp_shape
+        
+        if isinstance(video, torch.Tensor) and isinstance(events, torch.Tensor):
+            video = video.reshape(B * T, C_r, H, W)
+            events = events.reshape(B * T, C_e, H, W)
+            video = F.interpolate(video, tuple(interp_shape), mode="bilinear", align_corners=True)
+            events = F.interpolate(events, tuple(interp_shape), mode="bilinear", align_corners=True)
+            video = video.reshape(B, T, C_r, interp_shape[0], interp_shape[1])
+            events = events.reshape(B, T, C_e, interp_shape[0], interp_shape[1])
+        
+        queries[:, :, 1] *= (interp_shape[1] - 1) / (W - 1)
+        queries[:, :, 2] *= (interp_shape[0] - 1) / (H - 1)
+        
+        if self.single_point:
+            traj_e = torch.zeros((B, T, N, 2), device=device)
+            vis_e = torch.zeros((B, T, N), device=device)
+            conf_e = torch.zeros((B, T, N), device=device)
+            
+            for pind in range((N)):
+                querie = queries[:, pind : pind + 1]
+                traj_e_pind, vis_e_pind, conf_e_pind = self._process_one_point(video[:,:], events[:,:], querie, img_ifnew=img_ifnew)
+                traj_e[:, :, pind:pind+1] = traj_e_pind[:, :, :1]
+                vis_e[:, :, pind:pind+1] = vis_e_pind[:, :, :1]
+                conf_e[:, :, pind:pind+1] = conf_e_pind[:, :, :1]
+        else:
+            if self.grid_size > 0:
+                xy = get_points_on_a_grid(self.grid_size, video.shape[3:])
+                xy = torch.cat([torch.zeros_like(xy[:, :, :1]), xy], dim=2).to(device)  
+                queries = torch.cat([queries, xy], dim=1)  
+                
+            if self.num_uniformly_sampled_pts > 0:
+                xy = get_uniformly_sampled_pts(
+                    self.num_uniformly_sampled_pts,
+                    video.shape[1],
+                    video.shape[3:],
+                    device=device,
+                )
+                queries = torch.cat([queries, xy], dim=1)  
+                
+            sift_size = self.sift_size
+            if sift_size > 0:
+                xy = get_sift_sampled_pts(video, sift_size, T, [H, W], device=device)
+                if xy.shape[1] == sift_size:
+                    queries = torch.cat([queries, xy], dim=1)  #
+                else:
+                    sift_size = 0 
+        
+            preds = self.model(rgbs=video, events=events, queries=queries, iters=self.n_iters, img_ifnew=img_ifnew)
+            traj_e, vis_e = preds[0], preds[1]
+                
+            conf_e = preds[2]
+            if (
+                sift_size > 0
+                or self.grid_size > 0
+                or self.num_uniformly_sampled_pts > 0
+            ):
+                traj_e = traj_e[:, :, : -self.grid_size**2 - sift_size - self.num_uniformly_sampled_pts]
+                vis_e = vis_e[:, :, : -self.grid_size**2 - sift_size - self.num_uniformly_sampled_pts]
+                if conf_e is not None:
+                    conf_e = conf_e[:, :, : -self.grid_size**2 - sift_size - self.num_uniformly_sampled_pts]
+        
+        # if conf_e is not None:
+        #     vis_e = vis_e * conf_e
+        
+        if self.if_test:
+            thr = 0.9
+            vis_e = vis_e > thr
+            for i in range(len(queries)):
+                queries_t = queries[i, : traj_e.size(2), 0].to(torch.int64)
+                arange = torch.arange(0, len(queries_t))
+
+                # overwrite the predictions with the query points
+                traj_e[i, queries_t, arange] = queries[i, : traj_e.size(2), 1:]
+
+                # correct visibilities, the query points should be visible
+                vis_e[i, queries_t, arange] = True
+                
+        traj_e *= traj_e.new_tensor(
+            [(W - 1) / (self.interp_shape[1] - 1), (H - 1) / (self.interp_shape[0] - 1)]
+        )
+        
+        return traj_e, vis_e, conf_e
+                
+    def _process_one_point(self, video, events, query, img_ifnew):
+        B, T, C, H, W = video.shape
+        device = query.device
+        if self.local_grid_size > 0:
+            xy_target = get_points_on_a_grid(self.local_grid_size, (self.local_extent, self.local_extent), [query[0,0,2].item(), query[0,0,1].item()])
+            xy_target = torch.cat([torch.zeros_like(xy_target[:, :, :1]), xy_target], dim=2).to(device)
+            query = torch.cat([query, xy_target], dim=1)
+            
+        if self.grid_size > 0:
+            xy = get_points_on_a_grid(self.grid_size, video.shape[3:])
+            xy = torch.cat([torch.zeros_like(xy[:, :, :1]), xy], dim=2).to(device)  
+            query = torch.cat([query, xy], dim=1) 
+            
+        sift_size = self.sift_size
+        if sift_size > 0:
+            xy = get_sift_sampled_pts(video, sift_size, T, [H, W], device=device)
+            sift_size = xy.shape[1]
+            if sift_size > 0:
+                query = torch.cat([query, xy], dim=1)  #
+
+            num_uniformly_sampled_pts = self.sift_size - sift_size
+            if num_uniformly_sampled_pts > 0:
+                xy2 = get_uniformly_sampled_pts(
+                    num_uniformly_sampled_pts,
+                    video.shape[1],
+                    video.shape[3:],
+                    device=device,
+                )
+                query = torch.cat([query, xy2], dim=1)  #
+
+        if self.num_uniformly_sampled_pts > 0:
+            xy = get_uniformly_sampled_pts(
+                self.num_uniformly_sampled_pts,
+                video.shape[1],
+                video.shape[3:],
+                device=device,
+            )
+            query = torch.cat([query, xy], dim=1)  
+            
+        traj_e_pind, vis_e_pind, conf_e_pind = self.model(
+            rgbs=video, events=events, queries=query, iters=self.n_iters, img_ifnew=img_ifnew
+        )
+    
+        return traj_e_pind[..., :2], vis_e_pind, conf_e_pind
+        

LFE_TAP/evaluator/evaluator.py

@@ -0,0 +1,351 @@
+import os
+import numpy as np
+import torch
+import logging
+from tqdm import tqdm
+from typing import Optional, Mapping
+from torch.utils.tensorboard import SummaryWriter
+from LFE_TAP.models.tapfe import TAPFormer
+from LFE_TAP.utils.visualizer import Visualizer
+from LFE_TAP.utils.dataset_utils import dataclass_to_cuda_
+
+def get_error(est_data, gt_data):
+    # discard gt which happen after last est_data
+    # gt_data = gt_data[gt_data[:, 0] <= est_data[-1, 0]]
+
+    est_t, est_x, est_y = est_data.T
+    gt_t, gt_x, gt_y = gt_data.T
+
+    
+    if len(gt_t) == 0 or len(est_t) == 0:
+        return [], [], []
+
+    if len(est_t) < 2:
+        return gt_t, np.array([0]), np.array([0])
+
+    # find samples which have dt > threshold
+    error_x = np.interp(gt_t, est_t, est_x) - gt_x
+    error_y = np.interp(gt_t, est_t, est_y) - gt_y
+
+    return gt_t, error_x, error_y
+
+
+def compareTracks(tracks_pred, tracks_gt, max_distance=10):
+    pred_ids = np.unique(tracks_pred[:, 0])
+    gt_ids = np.unique(tracks_gt[:, 0])
+    pred_datas = {i: tracks_pred[tracks_pred[:, 0] == i, 1:] for i in pred_ids}
+    gt_datas = {i: tracks_gt[tracks_gt[:, 0] == i, 1:] for i in gt_ids}
+    
+    error_datas = np.zeros(shape=(0, 4))
+    errors = np.zeros(shape=(0, 2))
+    
+    for track_id, pred_data in tqdm(pred_datas.items(), disable=True):
+        gt_data = gt_datas[track_id]
+        init_time = gt_data[0, 0]
+        gt_data[:, 0] -= init_time
+        pred_data[:, 0] -= init_time
+        pred_data = np.concatenate((gt_data[0, :].reshape(1, -1), pred_data), axis=0)
+        
+        gt_t, error_x, error_y = get_error(pred_data, gt_data)
+        
+        if len(gt_t) != 0:
+            ids = (track_id * np.ones_like(error_x)).astype(int)
+            added_data = np.stack([ids, gt_t, error_x, error_y]).T
+            error_euclidean = np.sqrt(added_data[:, 2]**2 + added_data[:, 3]**2)
+            error_euclidean[0] = 0
+            idxs = np.where(error_euclidean > max_distance)[0]
+            if len(idxs) > 0:
+                feature_age = (gt_t[int(idxs[0])] - gt_t[0]) / (gt_t[-1] - gt_t[0])
+                if int(idxs[0]) <= 1:
+                    continue
+                error_euclidean = error_euclidean[:idxs[0]]
+                error = np.mean(error_euclidean)
+            else:
+                feature_age = 1
+                error = np.mean(error_euclidean)
+                
+            errors = np.concatenate([errors, [[error, feature_age]]])
+            error_datas = np.concatenate([error_datas, added_data])
+            
+    if errors.shape[0] == 0:
+        return [], [], [0, (gt_t[1]-gt_t[0])/(gt_t[-1] - gt_t[0]), 0]
+            
+    mean_error = np.mean(errors, axis=0)
+    ecpect_FA = (errors.shape[0]/len(gt_ids)) * mean_error[1]
+    mean_error = np.concatenate((mean_error, [ecpect_FA]))
+    # print("abandon ", len(gt_ids) - errors.shape[0], " points, remain ", errors.shape[0], " points")
+    # print("Mean Error: {:.4f}, Mean Feature Age: {:.4f}, Ecppect Feature Age: {:.4f}".format(mean_error[0], mean_error[1], mean_error[2]))
+    
+    return error_datas, errors, mean_error
+
+
+def compute_tapvid_metrics(
+    query_points: np.ndarray,
+    gt_occluded: np.ndarray,
+    gt_tracks: np.ndarray,
+    pred_occluded: np.ndarray,
+    pred_tracks: np.ndarray,
+    query_mode: str,
+) -> Mapping[str, np.ndarray]:
+    """Computes TAP-Vid metrics (Jaccard, Pts. Within Thresh, Occ. Acc.)
+    See the TAP-Vid paper for details on the metric computation.  All inputs are
+    given in raster coordinates.  The first three arguments should be the direct
+    outputs of the reader: the 'query_points', 'occluded', and 'target_points'.
+    The paper metrics assume these are scaled relative to 256x256 images.
+    pred_occluded and pred_tracks are your algorithm's predictions.
+    This function takes a batch of inputs, and computes metrics separately for
+    each video.  The metrics for the full benchmark are a simple mean of the
+    metrics across the full set of videos.  These numbers are between 0 and 1,
+    but the paper multiplies them by 100 to ease reading.
+    Args:
+       query_points: The query points, an in the format [t, y, x].  Its size is
+         [b, n, 3], where b is the batch size and n is the number of queries
+       gt_occluded: A boolean array of shape [b, n, t], where t is the number
+         of frames.  True indicates that the point is occluded.
+       gt_tracks: The target points, of shape [b, n, t, 2].  Each point is
+         in the format [x, y]
+       pred_occluded: A boolean array of predicted occlusions, in the same
+         format as gt_occluded.
+       pred_tracks: An array of track predictions from your algorithm, in the
+         same format as gt_tracks.
+       query_mode: Either 'first' or 'strided', depending on how queries are
+         sampled.  If 'first', we assume the prior knowledge that all points
+         before the query point are occluded, and these are removed from the
+         evaluation.
+    Returns:
+        A dict with the following keys:
+        occlusion_accuracy: Accuracy at predicting occlusion.
+        pts_within_{x} for x in [1, 2, 4, 8, 16]: Fraction of points
+          predicted to be within the given pixel threshold, ignoring occlusion
+          prediction.
+        jaccard_{x} for x in [1, 2, 4, 8, 16]: Jaccard metric for the given
+          threshold
+        average_pts_within_thresh: average across pts_within_{x}
+        average_jaccard: average across jaccard_{x}
+    """
+
+    metrics = {}
+    eye = np.eye(gt_tracks.shape[2], dtype=np.int32)
+
+    if query_mode == "first":
+        # evaluate frames after the query frame
+        query_frame_to_eval_frames = np.cumsum(eye, axis=1) - eye
+    elif query_mode == "strided":
+        # evaluate all frames except the query frame
+        query_frame_to_eval_frames = 1 - eye
+    else:
+        raise ValueError("Unknown query mode " + query_mode)
+
+    query_frame = query_points[..., 0]
+    query_frame = np.round(query_frame).astype(np.int32)
+    evaluation_points = query_frame_to_eval_frames[query_frame] > 0
+
+    # Occlusion accuracy is simply how often the predicted occlusion equals the
+    # ground truth.
+    occ_acc = np.sum(
+        np.equal(pred_occluded, gt_occluded) & evaluation_points,
+        axis=(1, 2),
+    ) / np.sum(evaluation_points)
+    metrics["occlusion_accuracy"] = occ_acc
+
+    # Next, convert the predictions and ground truth positions into pixel
+    # coordinates.
+    visible = np.logical_not(gt_occluded)
+    pred_visible = np.logical_not(pred_occluded)
+    all_frac_within = []
+    all_jaccard = []
+    for thresh in [1, 2, 4, 8, 16]:
+        # True positives are points that are within the threshold and where both
+        # the prediction and the ground truth are listed as visible.
+        within_dist = np.sum(
+            np.square(pred_tracks - gt_tracks),
+            axis=-1,
+        ) < np.square(thresh)
+        is_correct = np.logical_and(within_dist, visible)
+
+        # Compute the frac_within_threshold, which is the fraction of points
+        # within the threshold among points that are visible in the ground truth,
+        # ignoring whether they're predicted to be visible.
+        count_correct = np.sum(
+            is_correct & evaluation_points,
+            axis=(1, 2),
+        )
+        count_visible_points = np.sum(visible & evaluation_points, axis=(1, 2))
+        frac_correct = count_correct / count_visible_points
+        metrics["pts_within_" + str(thresh)] = frac_correct
+        all_frac_within.append(frac_correct)
+
+        true_positives = np.sum(
+            is_correct & pred_visible & evaluation_points, axis=(1, 2)
+        )
+
+        # The denominator of the jaccard metric is the true positives plus
+        # false positives plus false negatives.  However, note that true positives
+        # plus false negatives is simply the number of points in the ground truth
+        # which is easier to compute than trying to compute all three quantities.
+        # Thus we just add the number of points in the ground truth to the number
+        # of false positives.
+        #
+        # False positives are simply points that are predicted to be visible,
+        # but the ground truth is not visible or too far from the prediction.
+        gt_positives = np.sum(visible & evaluation_points, axis=(1, 2))
+        false_positives = (~visible) & pred_visible
+        false_positives = false_positives | ((~within_dist) & pred_visible)
+        false_positives = np.sum(false_positives & evaluation_points, axis=(1, 2))
+        jaccard = true_positives / (gt_positives + false_positives)
+        metrics["jaccard_" + str(thresh)] = jaccard
+        all_jaccard.append(jaccard)
+    metrics["average_jaccard"] = np.mean(
+        np.stack(all_jaccard, axis=1),
+        axis=1,
+    )
+    metrics["average_pts_within_thresh"] = np.mean(
+        np.stack(all_frac_within, axis=1),
+        axis=1,
+    )
+    return metrics
+
+class Evaluator:
+    def __init__(self, output_dir) -> None:
+        self.output_dir = output_dir
+        os.makedirs(output_dir, exist_ok=True)
+        
+    def compute_metrics(self, metrics, sample, pred_trajectory, dataset_name):
+        if isinstance(pred_trajectory, tuple):
+            pred_trajectory, pred_visibility = pred_trajectory
+        else:
+            pred_visibility = None
+            
+        if "kubric" in dataset_name:
+            B, T, N, D = sample.trajectory.shape
+            traj = sample.trajectory.clone()
+            thr = 0.6
+
+            if pred_visibility is None:
+                logging.warning("visibility is NONE")
+                pred_visibility = torch.zeros_like(sample.visibility)
+
+            if not pred_visibility.dtype == torch.bool:
+                pred_visibility = pred_visibility > thr
+
+            query_points = torch.cat(
+                    [
+                        torch.zeros_like(sample.trajectory[:, 0, :, :1]),
+                        sample.trajectory[:, 0],
+                    ],
+                    dim=2,
+                ).cpu().numpy()
+
+            pred_visibility = pred_visibility[:, :, :N]
+            pred_trajectory = pred_trajectory[:, :, :N]
+
+            gt_tracks = traj.permute(0, 2, 1, 3).cpu().numpy()
+            gt_occluded = (
+                torch.logical_not(sample.visibility.clone().permute(0, 2, 1))
+                .cpu()
+                .numpy()
+            )
+
+            pred_occluded = (
+                torch.logical_not(pred_visibility.clone().permute(0, 2, 1))
+                .cpu()
+                .numpy()
+            )
+            pred_tracks = pred_trajectory.permute(0, 2, 1, 3).cpu().numpy()
+
+            out_metrics = compute_tapvid_metrics(
+                query_points,
+                gt_occluded,
+                gt_tracks,
+                pred_occluded,
+                pred_tracks,
+                query_mode="strided" if "strided" in dataset_name else "first",
+            )
+
+            metrics[sample.seq_name[0]] = out_metrics
+            for metric_name in out_metrics.keys():
+                if "avg" not in metrics:
+                    metrics["avg"] = {}
+                metrics["avg"][metric_name] = np.mean(
+                    [v[metric_name] for k, v in metrics.items() if k != "avg"]
+                )
+
+            logging.info(f"Metrics: {out_metrics}")
+            logging.info(f"avg: {metrics['avg']}")
+            print("metrics", out_metrics)
+            print("avg", metrics["avg"])
+        elif "EC" or "EDS" in dataset_name:
+            traj = sample.trajectory.copy()
+            
+            mean_err_avg = []
+            for i in range(1, 31):
+                error_datas, errors, mean_err = compareTracks(pred_trajectory.cpu().numpy(), traj[0], sample.segmentation, i)
+                mean_err_avg.append(mean_err)
+            mean_err_avg = np.stack(mean_err_avg)
+            mean_err_avg = np.mean(mean_err_avg, axis=0)
+            print(sample.seq_name, "deep_ev mean error:", mean_err_avg[0], " mean age:", mean_err_avg[1], "expect age:", mean_err_avg[2])
+            
+    
+    @torch.no_grad()       
+    def evaluate_sequence(
+        self,
+        model,
+        test_dataloader: torch.utils.data.DataLoader,
+        dataset_name: str,
+        visualize_every: int = 1,
+        writer: Optional[SummaryWriter] = None,
+        step: Optional[int] = 0,
+    ):
+        metrics = {}
+        
+        vis = Visualizer(self.output_dir, fps=10 if "kubric" in dataset_name else 50)
+        
+        for ind, sample in enumerate(tqdm(test_dataloader)):
+            if isinstance(sample, tuple):
+                sample, gotit = sample
+                if not all(gotit):
+                    print(f"Skipping sample {ind} because gotit is {gotit}")
+                    continue
+                
+            if torch.cuda.is_available():
+                if "kubric" in dataset_name:
+                    dataclass_to_cuda_(sample)
+                device = torch.device("cuda")
+            else:
+                device = torch.device("cpu")
+        
+            if "kubric" in dataset_name:
+                queries = torch.cat(
+                    [
+                        torch.zeros_like(sample.trajectory[:, 0, :, :1]),
+                        sample.trajectory[:, 0],
+                    ],
+                    dim=2,
+                ).to(device)
+            elif "EC" or "EDS" in dataset_name:
+                queries = sample.query_points
+                queries = queries.to(device)
+            
+            pred_tracks = model(sample.video, sample.events, queries)
+            
+            if dataset_name == "EC" or dataset_name == "EDS":
+                seq_name = sample.seq_name[0]
+            else:
+                seq_name = str(ind)
+            if ind % visualize_every == 0:
+                vis.visualize(
+                    sample.video if isinstance(sample.video, torch.Tensor) else torch.from_numpy(sample.video).float(), 
+                    pred_tracks[0], 
+                    pred_tracks[1] > 0.8,
+                    filename=dataset_name + "_" + seq_name,
+                    writer=writer,
+                    step=step,
+                )
+            self.compute_metrics(metrics, sample, pred_tracks, dataset_name)
+        return metrics
+                
+            
+                
+            
+                
+            

LFE_TAP/evaluator/prediction.py

@@ -0,0 +1,311 @@
+import torch
+import time
+import torch.nn.functional as F
+import numpy as np
+
+from LFE_TAP.models.tapfe import TAPFormer, posenc
+from LFE_TAP.utils.model_utils import get_track_feat, normalize_voxels
+from LFE_TAP.models.embeddings import get_1d_sincos_pos_embed_from_grid
+
+torch.manual_seed(0)
+starter, ender = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
+starter1, ender1 = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
+
+class TAPFormer_online(TAPFormer):
+    def __init__(self, trained_model=None, window_size=16, stride=4, corr_radius=3, corr_levels=3, 
+                 backbone="basic", num_heads=8, hidden_size=384, space_depth=3, time_depth=3):
+        """
+        Initialize TAPFormer_online model for memory-efficient inference.
+        
+        Args:
+            trained_model: (optional) A trained TAPFormer model instance. If provided, parameters will be copied from it.
+            window_size, stride, corr_radius, corr_levels, backbone, num_heads, hidden_size, 
+            space_depth, time_depth: Model configuration parameters. Used when trained_model is None.
+        """
+        # If trained_model is provided, use it to initialize (backward compatibility)
+        if trained_model is not None:
+            super(TAPFormer_online, self).__init__(
+                window_size=trained_model.window_size,
+                stride=trained_model.stride,
+                corr_radius=trained_model.corr_radius,
+                corr_levels=trained_model.corr_levels,
+                backbone=trained_model.backbone,
+                hidden_size=trained_model.hidden_size,
+                space_depth=trained_model.space_depth,
+                time_depth=trained_model.time_depth
+            )
+            self.fusion_block = trained_model.fusion_block
+            self.updateformer2 = trained_model.updateformer2
+            self.corr_mlp = trained_model.corr_mlp
+        else:
+            # Direct initialization from parameters
+            super(TAPFormer_online, self).__init__(
+                window_size=window_size,
+                stride=stride,
+                corr_radius=corr_radius,
+                corr_levels=corr_levels,
+                backbone=backbone,
+                num_heads=num_heads,
+                hidden_size=hidden_size,
+                space_depth=space_depth,
+                time_depth=time_depth
+            )
+        
+        time_grid = torch.linspace(0, self.window_size - 1, self.window_size).reshape(1, self.window_size, 1)
+        self.register_buffer(
+            "time_emb", get_1d_sincos_pos_embed_from_grid(self.input_dim, time_grid[0])
+        )
+        self.corr_pyramid = []
+        
+    @torch.no_grad()  
+    def forward(self, rgbs, events, queries, iters=6, img_ifnew=None, feat_init=None, interp_shape=(384, 512), is_train=False):
+        # starter.record()
+        if self.backbone == "image":
+            self.updateformer2 = self.updateformer
+        B, T, C, H, W = events.shape
+        _, N, _ = queries.shape
+        _, T, C_img, _, _ = rgbs.shape
+        S = self.window_size
+        step = S // 2
+        device = queries.device
+        
+        queried_frames = queries[:, :, 0].long()
+        queried_coords = queries[..., 1:3]
+        queried_coords = queried_coords / self.stride
+        
+        coords_predicted = torch.zeros((B, T, N, 2), device=device)
+        vis_predicted= torch.zeros((B, T, N), device=device)
+        conf_predicted = torch.zeros((B, T, N), device=device)
+        
+        H_stride, W_stride = interp_shape[0] // self.stride, interp_shape[1] // self.stride
+        
+        coords_init = queries[:, :, 1:].reshape(B, 1, N, 2).repeat(1, self.window_size, 1, 1) / float(self.stride)
+        
+        vis_init = torch.zeros((B, S, N, 1), device=device).float()
+        conf_init = torch.zeros((B, S, N, 1), device=device).float()
+        coords_init = queried_coords.reshape(B, 1, N, 2).expand(B, S, N, 2).float()
+        
+        num_windows =  (T - S + step - 1) // step + 1
+        indices = range(0, step * num_windows, step)
+        
+        fmaps_fusion = None
+        first_window = True
+        track_feat_pyramid = []
+        track_feat_support_pyramid = []
+        
+        for ind in indices:
+            if ind > 0:
+                overlap = S - step
+                copy_over = (queried_frames < ind + overlap)[:, None, :, None]  # B 1 N 1
+                coords_prev = coords_predicted[:, ind : ind + overlap] / self.stride
+                padding_tensor = coords_prev[:, -1:, :, :].expand(-1, step, -1, -1)
+                coords_prev = torch.cat([coords_prev, padding_tensor], dim=1)
+
+                vis_prev = vis_predicted[:, ind : ind + overlap, :, None].clone()
+                padding_tensor = vis_prev[:, -1:, :, :].expand(-1, step, -1, -1)
+                vis_prev = torch.cat([vis_prev, padding_tensor], dim=1)
+
+                conf_prev = conf_predicted[:, ind : ind + overlap, :, None].clone()
+                padding_tensor = conf_prev[:, -1:, :, :].expand(-1, step, -1, -1)
+                conf_prev = torch.cat([conf_prev, padding_tensor], dim=1)
+
+                coords_init = torch.where(copy_over.expand_as(coords_init), coords_prev, coords_init)
+                vis_init = torch.where(copy_over.expand_as(vis_init), vis_prev, vis_init)
+                conf_init = torch.where(copy_over.expand_as(conf_init), conf_prev, conf_init)
+            
+            events_seq = events[:, ind : ind + S]
+            rgbs_seq = rgbs[:, ind : ind + S]
+            if img_ifnew is not None:
+                img_ifnew_seq = img_ifnew[ind : ind + S]
+            if not isinstance(events_seq, torch.Tensor):
+                events_seq = torch.from_numpy(events_seq).to(device).float()
+                rgbs_seq = torch.from_numpy(rgbs_seq).to(device).float()
+                events_seq = events_seq.contiguous()
+                rgbs_seq = rgbs_seq.contiguous()
+            
+            if ind + S > T:
+                pad = (S - rgbs_seq.shape[1]) % S
+                rgbs_seq = rgbs_seq.reshape(B, 1, (S - pad), C_img * H * W)
+                events_seq = events_seq.reshape(B, 1, (S - pad), C * H * W)
+                padding_tensor = rgbs_seq[:, :, -1:, :].expand(B, 1, pad, C_img * H * W)
+                rgbs_seq = torch.cat([rgbs_seq, padding_tensor], dim=2)
+                padding_tensor = events_seq[:, :, -1:, :].expand(B, 1, pad, C * H * W)
+                events_seq = torch.cat([events_seq, padding_tensor], dim=2)
+                if img_ifnew is not None:
+                    padding_numpy = np.ones(pad)
+                    img_ifnew_seq = np.concatenate((img_ifnew_seq, padding_numpy), axis=0)
+                # rgbs_seq = rgbs_seq.reshape(B, -1, C_img, H, W)
+                # events_seq = events_seq.reshape(B, -1, C, H, W)
+            
+            events_seq = events_seq.reshape(B * S, C, H, W)
+            rgbs_seq = rgbs_seq.reshape(B * S, C_img, H, W)
+            events_seq = F.interpolate(events_seq, tuple(interp_shape), mode='bilinear', align_corners=True)
+            rgbs_seq = F.interpolate(rgbs_seq, tuple(interp_shape), mode='bilinear', align_corners=True)
+            
+            events_seq = 2 * events_seq - 1.0
+            rgbs_seq = 2 * (rgbs_seq / 255.0) - 1.0
+            
+            dtype = rgbs_seq.dtype
+            
+            if fmaps_fusion is None:
+                fmaps_pyramid = self.fusion_block(rgbs_seq, events_seq, img_ifnew_seq if img_ifnew is not None else None)
+                
+                if isinstance(fmaps_pyramid, list):
+                    for i, fmaps_fusion in enumerate(fmaps_pyramid):
+                        fmaps_fusion = fmaps_fusion.permute(0, 2, 3, 1)
+                        fmaps_fusion = fmaps_fusion / torch.sqrt(
+                            torch.maximum(
+                                torch.sum(torch.square(fmaps_fusion), axis=-1, keepdims=True),
+                                torch.tensor(1e-12, device=fmaps_fusion.device),
+                            )
+                        )
+                        fmaps_fusion = fmaps_fusion.permute(0, 3, 1, 2).reshape(
+                            B, -1, self.latent_dim, int(H_stride / 2**i), int(W_stride / 2**i)
+                        )
+                        fmaps_fusion = fmaps_fusion.to(dtype)
+                        fmaps_pyramid[i] = fmaps_fusion
+                else:
+                    fmaps_fusion = fmaps_pyramid.permute(0, 2, 3, 1)
+                    fmaps_fusion = fmaps_fusion / torch.sqrt(
+                        torch.maximum(
+                            torch.sum(torch.square(fmaps_fusion), axis=-1, keepdims=True),
+                            torch.tensor(1e-12, device=fmaps_fusion.device),
+                        )
+                    )
+                    fmaps_fusion = fmaps_fusion.permute(0, 3, 1, 2).reshape(
+                        B, -1, self.latent_dim, H_stride, W_stride
+                    )
+            else:
+                rgbs_, events_ = rgbs_seq[S//2:], events_seq[S//2:]
+                fmaps_pyramid_last = self.fusion_block(rgbs_, events_, img_ifnew_seq[S//2:] if img_ifnew is not None else None)
+                
+                if isinstance(fmaps_pyramid_last, list):
+                    for i, fmaps_fusion_last in enumerate(fmaps_pyramid_last):
+                        fmaps_fusion_last = fmaps_fusion_last.permute(0, 2, 3, 1)
+                        fmaps_fusion_last = fmaps_fusion_last / torch.sqrt(
+                            torch.maximum(
+                                torch.sum(torch.square(fmaps_fusion_last), axis=-1, keepdims=True),
+                                torch.tensor(1e-12, device=fmaps_fusion.device),
+                            )
+                        )
+                        fmaps_fusion_last = fmaps_fusion_last.permute(0, 3, 1, 2).reshape(
+                            B, -1, self.latent_dim, int(H_stride / 2**i), int(W_stride / 2**i)
+                        )
+                        fmaps_fusion = torch.cat([fmaps_pyramid[i][:, S//2:], fmaps_fusion_last], dim=1)
+                        fmaps_fusion = fmaps_fusion.to(dtype)
+                        fmaps_pyramid[i] = fmaps_fusion
+                else:
+                    fmaps_fusion_last = fmaps_pyramid_last.permute(0, 2, 3, 1)
+                    fmaps_fusion_last = fmaps_fusion_last / torch.sqrt(
+                        torch.maximum(
+                            torch.sum(torch.square(fmaps_fusion_last), axis=-1, keepdims=True),
+                            torch.tensor(1e-12, device=fmaps_fusion_last.device),
+                        )
+                    )
+                    fmaps_fusion_last = fmaps_fusion_last.permute(0, 3, 1, 2).reshape(
+                        B, -1, self.latent_dim, H_stride, W_stride
+                    )
+                    fmaps_fusion = torch.cat([fmaps_pyramid[0][:, S//2:], fmaps_fusion_last], dim=1)
+                    fmaps_fusion = fmaps_fusion.to(dtype)
+                    fmaps_pyramid = None
+            
+            if not isinstance(fmaps_pyramid, list):
+                fmaps_pyramid = []
+                fmaps_pyramid.append(fmaps_fusion)
+                for i in range(self.corr_levels - 1):
+                    fmaps_ = fmaps_fusion.reshape(B * S, self.latent_dim, fmaps_fusion.shape[-2], fmaps_fusion.shape[-1])
+                    fmaps_ = F.avg_pool2d(fmaps_, 2, stride=2)
+                    fmaps_fusion = fmaps_.reshape(B, S, self.latent_dim, fmaps_.shape[-2], fmaps_.shape[-1])
+                    fmaps_pyramid.append(fmaps_fusion)
+            if first_window:
+                for i in range(self.corr_levels):
+                    track_feat, track_feat_support = get_track_feat(fmaps_pyramid[i], queried_frames, queried_coords/2**i, support_radius=self.corr_radius)
+                    track_feat_pyramid.append(track_feat.repeat(1, S, 1, 1))
+                    track_feat_support_pyramid.append(track_feat_support.unsqueeze(1))
+                first_window = False
+                
+                
+            attenstion_mask = (queried_frames < ind + S).reshape(B, 1, N) # B, 1, N
+            coords, viss, confs = self.forward_window(
+                fmaps_pyramid=[fmap for fmap in fmaps_pyramid],
+                coords=coords_init,
+                track_feat_support_pyramid=[attenstion_mask[:, None, :, :, None]*tfeat for tfeat in track_feat_support_pyramid],
+                corr_map_pyramid=self.corr_pyramid,
+                vis=vis_init,
+                conf=conf_init,
+                attenstion_mask=attenstion_mask.repeat(1, S, 1),
+                iters=iters,
+            )
+            S_trimmed = min(T - ind, S)     # accounts for last window duration
+            coords_predicted[:, ind : ind + S] = coords[-1][:, :S_trimmed]
+            vis_predicted[:, ind : ind + S] = viss[-1][:, :S_trimmed]
+            conf_predicted[:, ind : ind + S] = confs[-1][:, :S_trimmed]
+            
+        vis_predicted = torch.sigmoid(vis_predicted)
+        conf_predicted = torch.sigmoid(conf_predicted)
+        return coords_predicted, vis_predicted, conf_predicted
+    
+    def forward_window(self, fmaps_pyramid, coords, track_feat_support_pyramid, corr_map_pyramid, vis, conf, attenstion_mask, iters=6):
+        B, S, *_ = fmaps_pyramid[0].shape
+        N = coords.shape[2]
+        r = 2 * self.corr_radius + 1
+        
+        coord_preds, vis_preds, conf_preds = [], [], []
+        for it in range(iters):
+            coords = coords.detach()
+            coord_init = coords.view(B * S, N, 2)
+            corr_embs = []
+            for i in range(self.corr_levels):
+                corr_feat = self.get_correlation_feat(fmaps_pyramid[i], coord_init / 2 ** i)
+                track_feat_support = (
+                    track_feat_support_pyramid[i]
+                    .view(B, 1, r, r, N, self.latent_dim)
+                    .squeeze(1)
+                    .permute(0, 3, 1, 2, 4)
+                )
+                corr_volume = torch.einsum("btnhwc,bnijc->btnhwij", corr_feat, track_feat_support).reshape(B, S, N, r * r, r * r)
+                corr_emb = self.corr_mlp(corr_volume.reshape(B * S * N, r * r * r *r))
+                corr_embs.append(corr_emb)      
+                
+            corr_embs = torch.cat(corr_embs, dim=1)
+            corr_embs = corr_embs.view(B, S, N, corr_embs.shape[-1])
+            
+            transformer_input = [vis, conf, corr_embs]
+            
+            rel_coords_forward = coords[:, :-1] - coords[:, 1:]
+            rel_coords_backward = coords[:, 1:] - coords[:, :-1]
+            
+            rel_coords_forward = torch.nn.functional.pad(rel_coords_forward, (0, 0, 0, 0, 0, 1))
+            rel_coords_backward = torch.nn.functional.pad(rel_coords_backward, (0, 0, 0, 0, 1, 0))
+            
+            scale = (torch.tensor([self.model_resolution[1], self.model_resolution[0]], device=coords.device,) / self.stride)
+            rel_coords_forward = rel_coords_forward / scale
+            rel_coords_backward = rel_coords_backward / scale
+            
+            rel_pos_emb_input = posenc(torch.cat([rel_coords_forward, rel_coords_backward], dim=-1), min_deg=0, max_deg=10,)
+            transformer_input.append(rel_pos_emb_input)
+            
+            x = (torch.cat(transformer_input, dim=-1).permute(0, 2, 1, 3).reshape(B*N, S, -1))
+            
+            x = x + self.interpolate_time_embed(x, S)
+            x = x.view(B, N, S, -1)
+            delta = self.updateformer2(x)
+            
+            delta_coords = delta[..., :2].permute(0, 2, 1, 3)
+            delta_vis = delta[..., 2:3].permute(0, 2, 1, 3)
+            delta_conf = delta[..., 3:].permute(0, 2, 1, 3)
+            
+            vis = vis + delta_vis
+            conf = conf + delta_conf
+
+            coords = coords + delta_coords
+            coord_preds.append(coords[..., :2] * float(self.stride))
+
+            vis_preds.append(vis[..., 0])
+            conf_preds.append(conf[..., 0])
+        return coord_preds, vis_preds, conf_preds
+            
+    
+    def load_parameters(self, model):
+        # 从训练模型加载参数
+        self.load_state_dict(model.state_dict())

LFE_TAP/models/blocks.py

@@ -0,0 +1,994 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import collections
+
+from typing import Callable
+from itertools import repeat
+from functools import partial
+from einops import rearrange
+from timm.models.vision_transformer import Attention, Mlp
+from LFE_TAP.utils.model_utils import combine_tokens, recover_tokens, bilinear_sampler
+
+
+# From PyTorch internals
+def _ntuple(n):
+    def parse(x):
+        if isinstance(x, collections.abc.Iterable) and not isinstance(x, str):
+            return tuple(x)
+        return tuple(repeat(x, n))
+
+    return parse
+
+def exists(val):
+    return val is not None
+
+def default(val, d):
+    return val if exists(val) else d
+
+to_2tuple = _ntuple(2)
+class Mlp(nn.Module):
+    """MLP as used in Vision Transformer, MLP-Mixer and related networks"""
+
+    def __init__(
+        self,
+        in_features,
+        hidden_features=None,
+        out_features=None,
+        act_layer=nn.GELU,
+        norm_layer=None,
+        bias=True,
+        drop=0.0,
+        use_conv=False,
+    ):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        bias = to_2tuple(bias)
+        drop_probs = to_2tuple(drop)
+        linear_layer = partial(nn.Conv2d, kernel_size=1) if use_conv else nn.Linear
+
+        self.fc1 = linear_layer(in_features, hidden_features, bias=bias[0])
+        self.act = act_layer()
+        self.drop1 = nn.Dropout(drop_probs[0])
+        self.norm = (
+            norm_layer(hidden_features) if norm_layer is not None else nn.Identity()
+        )
+        self.fc2 = linear_layer(hidden_features, out_features, bias=bias[1])
+        self.drop2 = nn.Dropout(drop_probs[1])
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop1(x)
+        x = self.fc2(x)
+        x = self.drop2(x)
+        return x
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, in_planes, planes, norm_fn="group", stride=1, dilation=1):
+        super(ResidualBlock, self).__init__()
+
+        self.conv1 = nn.Conv2d(
+            in_planes,
+            planes,
+            kernel_size=3,
+            padding=dilation,
+            dilation=dilation,
+            stride=stride,
+            padding_mode="zeros",
+        )
+        self.conv2 = nn.Conv2d(
+            planes, planes, kernel_size=3, padding=dilation, dilation=dilation, padding_mode="zeros"
+        )
+        self.relu = nn.ReLU(inplace=True)
+
+        num_groups = planes // 8
+
+        if norm_fn == "group":
+            self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
+            self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
+            if not (in_planes == planes and stride == 1):
+                self.norm3 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
+
+        elif norm_fn == "batch":
+            self.norm1 = nn.BatchNorm2d(planes)
+            self.norm2 = nn.BatchNorm2d(planes)
+            if not (in_planes == planes and stride == 1):
+                self.norm3 = nn.BatchNorm2d(planes)
+
+        elif norm_fn == "instance":
+            self.norm1 = nn.InstanceNorm2d(planes)
+            self.norm2 = nn.InstanceNorm2d(planes)
+            if not (in_planes == planes and stride == 1):
+                self.norm3 = nn.InstanceNorm2d(planes)
+
+        elif norm_fn == "none":
+            self.norm1 = nn.Sequential()
+            self.norm2 = nn.Sequential()
+            if not (in_planes == planes and stride == 1):
+                self.norm3 = nn.Sequential()
+
+        if in_planes == planes and stride == 1:
+            self.downsample = None
+
+        else:
+            self.downsample = nn.Sequential(
+                nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride), self.norm3
+            )
+
+    def forward(self, x):
+        y = x
+        y = self.relu(self.norm1(self.conv1(y)))
+        y = self.relu(self.norm2(self.conv2(y)))
+
+        if self.downsample is not None:
+            x = self.downsample(x)
+
+        return self.relu(x + y)
+    
+
+class BasicEncoder(nn.Module):
+    def __init__(
+        self, input_dim=3, output_dim=128, stride=8, norm_fn="instance", dropout=0.0, shallow=False, in_planes=64, dilation=1,
+    ):
+        super(BasicEncoder, self).__init__()
+        self.stride = stride
+        self.norm_fn = norm_fn
+        self.in_planes = in_planes
+
+        if self.norm_fn == "group":
+            self.norm1 = nn.GroupNorm(num_groups=8, num_channels=self.in_planes)
+            self.norm2 = nn.GroupNorm(num_groups=8, num_channels=output_dim * 2)
+
+        elif self.norm_fn == "batch":
+            self.norm1 = nn.BatchNorm2d(self.in_planes)
+            self.norm2 = nn.BatchNorm2d(output_dim * 2)
+
+        elif self.norm_fn == "instance":
+            self.norm1 = nn.InstanceNorm2d(self.in_planes)
+            self.norm2 = nn.InstanceNorm2d(output_dim * 2)
+
+        elif self.norm_fn == "none":
+            self.norm1 = nn.Sequential()
+
+        self.conv1 = nn.Conv2d(
+            input_dim,
+            self.in_planes,
+            kernel_size=7,
+            stride=2,
+            padding=3,
+            padding_mode="zeros",
+        )
+        self.relu1 = nn.ReLU(inplace=True)
+
+        self.shallow = shallow
+        if self.shallow:
+            # self.layer1 = ResidualBlock(self.in_planes, 64, norm_fn=self.norm_fn, stride=1)
+            # self.layer2 = ResidualBlock(64, 96, self.norm_fn, stride=2)
+            # self.layer3 = ResidualBlock(96, 128, self.norm_fn, stride=2)
+            self.layer1 = self._make_layer(64, stride=1, dilation=dilation)
+            self.layer2 = self._make_layer(96, stride=2, dilation=dilation)
+            self.layer3 = self._make_layer(128, stride=2, dilation=dilation)
+            self.conv2 = nn.Conv2d(128 + 96 + 64, output_dim, kernel_size=1)
+        else:
+            self.layer1 = self._make_layer(64, stride=1)
+            self.layer2 = self._make_layer(96, stride=2)
+            self.layer3 = self._make_layer(128, stride=2)
+            self.layer4 = self._make_layer(128, stride=2)
+
+            self.conv2 = nn.Conv2d(
+                128 + 128 + 96 + 64,
+                output_dim * 2,
+                kernel_size=3,
+                padding=1,
+                padding_mode="zeros",
+            )
+            self.relu2 = nn.ReLU(inplace=True)
+            self.conv3 = nn.Conv2d(output_dim * 2, output_dim, kernel_size=1)
+
+        self.dropout = None
+        if dropout > 0:
+            self.dropout = nn.Dropout2d(p=dropout)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
+            elif isinstance(m, (nn.BatchNorm2d, nn.InstanceNorm2d, nn.GroupNorm)):
+                if m.weight is not None:
+                    nn.init.constant_(m.weight, 1)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+    def _make_layer(self, dim, stride=1, dilation=1):
+        layer1 = ResidualBlock(self.in_planes, dim, self.norm_fn, stride=stride, dilation=dilation)
+        layer2 = ResidualBlock(dim, dim, self.norm_fn, stride=1, dilation=dilation)
+        layers = (layer1, layer2)
+
+        self.in_planes = dim
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        _, _, H, W = x.shape
+
+        x = self.conv1(x)
+        x = self.norm1(x)
+        x = self.relu1(x)
+
+        if self.shallow:
+            a = self.layer1(x)
+            b = self.layer2(a)
+            c = self.layer3(b)
+            a = F.interpolate(
+                a,
+                (H // self.stride, W // self.stride),
+                mode="bilinear",
+                align_corners=True,
+            )
+            b = F.interpolate(
+                b,
+                (H // self.stride, W // self.stride),
+                mode="bilinear",
+                align_corners=True,
+            )
+            c = F.interpolate(
+                c,
+                (H // self.stride, W // self.stride),
+                mode="bilinear",
+                align_corners=True,
+            )
+            x = self.conv2(torch.cat([a, b, c], dim=1))
+        else:
+            a = self.layer1(x)
+            b = self.layer2(a)
+            c = self.layer3(b)
+            d = self.layer4(c)
+            a = F.interpolate(
+                a,
+                (H // self.stride, W // self.stride),
+                mode="bilinear",
+                align_corners=True,
+            )
+            b = F.interpolate(
+                b,
+                (H // self.stride, W // self.stride),
+                mode="bilinear",
+                align_corners=True,
+            )
+            c = F.interpolate(
+                c,
+                (H // self.stride, W // self.stride),
+                mode="bilinear",
+                align_corners=True,
+            )
+            d = F.interpolate(
+                d,
+                (H // self.stride, W // self.stride),
+                mode="bilinear",
+                align_corners=True,
+            )
+            x = self.conv2(torch.cat([a, b, c, d], dim=1))
+            x = self.norm2(x)
+            x = self.relu2(x)
+            x = self.conv3(x)
+
+        if self.training and self.dropout is not None:
+            x = self.dropout(x)
+        return x
+
+
+class FusionBlock_basic(nn.Module):
+    def __init__(self, img_in_dim=3, event_in_dim=10, output_dim=128, stride=8, norm_fn="instance", dropout=0.0):
+        super().__init__()
+        self.imgnet = BasicEncoder(input_dim=img_in_dim, output_dim=output_dim, stride=stride, norm_fn=norm_fn, dropout=dropout)
+        self.eventnet = BasicEncoder(input_dim=event_in_dim, output_dim=output_dim, stride=stride, norm_fn=norm_fn, dropout=dropout)
+        self.conv1 = nn.Conv2d(output_dim, 192, 1, padding=0)
+        self.conv2 = nn.Conv2d(output_dim, 192, 1, padding=0)
+        self.convo = nn.Conv2d(192*2, output_dim, 3, padding=1)
+
+        # for m in self.modules():
+        #     if isinstance(m, nn.Conv2d):
+        #         nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
+                
+    def forward(self, x_i, x_e, _):
+        x_i = self.imgnet(x_i)
+        x_e = self.eventnet(x_e)
+        c1 = F.relu(self.conv1(x_i))
+        c2 = F.relu(self.conv2(x_e))
+        out = torch.cat([c1, c2], dim=1)
+        out = F.relu(self.convo(out))
+        return x_i + out
+    
+class FusionBlock(nn.Module):
+    def __init__(self, img_in_dim=3, event_in_dim=10, output_dim=128, stride=8, norm_fn="instance", dropout=0.0):
+        super(FusionBlock, self).__init__()
+        self.stride = stride
+        self.norm_fn = norm_fn
+        self.in_planes = 32
+        
+        if self.norm_fn == "group":
+            self.norm1_i = nn.GroupNorm(num_groups=8, num_channels=self.in_planes)
+            self.norm1_e = nn.GroupNorm(num_groups=8, num_channels=self.in_planes)
+            self.norm2 = nn.GroupNorm(num_groups=8, num_channels=output_dim * 2)
+
+        elif self.norm_fn == "batch":
+            self.norm1_i = nn.BatchNorm2d(self.in_planes)
+            self.norm1_e = nn.BatchNorm2d(self.in_planes)
+            self.norm2 = nn.BatchNorm2d(output_dim * 2)
+
+        elif self.norm_fn == "instance":
+            self.norm1_i = nn.InstanceNorm2d(self.in_planes)
+            self.norm1_e = nn.InstanceNorm2d(self.in_planes)
+            self.norm2 = nn.InstanceNorm2d(output_dim * 2)
+
+        elif self.norm_fn == "none":
+            self.norm1_i = nn.Sequential()
+            self.norm1_e = nn.Sequential()
+            
+        self.conv1_i = nn.Conv2d(img_in_dim, self.in_planes, kernel_size=7, stride=2, padding=3, padding_mode="zeros")
+        self.conv1_e = nn.Conv2d(event_in_dim, self.in_planes, kernel_size=7, stride=2, padding=3, padding_mode="zeros")
+        self.relu1 = nn.ReLU(inplace=True)
+        
+        self.shallow = False
+        if self.shallow:
+            self.layer1_e = self._make_layer(self.in_planes, 64, stride=1)
+            self.layer2_e = self._make_layer(64, 96, stride=2)
+            self.layer3_e = self._make_layer(96, 128, stride=2)
+            # self.conv_half1 = self._half_conv(self.in_planes*2)
+            # self.conv_half2 = self._half_conv(64*2)
+            # self.conv_half3 = self._half_conv(96*2)
+            self.layer1_i = self._make_layer(self.in_planes*2, 64, stride=1)
+            self.layer2_i = self._make_layer(64*2, 96, stride=2)
+            self.layer3_i = self._make_layer(96*2, 128, stride=2)
+            self.conv2 = nn.Conv2d(128 + 96 + 64, output_dim, kernel_size=1)
+        else:
+            self.layer1_e = self._make_layer(self.in_planes, 64, stride=1)
+            self.layer2_e = self._make_layer(64, 96, stride=2)
+            self.layer3_e = self._make_layer(96, 128, stride=2)
+            self.layer4_e = self._make_layer(128, 128, stride=2)
+            # self.conv_half1 = self._half_conv(self.in_planes*2)
+            # self.conv_half2 = self._half_conv(64*2)
+            # self.conv_half3 = self._half_conv(96*2)
+            # self.conv_half4 = self._half_conv(128*2)
+            self.layer1_i = self._make_layer(self.in_planes*2, 64, stride=1)
+            self.layer2_i = self._make_layer(64*2, 96, stride=2)
+            self.layer3_i = self._make_layer(96*2, 128, stride=2)
+            self.layer4_i = self._make_layer(128*2, 128, stride=2)
+            self.conv2 = nn.Conv2d(128 + 128 + 96 + 64, output_dim, kernel_size=3, padding=1, padding_mode="zeros")
+            self.relu2 = nn.ReLU(inplace=True)
+            self.conv3 = nn.Conv2d(output_dim , output_dim, kernel_size=1)
+            
+        self.dropout = None
+        if dropout > 0:
+            self.dropout = nn.Dropout2d(p=dropout)
+            
+    def _half_conv(self, in_planes):
+        conv = nn.Conv2d(in_planes, in_planes // 2, kernel_size=1)
+        norm = nn.BatchNorm2d(in_planes // 2)
+        relu = nn.ReLU(inplace=True)
+        layers = (conv, norm, relu)
+        return nn.Sequential(*layers)
+        
+    def _make_layer(self, input_dim, output_dim, stride=1):
+        layer1 = ResidualBlock(input_dim, output_dim, self.norm_fn, stride=stride)
+        layer2 = ResidualBlock(output_dim, output_dim, self.norm_fn, stride=1)
+        layers = (layer1, layer2)
+
+        return nn.Sequential(*layers)
+    
+    def forward(self, x_i, x_e):
+        _, _, H, W = x_i.size()
+        
+        x_i = self.conv1_i(x_i)
+        x_e = self.conv1_e(x_e)
+        x_i = self.relu1(self.norm1_i(x_i))
+        x_e = self.relu1(self.norm1_e(x_e))
+        
+        if self.shallow:
+            x_e_a = self.layer1_e(x_e)
+            x_e_b = self.layer2_e(x_e_a)
+            x_e_c = self.layer3_e(x_e_b)
+            # x_i_a = self.layer1_i(self.conv_half1(torch.cat((x_i, x_e), dim=1)))
+            # x_i_b = self.layer2_i(self.conv_half2(torch.cat((x_i_a, x_e_a), dim=1)))
+            # x_i_c = self.layer3_i(self.conv_half3(torch.cat((x_i_b, x_e_b), dim=1)))
+            x_i_a = self.layer1_i(torch.cat((x_i, x_e), dim=1))
+            x_i_b = self.layer2_i(torch.cat((x_i_a, x_e_a), dim=1))
+            x_i_c = self.layer3_i(torch.cat((x_i_b, x_e_b), dim=1))
+            x_i_a = F.interpolate(x_i_a, size=(H // self.stride, W // self.stride), mode='bilinear', align_corners=True)
+            x_i_b = F.interpolate(x_i_b, size=(H // self.stride, W // self.stride), mode='bilinear', align_corners=True)
+            x_i_c = F.interpolate(x_i_c, size=(H // self.stride, W // self.stride), mode='bilinear', align_corners=True)
+            x = self.conv2(torch.cat((x_i_a, x_i_b, x_i_c), dim=1))
+        else:
+            x_e_a = self.layer1_e(x_e)
+            x_e_b = self.layer2_e(x_e_a)
+            x_e_c = self.layer3_e(x_e_b)
+            # x_e_d = self.layer4_e(x_e_c)
+            # x_i_a = self.layer1_i(self.conv_half1(torch.cat((x_i, x_e), dim=1)))
+            # x_i_b = self.layer2_i(self.conv_half2(torch.cat((x_i_a, x_e_a), dim=1)))
+            # x_i_c = self.layer3_i(self.conv_half3(torch.cat((x_i_b, x_e_b), dim=1)))
+            # x_i_d = self.layer4_i(self.conv_half4(torch.cat((x_i_c, x_e_c), dim=1)))
+            x_i_a = self.layer1_i(torch.cat((x_i, x_e), dim=1))
+            x_i_b = self.layer2_i(torch.cat((x_i_a, x_e_a), dim=1))
+            x_i_c = self.layer3_i(torch.cat((x_i_b, x_e_b), dim=1))
+            x_i_d = self.layer4_i(torch.cat((x_i_c, x_e_c), dim=1))
+            x_i_a = F.interpolate(x_i_a, size=(H // self.stride, W // self.stride), mode='bilinear', align_corners=True)
+            x_i_b = F.interpolate(x_i_b, size=(H // self.stride, W // self.stride), mode='bilinear', align_corners=True)
+            x_i_c = F.interpolate(x_i_c, size=(H // self.stride, W // self.stride), mode='bilinear', align_corners=True)
+            x_i_d = F.interpolate(x_i_d, size=(H // self.stride, W // self.stride), mode='bilinear', align_corners=True)
+            x = self.conv2(torch.cat((x_i_a, x_i_b, x_i_c, x_i_d), dim=1))
+            x = self.relu2(self.norm2(x))
+            x = self.conv3(x)
+
+        if self.training and self.dropout is not None:
+            x = self.dropout(x)
+            
+        return x
+
+
+class AttnBlock(nn.Module):
+    """
+    A DiT block with adaptive layer norm zero (adaLN-Zero) conditioning.
+    """
+
+    def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, **block_kwargs):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.attn = Attention(
+            hidden_size, num_heads=num_heads, qkv_bias=True, **block_kwargs
+        )
+
+        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        approx_gelu = lambda: nn.GELU(approximate="tanh")
+        self.mlp = Mlp(
+            in_features=hidden_size,
+            hidden_features=mlp_hidden_dim,
+            act_layer=approx_gelu,
+            drop=0,
+        )
+
+    def forward(self, x):
+        x = x + self.attn(self.norm1(x))
+        x = x + self.mlp(self.norm2(x))
+        return x
+    
+
+class UpdateFormer(nn.Module):
+    def __init__(self, space_depth=12, time_depth=12, input_dim=320, hidden_size=384, num_heads=8, output_dim=130, mlp_ratio=4.0):
+        super(UpdateFormer, self).__init__()
+        self.hidden_size = hidden_size
+        self.input_transform = nn.Linear(input_dim, hidden_size, bias=True)
+        self.flow_head = nn.Linear(hidden_size, output_dim, bias=True)
+        
+        self.time_blocks = nn.ModuleList(
+            [
+                AttnBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio)
+                for _ in range(time_depth)
+            ]
+        )
+        
+        self.space_blocks = nn.ModuleList(
+            [
+                AttnBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio)
+                for _ in range(space_depth)
+            ]
+        )
+        assert len(self.time_blocks) >= len(self.space_blocks)
+        self.initialize_weights()
+        
+    def initialize_weights(self):
+        def _basic_init(module):
+            if isinstance(module, nn.Linear):
+                torch.nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.constant_(module.bias, 0)
+
+        self.apply(_basic_init)
+        
+    def forward(self, x):
+        x = self.input_transform(x)
+        j = 0
+        for i in range(len(self.time_blocks)):
+            B, N, T, _ = x.shape
+            x_time = rearrange(x, "b n t c -> (b n) t c", b=B, t=T, n=N)
+            x_time = self.time_blocks[i](x_time)
+            
+            x = rearrange(x_time, "(b n) t c -> b n t c", b=B, t=T, n=N)
+            if self.add_space_attn and (i % (len(self.time_blocks) // len(self.space_blocks)) == 0):
+                x_space = rearrange(x, "b n t c -> (b n) t c", b=B, t=T, n=N)
+                x_space = self.space_blocks[j](x_space)
+                x = rearrange(x_space, "(b n) t c -> b n t c", b=B, t=T, n=N)
+                j += 1
+                
+        flow = self.flow_head(x)
+        return flow
+
+
+class Attention2(nn.Module):
+    def __init__(
+        self, query_dim, context_dim=None, num_heads=8, dim_head=48, qkv_bias=False
+    ):
+        super().__init__()
+        inner_dim = dim_head * num_heads
+        context_dim = default(context_dim, query_dim)
+        self.scale = dim_head**-0.5
+        self.heads = num_heads
+
+        self.to_q = nn.Linear(query_dim, inner_dim, bias=qkv_bias)
+        self.to_kv = nn.Linear(context_dim, inner_dim * 2, bias=qkv_bias)
+        self.to_out = nn.Linear(inner_dim, query_dim)
+
+    def forward(self, x, context=None, attn_bias=None):
+        B, N1, C = x.shape
+        h = self.heads
+
+        q = self.to_q(x).reshape(B, N1, h, C // h).permute(0, 2, 1, 3)
+        context = default(context, x)
+        k, v = self.to_kv(context).chunk(2, dim=-1)
+
+        N2 = context.shape[1]
+        k = k.reshape(B, N2, h, C // h).permute(0, 2, 1, 3)
+        v = v.reshape(B, N2, h, C // h).permute(0, 2, 1, 3)
+
+        sim = (q @ k.transpose(-2, -1)) * self.scale
+
+        if attn_bias is not None:
+            sim = sim + attn_bias
+        attn = sim.softmax(dim=-1)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N1, C)
+        return self.to_out(x)
+
+ 
+class CrossAttnBlock(nn.Module):
+    def __init__(
+        self, hidden_size, context_dim, num_heads=1, mlp_ratio=4.0, **block_kwargs
+    ):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.norm_context = nn.LayerNorm(hidden_size)
+        self.cross_attn = Attention2(
+            hidden_size,
+            context_dim=context_dim,
+            num_heads=num_heads,
+            qkv_bias=True,
+            **block_kwargs
+        )
+
+        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        approx_gelu = lambda: nn.GELU(approximate="tanh")
+        self.mlp = Mlp(
+            in_features=hidden_size,
+            hidden_features=mlp_hidden_dim,
+            act_layer=approx_gelu,
+            drop=0,
+        )
+
+    def forward(self, x, context, mask=None):
+        attn_bias = None
+        if mask is not None:
+            if mask.shape[1] == x.shape[1]:
+                mask = mask[:, None, :, None].expand(
+                    -1, self.cross_attn.heads, -1, context.shape[1]
+                )
+            else:
+                mask = mask[:, None, None].expand(
+                    -1, self.cross_attn.heads, x.shape[1], -1
+                )
+
+            max_neg_value = -torch.finfo(x.dtype).max
+            attn_bias = (~mask) * max_neg_value
+        x = x + self.cross_attn(
+            self.norm1(x), context=self.norm_context(context), attn_bias=attn_bias
+        )
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+
+class AttnBlock2(nn.Module):
+    def __init__(
+        self,
+        hidden_size,
+        num_heads,
+        attn_class: Callable[..., nn.Module] = Attention2,
+        mlp_ratio=4.0,
+        **block_kwargs
+    ):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.attn = attn_class(
+            hidden_size, num_heads=num_heads, qkv_bias=True, **block_kwargs
+        )
+
+        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        approx_gelu = lambda: nn.GELU(approximate="tanh")
+        self.mlp = Mlp(
+            in_features=hidden_size,
+            hidden_features=mlp_hidden_dim,
+            act_layer=approx_gelu,
+            drop=0,
+        )
+
+    def forward(self, x, mask=None):
+        attn_bias = mask
+        if mask is not None:
+            mask = (
+                (mask[:, None] * mask[:, :, None])
+                .unsqueeze(1)
+                .expand(-1, self.attn.num_heads, -1, -1)
+            )
+            max_neg_value = -torch.finfo(x.dtype).max
+            attn_bias = (~mask) * max_neg_value
+        x = x + self.attn(self.norm1(x), attn_bias=attn_bias)
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+    
+class EfficientUpdateFormer(nn.Module):
+    """
+    Transformer model that updates track estimates.
+    """
+
+    def __init__(
+        self,
+        space_depth=6,
+        time_depth=6,
+        input_dim=320,
+        hidden_size=384,
+        num_heads=8,
+        output_dim=130,
+        mlp_ratio=4.0,
+        num_virtual_tracks=32,
+        add_space_attn=True,
+        linear_layer_for_vis_conf=False,
+    ):
+        super().__init__()
+        self.out_channels = 2
+        self.num_heads = num_heads
+        self.hidden_size = hidden_size
+        self.input_transform = torch.nn.Linear(input_dim, hidden_size, bias=True)
+        if linear_layer_for_vis_conf:
+            self.flow_head = torch.nn.Linear(hidden_size, output_dim - 2, bias=True)
+            self.vis_conf_head = torch.nn.Linear(hidden_size, 2, bias=True)
+        else:
+            self.flow_head = torch.nn.Linear(hidden_size, output_dim, bias=True)
+        self.num_virtual_tracks = num_virtual_tracks
+        self.virual_tracks = nn.Parameter(
+            torch.randn(1, num_virtual_tracks, 1, hidden_size)
+        )
+        self.add_space_attn = add_space_attn
+        self.linear_layer_for_vis_conf = linear_layer_for_vis_conf
+        self.time_blocks = nn.ModuleList(
+            [
+                AttnBlock2(
+                    hidden_size,
+                    num_heads,
+                    mlp_ratio=mlp_ratio,
+                    attn_class=Attention2,
+                )
+                for _ in range(time_depth)
+            ]
+        )
+
+        if add_space_attn:
+            self.space_virtual_blocks = nn.ModuleList(
+                [
+                    AttnBlock2(
+                        hidden_size,
+                        num_heads,
+                        mlp_ratio=mlp_ratio,
+                        attn_class=Attention2,
+                    )
+                    for _ in range(space_depth)
+                ]
+            )
+            self.space_point2virtual_blocks = nn.ModuleList(
+                [
+                    CrossAttnBlock(
+                        hidden_size, hidden_size, num_heads, mlp_ratio=mlp_ratio
+                    )
+                    for _ in range(space_depth)
+                ]
+            )
+            self.space_virtual2point_blocks = nn.ModuleList(
+                [
+                    CrossAttnBlock(
+                        hidden_size, hidden_size, num_heads, mlp_ratio=mlp_ratio
+                    )
+                    for _ in range(space_depth)
+                ]
+            )
+            assert len(self.time_blocks) >= len(self.space_virtual2point_blocks)
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        def _basic_init(module):
+            if isinstance(module, nn.Linear):
+                torch.nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.constant_(module.bias, 0)
+            torch.nn.init.trunc_normal_(self.flow_head.weight, std=0.001)
+            if self.linear_layer_for_vis_conf:
+                torch.nn.init.trunc_normal_(self.vis_conf_head.weight, std=0.001)
+
+        def _trunc_init(module):
+            """ViT weight initialization, original timm impl (for reproducibility)"""
+            if isinstance(module, nn.Linear):
+                torch.nn.init.trunc_normal_(module.weight, std=0.02)
+                if module.bias is not None:
+                    nn.init.zeros_(module.bias)
+
+        self.apply(_basic_init)
+
+    def forward(self, input_tensor, mask=None, add_space_attn=True):
+        tokens = self.input_transform(input_tensor)
+
+        B, _, T, _ = tokens.shape
+        virtual_tokens = self.virual_tracks.repeat(B, 1, T, 1)
+        tokens = torch.cat([tokens, virtual_tokens], dim=1)
+
+        _, N, _, _ = tokens.shape
+        j = 0
+        layers = []
+        for i in range(len(self.time_blocks)):
+            time_tokens = tokens.contiguous().view(B * N, T, -1)  # B N T C -> (B N) T C
+            time_tokens = self.time_blocks[i](time_tokens)
+
+            tokens = time_tokens.view(B, N, T, -1)  # (B N) T C -> B N T C
+            if (
+                add_space_attn
+                and hasattr(self, "space_virtual_blocks")
+                and (i % (len(self.time_blocks) // len(self.space_virtual_blocks)) == 0)
+            ):
+                space_tokens = (
+                    tokens.permute(0, 2, 1, 3).contiguous().view(B * T, N, -1)
+                )  # B N T C -> (B T) N C
+
+                point_tokens = space_tokens[:, : N - self.num_virtual_tracks]
+                virtual_tokens = space_tokens[:, N - self.num_virtual_tracks :]
+
+                virtual_tokens = self.space_virtual2point_blocks[j](
+                    virtual_tokens, point_tokens, mask=mask
+                )
+
+                virtual_tokens = self.space_virtual_blocks[j](virtual_tokens)
+                point_tokens = self.space_point2virtual_blocks[j](
+                    point_tokens, virtual_tokens, mask=mask
+                )
+
+                space_tokens = torch.cat([point_tokens, virtual_tokens], dim=1)
+                tokens = space_tokens.view(B, T, N, -1).permute(
+                    0, 2, 1, 3
+                )  # (B T) N C -> B N T C
+                j += 1
+        tokens = tokens[:, : N - self.num_virtual_tracks]
+
+        flow = self.flow_head(tokens)
+        if self.linear_layer_for_vis_conf:
+            vis_conf = self.vis_conf_head(tokens)
+            flow = torch.cat([flow, vis_conf], dim=-1)
+
+        return flow
+    
+    
+class BaseBackbone(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.patch_size = 16
+
+        self.pos_embed_x = None
+
+        self.return_inter = False
+
+    def finetune_track(self, img_size):
+        new_patch_size = 16
+
+        patch_pos_embed = self.absolute_pos_embed
+        patch_pos_embed = patch_pos_embed.transpose(1, 2)
+        B, E, Q = patch_pos_embed.shape
+        P_H, P_W = img_size[0] // self.patch_size, img_size[1] // self.patch_size
+        patch_pos_embed = patch_pos_embed.view(B, E, P_H, P_W)
+
+        # for search region
+        H, W = img_size
+        new_P_H, new_P_W = H // new_patch_size, W // new_patch_size
+        img_patch_pos_embed = nn.functional.interpolate(patch_pos_embed, size=(new_P_H, new_P_W), mode='bicubic',
+                                                           align_corners=False)
+        img_patch_pos_embed = img_patch_pos_embed.flatten(2).transpose(1, 2)
+
+        self.pos_embed_x = nn.Parameter(img_patch_pos_embed)
+        
+        # if self.return_inter:
+        #     for i_layer in self.fpn_stage:
+        #         if i_layer != 11:
+        #             norm_layer = partial(nn.LayerNorm, eps=1e-6)
+        #             layer = norm_layer(self.embed_dim)
+        #             layer_name = f'norm{i_layer}'
+        #             self.add_module(layer_name, layer)
+
+
+    def forward(self, x, **kwargs):
+        """
+        Joint feature extraction and relation modeling for the basic HiViT backbone.
+        Args:
+            z (torch.Tensor): template feature, [B, C, H_z, W_z]
+            x (torch.Tensor): search region feature, [B, C, H_x, W_x]
+
+        Returns:
+            x (torch.Tensor): merged template and search region feature, [B, L_z+L_x, C]
+            attn : None
+        """
+        B = x.shape[0]
+    
+        x = self.patch_embed(x)
+    
+        for blk in self.blocks[:-self.num_main_blocks]:
+            x = blk(x)
+
+        x = x[..., 0, 0, :]
+
+        x += self.pos_embed_x       # 添加位置编码
+        
+        x = self.pos_drop(x)
+
+        for blk in self.blocks[-self.num_main_blocks:]:
+            x = blk(x)
+
+        aux_dict = {"attn": None}
+        x = self.norm_(x)
+
+        return x, aux_dict
+
+
+class CorrBlock:
+    def __init__(
+        self,
+        fmaps,
+        num_levels=4,
+        radius=4,
+        multiple_track_feats=False,
+        padding_mode="zeros",
+        appearance_fact_flow_dim=None,
+    ):
+        B, S, C, H, W = fmaps.shape
+        self.S, self.C, self.H, self.W = S, C, H, W
+        self.padding_mode = padding_mode
+        self.num_levels = num_levels
+        self.radius = radius
+        self.fmaps_pyramid = []
+        self.multiple_track_feats = multiple_track_feats
+        self.appearance_fact_flow_dim = appearance_fact_flow_dim
+
+        self.fmaps_pyramid.append(fmaps)
+        for i in range(self.num_levels - 1):
+            fmaps_ = fmaps.reshape(B * S, C, H, W)
+            fmaps_ = F.avg_pool2d(fmaps_, 2, stride=2)
+            _, _, H, W = fmaps_.shape
+            fmaps = fmaps_.reshape(B, S, C, H, W)
+            self.fmaps_pyramid.append(fmaps)
+
+    def sample(self, coords):
+        r = self.radius
+        B, S, N, D = coords.shape
+        assert D == 2
+
+        H, W = self.H, self.W
+        out_pyramid = []
+        for i in range(self.num_levels):
+            corrs = self.corrs_pyramid[i]  # B, S, N, H, W
+            *_, H, W = corrs.shape
+
+            dx = torch.linspace(-r, r, 2 * r + 1)
+            dy = torch.linspace(-r, r, 2 * r + 1)
+            delta = torch.stack(torch.meshgrid(dy, dx, indexing="ij"), axis=-1).to(coords.device)
+
+            centroid_lvl = coords.reshape(B * S * N, 1, 1, 2) / 2**i
+            delta_lvl = delta.view(1, 2 * r + 1, 2 * r + 1, 2)
+            coords_lvl = centroid_lvl + delta_lvl
+
+            corrs = bilinear_sampler(
+                corrs.reshape(B * S * N, 1, H, W),
+                coords_lvl,
+                padding_mode=self.padding_mode,
+            )
+            corrs = corrs.view(B, S, N, -1)
+            out_pyramid.append(corrs)
+
+        out = torch.cat(out_pyramid, dim=-1)  # B, S, N, LRR*2
+        out = out.permute(0, 2, 1, 3).contiguous().view(B * N, S, -1).float()
+        return out
+
+    def corr(self, targets, coords=None, gt_flow=None, use_gt_flow=False, use_flow_tokens=False,
+             use_af_high_dim=False, interaction_network=None):
+        assert sum([coords is not None,
+            use_gt_flow,
+            use_flow_tokens,
+            use_af_high_dim,
+            interaction_network is not None]) <= 1, \
+            "Exactly one of coords, use_gt_flow, use_flow_tokens, use_af_high_dim, or interaction_network must be specified."
+        assert not use_flow_tokens or self.appearance_fact_flow_dim is not None
+
+        # Appearance factorization
+        if coords is not None:
+            B, S, N, D2 = targets.shape
+            targets = targets.reshape(B, S, N, 2, D2 // 2)
+            if use_gt_flow:
+                flow = gt_flow
+            else:
+                flow = coords[:, 1:] - coords[:, :-1]
+                flow = torch.cat([flow[:, 0:1], flow], dim=1)
+            flow = gt_flow if use_gt_flow else flow
+            targets = flow[..., 0:1] * targets[..., 0, :] + flow[..., 1:2] * targets[..., 1, :]
+
+            ###### DEBUG ########
+            # flow: [B, S, N, 2]
+            # targets: [B, S, N, 2, feat_dim], remember feat_dim == latent_dim / 2
+
+        # Appearane factorization with flow tokens
+        if use_flow_tokens:
+            fdim = self.appearance_fact_flow_dim
+            flow = targets[..., -fdim:]
+            targets = targets[..., :-fdim]
+            B, S, N, D2 = targets.shape
+            targets = targets.reshape(B, S, N, fdim, D2 // fdim)
+            #targets = flow[..., 0:1] * targets[..., 0, :] + flow[..., 1:2] * targets[..., 1, :]
+            targets = torch.einsum('btnji,btnj->btni', targets, flow)
+
+        if use_af_high_dim:
+            fdim = self.appearance_fact_flow_dim
+            flow = targets[..., -fdim:]
+            targets = targets[..., :-fdim]
+            targets = flow * targets
+
+        #if interaction_network is not None:
+        #    flow_feat = targets[...,  :self.appearance_fact_flow_dim]
+        #    targets = targets[..., self.appearance_fact_flow_dim:]
+
+        # Correlation
+        B, S, N, C = targets.shape
+        #assert C == self.C
+        assert S == self.S
+
+        fmap1 = targets
+
+        self.corrs_pyramid = []
+        for _, fmaps in enumerate(self.fmaps_pyramid):
+            *_, H, W = fmaps.shape
+            fmap2s = fmaps.view(B, S, C, H * W)  # B S C H W ->  B S C (H W)
+            corrs = torch.matmul(fmap1, fmap2s)
+            corrs = corrs.view(B, S, N, H, W)  # B S N (H W) -> B S N H W
+            corrs = corrs / torch.sqrt(torch.tensor(C).float())
+            self.corrs_pyramid.append(corrs)
+
+    def sample_fmap(self, coords):
+        '''Sample at coords directly from scaled feature map.
+        '''
+        B, S, N, D = coords.shape
+        assert D == 2
+
+        H, W = self.H, self.W
+        out_pyramid = []
+        for i in range(self.num_levels):
+
+            fmap = self.fmaps_pyramid[i]  # B, S, C, H, W
+            B, S, C, H, W = fmap.shape
+
+            coords_lvl = coords / 2**i
+
+            fmap = fmap.reshape(B * S, C, H, W)
+            coords_lvl = coords_lvl.reshape(B * S, N, 2)
+
+            coords_normalized = coords_lvl.clone()
+            coords_normalized[..., 0] = coords_normalized[..., 0] / (W - 1) * 2 - 1
+            coords_normalized[..., 1] = coords_normalized[..., 1] / (H - 1) * 2 - 1
+            coords_normalized = coords_normalized.unsqueeze(1)
+            feature_at_coords = F.grid_sample(fmap, coords_normalized,
+                                         mode='bilinear', align_corners=True)
+            feature_at_coords = feature_at_coords.permute(0, 2, 3, 1).view(B, S, N, C)
+            out_pyramid.append(feature_at_coords)
+
+        out = torch.cat(out_pyramid, dim=-1)  # B, S, N, LRR*2
+        return out

LFE_TAP/models/embeddings.py

@@ -0,0 +1,110 @@
+import torch
+from typing import Tuple, Union
+
+def get_1d_sincos_pos_embed_from_grid(
+    embed_dim: int, pos: torch.Tensor
+) -> torch.Tensor:
+    """
+    This function generates a 1D positional embedding from a given grid using sine and cosine functions.
+
+    Args:
+    - embed_dim: The embedding dimension.
+    - pos: The position to generate the embedding from.
+
+    Returns:
+    - emb: The generated 1D positional embedding.
+    """
+    assert embed_dim % 2 == 0
+    omega = torch.arange(embed_dim // 2, dtype=torch.double)
+    omega /= embed_dim / 2.0
+    omega = 1.0 / 10000**omega  # (D/2,)
+
+    pos = pos.reshape(-1)  # (M,)
+    out = torch.einsum("m,d->md", pos, omega)  # (M, D/2), outer product
+
+    emb_sin = torch.sin(out)  # (M, D/2)
+    emb_cos = torch.cos(out)  # (M, D/2)
+
+    emb = torch.cat([emb_sin, emb_cos], dim=1)  # (M, D)
+    return emb[None].float()
+
+def get_2d_embedding(xy: torch.Tensor, C: int, cat_coords: bool = True) -> torch.Tensor:
+    """
+    This function generates a 2D positional embedding from given coordinates using sine and cosine functions.
+
+    Args:
+    - xy: The coordinates to generate the embedding from.
+    - C: The size of the embedding.
+    - cat_coords: A flag to indicate whether to concatenate the original coordinates to the embedding.
+
+    Returns:
+    - pe: The generated 2D positional embedding.
+    """
+    B, N, D = xy.shape
+    assert D == 2
+
+    x = xy[:, :, 0:1]
+    y = xy[:, :, 1:2]
+    div_term = (
+        torch.arange(0, C, 2, device=xy.device, dtype=torch.float32) * (1000.0 / C)
+    ).reshape(1, 1, int(C / 2))
+
+    pe_x = torch.zeros(B, N, C, device=xy.device, dtype=torch.float32)
+    pe_y = torch.zeros(B, N, C, device=xy.device, dtype=torch.float32)
+
+    pe_x[:, :, 0::2] = torch.sin(x * div_term)
+    pe_x[:, :, 1::2] = torch.cos(x * div_term)
+
+    pe_y[:, :, 0::2] = torch.sin(y * div_term)
+    pe_y[:, :, 1::2] = torch.cos(y * div_term)
+
+    pe = torch.cat([pe_x, pe_y], dim=2)  # (B, N, C*3)
+    if cat_coords:
+        pe = torch.cat([xy, pe], dim=2)  # (B, N, C*3+3)
+    return pe
+
+def get_2d_sincos_pos_embed_from_grid(
+    embed_dim: int, grid: torch.Tensor
+) -> torch.Tensor:
+    """
+    This function generates a 2D positional embedding from a given grid using sine and cosine functions.
+
+    Args:
+    - embed_dim: The embedding dimension.
+    - grid: The grid to generate the embedding from.
+
+    Returns:
+    - emb: The generated 2D positional embedding.
+    """
+    assert embed_dim % 2 == 0
+
+    # use half of dimensions to encode grid_h
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)
+
+    emb = torch.cat([emb_h, emb_w], dim=2)  # (H*W, D)
+    return emb
+
+def get_2d_sincos_pos_embed(
+    embed_dim: int, grid_size: Union[int, Tuple[int, int]]
+) -> torch.Tensor:
+    """
+    This function initializes a grid and generates a 2D positional embedding using sine and cosine functions.
+    It is a wrapper of get_2d_sincos_pos_embed_from_grid.
+    Args:
+    - embed_dim: The embedding dimension.
+    - grid_size: The grid size.
+    Returns:
+    - pos_embed: The generated 2D positional embedding.
+    """
+    if isinstance(grid_size, tuple):
+        grid_size_h, grid_size_w = grid_size
+    else:
+        grid_size_h = grid_size_w = grid_size
+    grid_h = torch.arange(grid_size_h, dtype=torch.float)
+    grid_w = torch.arange(grid_size_w, dtype=torch.float)
+    grid = torch.meshgrid(grid_w, grid_h, indexing="xy")
+    grid = torch.stack(grid, dim=0)
+    grid = grid.reshape([2, 1, grid_size_h, grid_size_w])
+    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+    return pos_embed.reshape(1, grid_size_h, grid_size_w, -1).permute(0, 3, 1, 2)

LFE_TAP/models/fusionFormer.py

@@ -0,0 +1,253 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from einops.layers.torch import Rearrange
+from LFE_TAP.models.blocks import ResidualBlock, CrossAttnBlock, AttnBlock2, BasicEncoder
+from timm.models.vision_transformer import Mlp
+
+
+class ST_Transformer(nn.Module):
+    def __init__(self, dim, heads, mlp_dim=512, dropout=0., mlp_ratio=4.0):
+        super().__init__()
+        self.event_t = CrossAttnBlock(128, 128, num_heads=8, mlp_ratio=4.0, dim_head=16)
+        self.fe_space = CrossAttnBlock(128, 128, num_heads=8, mlp_ratio=4.0, dim_head=16)
+        
+    def forward(self, x_i, x_e, x_e_pre):
+        x_q = self.event_t(x_e, x_e_pre)
+        x = self.fe_space(x_q, x_i)
+        
+        return x, x_q
+        
+class downsample(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.layer = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=2, padding=1),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+        )
+    def forward(self, x):
+        return self.layer(x)
+    
+
+class upsample(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.layer = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+        )
+        self.cov = nn.Sequential(
+            nn.Conv2d(out_channels * 2, out_channels, kernel_size=1, stride=1, padding=0),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+        )
+    def forward(self, x1, x2):
+        x1 = F.interpolate(x1, scale_factor=2, mode="bilinear", align_corners=True)
+        x_ = self.layer(x1)
+        if x2 is not None:
+            x_ = torch.cat([x2, x_], dim=1)
+            x_ = self.cov(x_)
+        else:
+            x_ = self.cov(torch.cat([x_, x_], dim=1))
+        return x_
+        
+
+class Fusionformer(nn.Module):
+    def __init__(self, image_size=(384, 512), out_dim=128, mlp_dim=512, depth=6, stride=8, dropout=0.):
+        super().__init__()
+        img_h, img_w = image_size
+        self.stride = stride
+        self.in_planes = 32
+
+        self.fnet_img = BasicEncoder(
+            output_dim=128, norm_fn="instance", dropout=0, stride=stride, shallow=True, in_planes=32
+        )
+        self.fnet_event = BasicEncoder(
+            input_dim=10, output_dim=128, norm_fn="instance", dropout=0, stride=stride, shallow=True, in_planes=32, dilation=1
+        )
+        
+        self.transunet = CLWF(128, out_dim, image_size, stride, mlp_dim, depth, dropout)
+        
+        # self.resnet = ResidualBlock(128, out_dim, stride=1)
+        
+    def forward(self, x_i, x_e, img_ifnew=None, feature_teacher=None):
+        _, _, H, W = x_i.size()
+        
+        x_i = self.fnet_img(x_i)
+        x_e = self.fnet_event(x_e)
+        
+        x_out = self.transunet(x_i, x_e, img_ifnew, feature_teacher)
+        return x_out
+        
+
+class CLWF(nn.Module):
+    def __init__(self, in_channels, out_channels, image_size=(384, 512), stride=8, mlp_dim=512, depth=3, dropout=0.):
+        super().__init__()
+        img_h, img_w = image_size
+        self.patches_resolution = (img_h // (stride * 4), img_w // (stride * 4))
+        num_patches = self.patches_resolution[0] * self.patches_resolution[1]
+        
+        self.down1 = downsample(in_channels, 192)
+        self.down2 = downsample(192, 256)
+
+        self.up1 = upsample(256, 192)
+        self.up2 = upsample(192, out_channels)
+        
+        self.pos_embedding = nn.Parameter(torch.randn(1, num_patches, 256))
+        self.dropout_i = nn.Dropout(dropout)
+        self.dropout_e = nn.Dropout(dropout)
+        
+        self.xe_history = []
+        self.x_e_pre = None
+        self.x_out_ = None
+        # self.x_i_last = None
+        
+        self.cov_out1 = nn.Conv2d(256, 128, kernel_size=1)
+        self.cov_out2 = nn.Conv2d(192, 128, kernel_size=1)
+        
+        self.layers = nn.ModuleList(
+            [
+                CrossAttnBlock(256, 256, num_heads=8, mlp_ratio=4.0, dim_head=32)
+                for _ in range(depth)
+            ]
+        )
+        self.TemporalAdapter = nn.ModuleList(
+            [
+                AttnBlock2(256, num_heads=8, mlp_ratio=4.0, dim_head=32)
+                for _ in range(depth)
+            ]
+        )
+        
+    def forward(self, x_i, x_e, img_ifnew=None, feature_teacher=None):
+        x1_i = self.down1(x_i)
+        x2_i = self.down2(x1_i)
+        
+        x1_e = self.down1(x_e)
+        x2_e = self.down2(x1_e)
+        
+        x2_i_ = x2_i.clone()
+        x2_e_ = x2_e.clone()
+        x2_i_ = rearrange(x2_i_, 'b c h w -> b (h w) c')
+        x2_e_ = rearrange(x2_e_, 'b c h w -> b (h w) c')
+        
+        b, n, _ = x2_i_.shape
+        
+        x2_i_ = self.dropout_i(x2_i_)
+        x2_e_ = self.dropout_e(x2_e_)
+        
+        x_out = []
+        for time in range(len(x_i)):
+            x_i_t = x2_i_[time:time+1]
+            x_e_t = x2_e_[time:time+1]
+
+            if img_ifnew[time] == 1:
+                for st_attn in self.layers:
+                    x_i_t = st_attn(x_i_t, x_e_t)
+                self.x_out_ = x_i_t
+            else:
+                x_i_t = self.x_out_
+                for st_attn in self.layers:
+                    x_i_t = st_attn(x_i_t, x_e_t)
+                self.x_out_ = x_i_t
+
+            x_out.append(rearrange(self.x_out_, 'b (h w) out_dim -> b out_dim h w', h=self.patches_resolution[0], w=self.patches_resolution[1]))
+        
+        x_out = torch.cat(x_out, dim=0)
+        x_out = rearrange(x_out, 't c h w -> (h w) t c')
+        for time_attn in self.TemporalAdapter:
+            x_out = time_attn(x_out)
+        x_out1 = rearrange(x_out, '(h w) t c -> t c h w', h=self.patches_resolution[0], w=self.patches_resolution[1])
+        x_out2 = self.up1(x_out1, torch.tensor(img_ifnew)[:,None,None,None].to(x1_e.device).float() * (x1_i) + torch.tensor(1-img_ifnew)[:,None,None,None].to(x1_e.device).float() * (x1_e))
+        x_out3 = self.up2(x_out2, torch.tensor(img_ifnew)[:,None,None,None].to(x1_e.device).float() * (x_i) + torch.tensor(1-img_ifnew)[:,None,None,None].to(x1_e.device).float() * (x_e))
+        x_out1 = self.cov_out1(x_out1)
+        x_out2 = self.cov_out2(x_out2)
+        x_out_pyramid = [x_out3, x_out2, x_out1]
+        return x_out_pyramid
+    
+
+class Unet_Transformer(nn.Module):
+    def __init__(self, input_dim=3, image_size=(384, 512), out_dim=128, mlp_dim=512, depth=6, stride=8, dropout=0.):
+        super().__init__()
+        img_h, img_w = image_size
+        self.stride = stride
+        self.in_planes = 32
+
+        self.fnet = BasicEncoder(
+            input_dim=input_dim, output_dim=128, norm_fn="instance", dropout=0, stride=stride, shallow=True, in_planes=32
+        )
+        
+        self.transunet = TransUnet_pyramid_onemod(128, out_dim, image_size, stride, mlp_dim, depth, dropout)
+        
+        # self.resnet = ResidualBlock(128, out_dim, stride=1)
+        
+    def forward(self, x, feature_teacher=None):
+        _, _, H, W = x.size()
+        
+        x = self.fnet(x)
+        
+        x_out = self.transunet(x, feature_teacher)
+        
+        # x_out = self.resnet(x_out)
+            
+        return x_out
+    
+
+class TransUnet_pyramid_onemod(nn.Module):
+    def __init__(self, in_channels, out_channels, image_size=(384, 512), stride=8, mlp_dim=512, depth=3, dropout=0.):
+        super().__init__()
+        img_h, img_w = image_size
+        self.patches_resolution = (img_h // (stride * 4), img_w // (stride * 4))
+        num_patches = self.patches_resolution[0] * self.patches_resolution[1]
+        
+        self.down1 = downsample(in_channels, 192)
+        self.down2 = downsample(192, 256)
+
+        self.up1 = upsample(256, 192)
+        self.up2 = upsample(192, out_channels)
+        
+        self.pos_embedding = nn.Parameter(torch.randn(1, num_patches, 256))
+        self.dropout_i = nn.Dropout(dropout)
+        self.dropout_e = nn.Dropout(dropout)
+        
+        self.xe_history = []
+        self.x_e_pre = None
+        self.x_out_ = None
+        
+        self.cov_out1 = nn.Conv2d(256, 128, kernel_size=1)
+        self.cov_out2 = nn.Conv2d(192, 128, kernel_size=1)
+        
+
+        self.TemporalAdapter = nn.ModuleList(
+            [
+                AttnBlock2(256, num_heads=8, mlp_ratio=4.0, dim_head=32)
+                for _ in range(depth)
+            ]
+        )
+        
+    def forward(self, x, feature_teacher=None):
+        x1 = self.down1(x)
+        x2 = self.down2(x1)
+        
+        x2_ = x2.clone()
+        x2_ = rearrange(x2_, 'b c h w -> b (h w) c')
+        
+        b, n, _ = x2_.shape
+        
+        x2_ = self.dropout_i(x2_)
+        
+        x2_ = x2_.permute(1,0,2)
+        
+        for time_attn in self.TemporalAdapter:
+            x_out = time_attn(x2_)
+        x_out1 = rearrange(x_out, '(h w) t c -> t c h w', h=self.patches_resolution[0], w=self.patches_resolution[1])
+        x_out2 = self.up1(x_out1, x1)
+        x_out3 = self.up2(x_out2, x)
+        x_out1 = self.cov_out1(x_out1)
+        x_out2 = self.cov_out2(x_out2)
+        x_out_pyramid = [x_out3, x_out2, x_out1]
+        
+        
+        return x_out_pyramid

LFE_TAP/models/tapformer.py

@@ -0,0 +1,308 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import time
+
+from LFE_TAP.models.blocks import (
+    BasicEncoder,
+    FusionBlock,
+    FusionBlock_basic,
+    EfficientUpdateFormer,
+    UpdateFormer,
+    Mlp,
+)
+from LFE_TAP.models.fusionFormer import Fusionformer
+from LFE_TAP.utils.model_utils import get_track_feat, bilinear_sampler, get_support_points
+from LFE_TAP.models.embeddings import get_1d_sincos_pos_embed_from_grid
+
+torch.manual_seed(0)
+
+
+def posenc(x, min_deg, max_deg):
+    """Cat x with a positional encoding of x with scales 2^[min_deg, max_deg-1].
+    Instead of computing [sin(x), cos(x)], we use the trig identity
+    cos(x) = sin(x + pi/2) and do one vectorized call to sin([x, x+pi/2]).
+    Args:
+      x: torch.Tensor, variables to be encoded. Note that x should be in [-pi, pi].
+      min_deg: int, the minimum (inclusive) degree of the encoding.
+      max_deg: int, the maximum (exclusive) degree of the encoding.
+      legacy_posenc_order: bool, keep the same ordering as the original tf code.
+    Returns:
+      encoded: torch.Tensor, encoded variables.
+    """
+    if min_deg == max_deg:
+        return x
+    scales = torch.tensor(
+        [2**i for i in range(min_deg, max_deg)], dtype=x.dtype, device=x.device
+    )
+
+    xb = (x[..., None, :] * scales[:, None]).reshape(list(x.shape[:-1]) + [-1])
+    four_feat = torch.sin(torch.cat([xb, xb + 0.5 * torch.pi], dim=-1))
+    return torch.cat([x] + [four_feat], dim=-1)
+
+
+class TAPFormer(nn.Module):
+    def __init__(self, window_size=16, stride=8, corr_radius=3, corr_levels=3, backbone="basic", num_heads=8, hidden_size=384, space_depth=3, time_depth=3):
+        super(TAPFormer, self).__init__()
+        self.window_size = window_size
+        self.stride = stride
+        self.corr_radius = corr_radius
+        self.corr_levels = corr_levels
+        self.hidden_size = hidden_size
+        self.space_depth = space_depth
+        self.time_depth = time_depth
+        self.latent_dim = 128
+        self.backbone = backbone
+        self.model_resolution = (384, 512)
+        self.mlp_output_dim = 256
+        self.input_dim = 2 + 84 + self.mlp_output_dim * self.corr_levels
+        num_virtual_tracks = 32
+        
+
+        self.fusion_block = Fusionformer(image_size=self.model_resolution, out_dim=self.latent_dim, mlp_dim=512, stride=self.stride, depth=2)
+        self.updateformer2 = EfficientUpdateFormer(
+            space_depth=space_depth,
+            time_depth=time_depth,
+            input_dim=self.input_dim,
+            hidden_size=hidden_size,
+            output_dim=4,
+            mlp_ratio=4.0,
+            num_virtual_tracks=num_virtual_tracks,
+            linear_layer_for_vis_conf=True,
+        )
+        
+        # self.norm = nn.GroupNorm(1, self.latent_dim)
+        # self.ffeat_updater = nn.Sequential(
+        #     nn.Linear(self.latent_dim, self.latent_dim),
+        #     nn.ReLU(),
+        # )
+        self.corr_mlp = Mlp(in_features=(2*corr_radius+1) ** 4, hidden_features=384, out_features=256)
+        
+        time_grid = torch.linspace(0, window_size - 1, window_size).reshape(1, window_size, 1)
+        self.register_buffer(
+            "time_emb", get_1d_sincos_pos_embed_from_grid(self.input_dim, time_grid[0])
+        )
+        
+    def interpolate_time_embed(self, x, t):
+        previous_dtype = x.dtype
+        T = self.time_emb.shape[1]
+
+        if t == T:
+            return self.time_emb
+
+        time_emb = self.time_emb.float()
+        time_emb = F.interpolate(
+            time_emb.permute(0, 2, 1), size=t, mode="linear"
+        ).permute(0, 2, 1)
+        return time_emb.to(previous_dtype)    
+        
+    def get_correlation_feat(self, fmaps, queried_coords):
+        B, T, D, H_, W_ = fmaps.shape
+        N = queried_coords.shape[1]
+        r = self.corr_radius
+        sample_coords = torch.cat(
+            [torch.zeros_like(queried_coords[..., :1]), queried_coords], dim=-1
+        )[:, None]
+        support_points = get_support_points(sample_coords, r, reshape_back=False)
+        correlation_feat = bilinear_sampler(
+            fmaps.reshape(B * T, D, 1, H_, W_), support_points
+        )
+        return correlation_feat.view(B, T, D, N, (2 * r + 1), (2 * r + 1)).permute(
+            0, 1, 3, 4, 5, 2
+        )
+        
+    def forward_window(self, fmaps_pyramid, coords, track_feat_support_pyramid, vis, conf, attenstion_mask, iters=4):
+        B, S, D, *_ = fmaps_pyramid[0].shape
+        N = coords.shape[2]
+        r = 2 * self.corr_radius + 1
+        
+        coord_preds, vis_preds, conf_preds = [], [], []
+        for it in range(iters):
+            coords = coords.detach()
+            coord_init = coords.view(B * S, N, 2)
+            corr_embs = []
+            for i in range(self.corr_levels):
+                corr_feat = self.get_correlation_feat(fmaps_pyramid[i], coord_init / 2**i)
+                track_feat_support = (
+                    track_feat_support_pyramid[i]
+                    .view(B, 1, r, r, N, self.latent_dim)
+                    .squeeze(1)
+                    .permute(0, 3, 1, 2, 4)
+                )
+                corr_volume = torch.einsum("btnhwc,bnijc->btnhwij", corr_feat, track_feat_support)
+                corr_emb = self.corr_mlp(corr_volume.reshape(B * S * N, r * r * r *r))
+                # del corr_volume, corr_feat
+                # torch.cuda.empty_cache()
+                corr_embs.append(corr_emb)        
+                
+            corr_embs = torch.cat(corr_embs, dim=1)
+            corr_embs = corr_embs.view(B, S, N, corr_embs.shape[-1])
+            
+            transformer_input = [vis, conf, corr_embs]
+            
+            rel_coords_forward = coords[:, :-1] - coords[:, 1:]
+            rel_coords_backward = coords[:, 1:] - coords[:, :-1]
+            
+            rel_coords_forward = torch.nn.functional.pad(rel_coords_forward, (0, 0, 0, 0, 0, 1))
+            rel_coords_backward = torch.nn.functional.pad(rel_coords_backward, (0, 0, 0, 0, 1, 0))
+            
+            scale = (torch.tensor([self.model_resolution[1], self.model_resolution[0]], device=coords.device,) / self.stride)
+            rel_coords_forward = rel_coords_forward / scale     # 归一化到[-1, 1]
+            rel_coords_backward = rel_coords_backward / scale
+            
+            rel_pos_emb_input = posenc(torch.cat([rel_coords_forward, rel_coords_backward], dim=-1), min_deg=0, max_deg=10,)
+            transformer_input.append(rel_pos_emb_input)
+            
+            x = (torch.cat(transformer_input, dim=-1).permute(0, 2, 1, 3).reshape(B*N, S, -1))
+            
+            x = x + self.interpolate_time_embed(x, S)
+            x = x.view(B, N, S, -1)
+            
+            delta = self.updateformer2(x)
+            
+            delta_coords = delta[..., :2].permute(0, 2, 1, 3)
+            delta_vis = delta[..., 2:3].permute(0, 2, 1, 3)
+            delta_conf = delta[..., 3:].permute(0, 2, 1, 3)
+            
+            vis = vis + delta_vis
+            conf = conf + delta_conf
+
+            coords = coords + delta_coords
+            coord_preds.append(coords[..., :2] * float(self.stride))
+
+            vis_preds.append(vis[..., 0])
+            conf_preds.append(conf[..., 0])
+        return coord_preds, vis_preds, conf_preds
+    
+    def forward(self, rgbs, events, queries, iters=4, img_ifnew=None, feat_init=None, is_train=False):
+        B, T, C, H, W = events.shape
+        B, N, _ = queries.shape
+        _, _, C_img, _, _ = rgbs.shape
+        S = self.window_size
+        step = S // 2
+        device = events.device
+        assert H % self.stride == 0 and W % self.stride == 0
+        assert B == 1, "batch size should be 1"
+        
+        queried_frames = queries[:, :, 0].long()
+        queried_coords = queries[..., 1:3]
+        queried_coords = queried_coords / self.stride
+        
+        coords_predicted = torch.zeros((B, T, N, 2), device=device)
+        vis_predicted= torch.zeros((B, T, N), device=device)
+        conf_predicted = torch.zeros((B, T, N), device=device)
+        
+        all_coords_predictions, all_vis_predictions, all_confidence_predictions = ([], [], [])
+        H_stride, W_stride = H // self.stride, W // self.stride
+        
+        rgbs = 2 * (rgbs / 255.0) - 1.0
+        # events = torch.sigmoid(events)
+        events = 2 * events - 1.0
+        dtype = rgbs.dtype
+        
+        # fusion event and image to get fusion feature
+        # start = time.time()
+        fmaps = self.fusion_block(rgbs.reshape(-1, C_img, H, W), events.reshape(-1, C, H, W), img_ifnew)
+        # print(f"fusion block time: {time.time() - start}")
+        fmaps = fmaps.permute(0, 2, 3, 1)
+        fmaps = fmaps / torch.sqrt(
+            torch.maximum(
+                torch.sum(torch.square(fmaps), axis=-1, keepdims=True),
+                torch.tensor(1e-12, device=fmaps.device),
+            )
+        )
+        fmaps = fmaps.permute(0, 3, 1, 2).reshape(
+            B, -1, self.latent_dim, H_stride, W_stride
+        )
+        fmaps = fmaps.to(dtype)
+        
+        # compute queries point feature
+        fmaps_pyramid = []
+        track_feat_pyramid = []
+        track_feat_support_pyramid = []
+        fmaps_pyramid.append(fmaps)
+        for i in range(self.corr_levels - 1):
+            fmaps_ = fmaps.reshape(B * T, self.latent_dim, fmaps.shape[-2], fmaps.shape[-1])
+            fmaps_ = F.avg_pool2d(fmaps_, 2, stride=2)
+            fmaps = fmaps_.reshape(B, T, self.latent_dim, fmaps_.shape[-2], fmaps_.shape[-1])
+            fmaps_pyramid.append(fmaps)
+            
+        for i in range(self.corr_levels):
+            track_feat, track_feat_support = get_track_feat(fmaps_pyramid[i], queried_frames, queried_coords/2**i, support_radius=self.corr_radius)
+            track_feat_pyramid.append(track_feat.repeat(1, T, 1, 1))
+            track_feat_support_pyramid.append(track_feat_support.unsqueeze(1))
+            
+        vis_init = torch.zeros((B, S, N, 1), device=device).float()
+        conf_init = torch.zeros((B, S, N, 1), device=device).float()
+        coords_init = queried_coords.reshape(B, 1, N, 2).expand(B, S, N, 2).float()
+        
+        num_windows =  (T - S + step - 1) // step + 1
+        indices = range(0, step * num_windows, step)
+        
+        for ind in indices:
+            if ind > 0:
+                overlap = S - step
+                copy_over = (queried_frames < ind + overlap)[:, None, :, None]   # B, 1, N, 1
+                coords_prev = coords_predicted[:, ind : ind + overlap] / self.stride
+                padding_tensor = coords_prev[:, -1:, :, :].expand(-1, step, -1, -1)  # 将上一个时刻的坐标作为待优化坐标的初始值
+                coords_prev = torch.cat([coords_prev, padding_tensor], dim=1)
+                
+                vis_prev = vis_predicted[:, ind : ind + overlap, :, None].clone()
+                padding_tensor = vis_prev[:, -1:, :, :].expand(-1, step, -1, -1)
+                vis_prev = torch.cat([vis_prev, padding_tensor], dim=1)
+                
+                conf_prev = conf_predicted[:, ind : ind + overlap, :, None].clone()
+                padding_tensor = conf_prev[:, -1:, :, :].expand(-1, step, -1, -1)
+                conf_prev = torch.cat([conf_prev, padding_tensor], dim=1)
+                
+                coords_init = torch.where(copy_over.expand_as(coords_init), coords_prev, coords_init)
+                vis_init = torch.where(copy_over.expand_as(vis_init), vis_prev, vis_init)
+                conf_init = torch.where(copy_over.expand_as(conf_init), conf_prev, conf_init)
+                
+            attenstion_mask = (queried_frames < ind + S).reshape(B, 1, N) # B, 1, N
+            # start = time.time()
+            coords, viss, confs = self.forward_window(
+                fmaps_pyramid=[fmap[:, ind : ind +S] for fmap in fmaps_pyramid],
+                coords=coords_init,
+                track_feat_support_pyramid=[attenstion_mask[:, None, :, :, None]*tfeat for tfeat in track_feat_support_pyramid],
+                vis=vis_init,
+                conf=conf_init,
+                attenstion_mask=attenstion_mask.repeat(1, S, 1),
+                iters=iters,
+            )
+            # print(f"forward window time: {time.time() - start}")
+            S_trimmed = min(T - ind, S)     # accounts for last window duration
+            coords_predicted[:, ind : ind + S] = coords[-1][:, :S_trimmed]
+            vis_predicted[:, ind : ind + S] = viss[-1][:, :S_trimmed]
+            conf_predicted[:, ind : ind + S] = confs[-1][:, :S_trimmed]
+            if is_train:
+                all_coords_predictions.append(
+                    [coord[:, :S_trimmed] for coord in coords]
+                )
+                all_vis_predictions.append(
+                    [torch.sigmoid(vis[:, :S_trimmed]) for vis in viss]
+                )
+                all_confidence_predictions.append(
+                    [torch.sigmoid(conf[:, :S_trimmed]) for conf in confs]
+                )
+                
+        vis_predicted = torch.sigmoid(vis_predicted)
+        conf_predicted = torch.sigmoid(conf_predicted)
+        
+        if is_train:
+            valid_mask = (
+                queried_frames[:, None]
+                <= torch.arange(0, T, device=device)[None, :, None]
+            )
+            train_data = (
+                all_coords_predictions,
+                all_vis_predictions,
+                all_confidence_predictions,
+                valid_mask,
+            )
+        else:
+            train_data = None
+            
+        return coords_predicted, vis_predicted, conf_predicted, train_data
+        
+        

LFE_TAP/utils/dataset_utils.py

@@ -0,0 +1,115 @@
+import torch
+import dataclasses
+import numpy as np
+import torch.nn.functional as F
+from typing import Optional, Any, Union
+
+@dataclasses.dataclass(eq=False)
+class FrameEventData:
+    """ Data class for frame, event, tracks data. """
+    video: torch.Tensor         # (B, S, C_i, H, W)
+    events: torch.Tensor        # (B, S, C_e, H, W)
+    segmentation: torch.Tensor  # (B, S, 1, H, W)
+    trajectory: torch.Tensor    # (B, S, N, 2)
+    visibility: torch.Tensor    # (B, S, N)
+    img_ifnew: torch.Tensor = None 
+    # optional daa
+    clear_video: Optional[torch.Tensor] = None         # (B, S, C_i, H, W)
+    valid: Optional[torch.Tensor] = None  # (B, S, N)
+    seq_name: Optional[torch.Tensor] = None 
+    
+    
+@dataclasses.dataclass(eq=False)
+class FrameEventData_test:
+    video: np.array
+    events: Union[np.array, list]
+    segmentation: np.array
+    trajectory: np.array
+    query_points: torch.Tensor
+    # optional daa
+    visibility: Optional[torch.Tensor] = None
+    valid: Optional[torch.Tensor] = None  # (B, S, N)
+    seq_name: Optional[torch.Tensor] = None  
+    img_ifnew: Optional[np.array] = None
+    img_ifnew_full: Optional[np.array] = None
+    rgb_timestamp: Optional[np.array] = None
+    
+    
+def collate_fn(batch):
+    """ Collate function for frame, event, tracks data. """
+    video = torch.stack([b.video for b, in batch], dim=0)
+    events = torch.stack([b.events for b in batch], dim=0)
+    segmentation = torch.stack([b.segmentation for b in batch], dim=0)
+    trajectory = torch.stack([b.trajectory for b in batch], dim=0)
+    visibility = torch.stack([b.visibility for b in batch], dim=0)
+    
+    seq_name = [b.seq_name for b in batch]
+    
+    return FrameEventData(video, events, segmentation, trajectory, visibility, seq_name=seq_name)
+    
+    
+def collate_fn_train(batch):
+    """ Collate function for training data. """
+    gotit = [gotit for _, gotit in batch]
+    video = torch.stack([b.video for b, _ in batch], dim=0)
+    events = torch.stack([b.events for b, _ in batch], dim=0)
+    segmentation = torch.stack([b.segmentation for b, _ in batch], dim=0)
+    trajectory = torch.stack([b.trajectory for b, _ in batch], dim=0)
+    visibility = torch.stack([b.visibility for b, _ in batch], dim=0)
+    clear_video = torch.stack([b.clear_video for b, _ in batch], dim=0)
+    valid = torch.stack([b.valid for b, _ in batch], dim=0)
+    seq_name = [b.seq_name for b, _ in batch]
+    for b, _ in batch:
+        if b.img_ifnew is None:
+            gotit = [False]
+            img_ifnew = []
+        else:
+            img_ifnew = torch.stack([b.img_ifnew for b, _ in batch], dim=0)
+    return (FrameEventData(video, events, segmentation, trajectory, visibility, clear_video=clear_video, valid=valid, seq_name=seq_name, img_ifnew=img_ifnew), gotit)
+
+
+def collate_fn_EDS(batch):
+    video = np.stack([b.video for b, _ in batch], axis=0)
+    events = np.stack([b.events for b, _ in batch], axis=0)
+    segmentation = torch.stack([b.segmentation for b, _ in batch], axis=0)
+    trajectory = np.stack([b.trajectory for b, _ in batch], axis=0)
+    query_points = None
+    if batch[0][0].query_points is not None:
+        query_points = torch.stack([b.query_points for b, _ in batch], dim=0)
+    
+    seq_name = [b.seq_name for b, _ in batch]
+    rgb_timestamp = [b.rgb_timestamp for b, _ in batch]
+    
+    return FrameEventData_test(video, events, segmentation, trajectory, query_points, seq_name=seq_name, rgb_timestamp=rgb_timestamp)
+
+
+def try_to_cuda(t: Any) -> Any:
+    """
+    Try to move the input variable `t` to a cuda device.
+
+    Args:
+        t: Input.
+
+    Returns:
+        t_cuda: `t` moved to a cuda device, if supported.
+    """
+    try:
+        t = t.float().cuda()
+    except AttributeError:
+        pass
+    return t
+
+
+def dataclass_to_cuda_(obj):
+    """
+    Move all contents of a dataclass to cuda inplace if supported.
+
+    Args:
+        batch: Input dataclass.
+
+    Returns:
+        batch_cuda: `batch` moved to a cuda device, if supported.
+    """
+    for f in dataclasses.fields(obj):
+        setattr(obj, f.name, try_to_cuda(getattr(obj, f.name)))
+    return obj

LFE_TAP/utils/event/representations.py

@@ -0,0 +1,312 @@
+import torch
+import numpy as np
+import cv2
+from enum import Enum, auto
+
+import hdf5plugin
+import h5py
+
+
+class EventRepresentationTypes(Enum):
+    time_surface = 0
+    voxel_grid = 1
+    event_stack = 2
+
+
+class EventRepresentation:
+    def __init__(self):
+        pass
+
+    def convert(self, events):
+        raise NotImplementedError
+
+
+class TimeSurface(EventRepresentation):
+    def __init__(self, input_size: tuple, p, t, x, y):
+        assert len(input_size) == 3
+        t = t.astype('float32')
+        x = x.astype('float32')
+        y = y.astype('float32')
+        pol = p.astype('float32')
+        self.input_size = input_size
+        self.time_surface = torch.zeros(input_size, dtype=torch.float, requires_grad=False)
+        self.n_bins = input_size[0] // 2
+
+    def convert(self, events):
+        _, H, W = self.time_surface.shape
+        with torch.no_grad():
+            self.time_surface = torch.zeros(self.input_size, dtype=torch.float, requires_grad=False,
+                                            device=events['p'].device)
+            time_surface = self.time_surface.clone()
+
+            t = events['t'].cpu().numpy()
+            dt_bin = 1. / self.n_bins
+            x0 = events['x'].int()
+            y0 = events['y'].int()
+            p0 = events['p'].int()
+            t0 = events['t']
+
+            # iterate over bins
+            for i_bin in range(self.n_bins):
+                t0_bin = i_bin * dt_bin
+                t1_bin = t0_bin + dt_bin
+
+                # mask_t = np.logical_and(time > t0_bin, time <= t1_bin)
+                # x_bin, y_bin, p_bin, t_bin = x[mask_t], y[mask_t], p[mask_t], time[mask_t]
+                idx0 = np.searchsorted(t, t0_bin, side='left')
+                idx1 = np.searchsorted(t, t1_bin, side='right')
+                x_bin = x0[idx0:idx1]
+                y_bin = y0[idx0:idx1]
+                p_bin = p0[idx0:idx1]
+                t_bin = t0[idx0:idx1]
+
+                n_events = len(x_bin)
+                for i in range(n_events):
+                    if 0 <= x_bin[i] < W and 0 <= y_bin[i] < H:
+                        time_surface[2*i_bin+p_bin[i], y_bin[i], x_bin[i]] = t_bin[i]
+
+        return time_surface
+    
+    
+class TimeOrderSurface():
+    def __init__(self, input_size: tuple, x, y, p, t):
+        assert len(input_size) == 3
+        H, W, C = input_size
+        self.t = torch.from_numpy(t.astype('int32'))
+        self.x = torch.from_numpy(x.astype('int32'))
+        self.y = torch.from_numpy(y.astype('int32'))
+        self.pol = torch.from_numpy(p.astype('int32'))
+        
+        mask = (self.x >= 3) & (self.x < W - 3) & (self.y >= 3) & (self.y < H - 3) & (self.t >= 370000) & (self.t <= 900000)
+        self.x = self.x[mask]
+        self.y = self.y[mask]
+        self.pol = self.pol[mask]
+        self.t = self.t[mask]
+        
+        self.index = torch.tensor(0 , device=self.t.device)
+        self.input_size = input_size
+        self.n_bins = input_size[2] // 2
+        self.tos = torch.zeros((input_size[0], input_size[1], 2), dtype=torch.float, requires_grad=False, device=self.t.device)
+        self.sae = torch.zeros((input_size[0], input_size[1], 2), dtype=torch.float, requires_grad=False, device=self.t.device)
+        self.sae_latest = torch.zeros((input_size[0], input_size[1], 2), dtype=torch.float, requires_grad=False, device=self.t.device)
+        self.TOS_bins_001 = torch.zeros(self.input_size, dtype=torch.float, requires_grad=False,
+                                        device=self.t.device)
+        self.TOS_bins_002 = torch.zeros(self.input_size, dtype=torch.float, requires_grad=False,
+                                        device=self.t.device)
+        
+    def convert(self, time, n_bins, type):
+        assert n_bins < self.n_bins
+        with torch.no_grad():
+            # TOS_bins_001 = self.TOS_bins_001.clone()
+            # TOS_bins_002 = self.TOS_bins_002.clone()
+            
+            while self.index < len(self.t) and self.t[self.index] <= time:
+                pol = 1 if self.pol[self.index] else 0
+                pol_inv = 0 if self.pol[self.index] else 1
+                if ((self.t[self.index] > self.sae_latest[self.y[self.index]][self.x[self.index]][pol] + 20000) or
+                        (self.sae_latest[self.y[self.index]][self.x[self.index]][pol_inv] > self.sae_latest[self.y[self.index]][self.x[self.index]][pol])):
+                    self.sae_latest[self.y[self.index]][self.x[self.index]][pol] = self.t[self.index]
+                    self.sae[self.y[self.index]][self.x[self.index]][pol] = self.t[self.index]
+                    self.tos[self.y[self.index] - 3:self.y[self.index] + 4,
+                    self.x[self.index] - 3:self.x[self.index] + 4, self.pol[self.index]] -= 1
+                    self.tos[self.y[self.index] - 3:self.y[self.index] + 4, self.x[self.index] - 3:self.x[self.index] + 4, self.pol[self.index]][
+                        self.tos[self.y[self.index] - 3:self.y[self.index] + 4, self.x[self.index] - 3:self.x[self.index] + 4, self.pol[self.index]] < 241] = 0
+                    self.tos[self.y[self.index], self.x[self.index], self.pol[self.index]] = 255
+                else:
+                    self.sae_latest[self.y[self.index]][self.x[self.index]][pol] = self.t[self.index]
+
+                self.index += 1
+            
+            # while self.t[self.index] <= 900000 and self.t[self.index] <= time:
+            #     for x0 in range(self.x[self.index]-3, self.x[self.index]+4):
+            #         for y0 in range(self.y[self.index]-3, self.y[self.index]+4):
+            #             if self.tos[y0, x0, self.pol[self.index]] != 0:
+            #                 self.tos[y0, x0, self.pol[self.index]] -= 1
+            #                 if self.tos[y0, x0, self.pol[self.index]] < 241:
+            #                     self.tos[y0, x0, self.pol[self.index]] = 0
+            #                 # self.tos[y0, x0, self.pol[self.index]].clamp_(min=0, max=240)  # 使用clamp函数限制值的范围
+            #     self.tos[self.y[self.index], self.x[self.index], self.pol[self.index]] = 255
+            #     self.index += 1
+            
+            # time_threshold = 900000
+            #
+            # # 使用torch.where函数实现条件操作
+            # while torch.any(torch.logical_and(self.t[self.index] <= time_threshold, self.t[self.index] <= time)):
+            #     y_slice = slice(self.y[self.index]-3, self.y[self.index]+4)
+            #     x_slice = slice(self.x[self.index]-3, self.x[self.index]+4)
+            #     mask = (self.tos[y_slice, x_slice, self.pol[self.index]] != 0)
+            #     self.tos[y_slice, x_slice, self.pol[self.index]].masked_scatter_(mask, self.tos[y_slice, x_slice, self.pol[self.index]] - 1)
+            #     self.tos[y_slice, x_slice, self.pol[self.index]] = torch.where(self.tos[y_slice, x_slice, self.pol[self.index]] < 241,
+            #                                                                    torch.tensor(0), self.tos[y_slice, x_slice, self.pol[self.index]])
+            #     self.tos[self.y[self.index], self.x[self.index], self.pol[self.index]] = 255
+            #     self.index += 1
+
+            if type == 0:
+                self.TOS_bins_001[:, :, 2*n_bins] = self.tos[:, :, 0]
+                self.TOS_bins_001[:, :, 2*n_bins+1] = self.tos[:, :, 1]
+            elif type == 1:
+                self.TOS_bins_002[:, :, 2*n_bins] = self.tos[:, :, 0]
+                self.TOS_bins_002[:, :, 2*n_bins+1] = self.tos[:, :, 1]
+                
+    def get_time_order_surface(self, type):
+        # tos = self.tos.numpy()
+        # cv2.imshow("p_tos", tos[:, :, 0])
+        # cv2.imshow("n_tos", tos[:, :, 1])
+        # cv2.waitKey(0)
+        if type == 0:
+            return self.TOS_bins_001
+        elif type == 1:
+            return self.TOS_bins_002
+                         
+
+class VoxelGrid(EventRepresentation):
+    def __init__(self, input_size: tuple, normalize: bool):
+        assert len(input_size) == 3
+        self.voxel_grid = torch.zeros((input_size), dtype=torch.float, requires_grad=False)
+        self.input_size = input_size
+        self.nb_channels = input_size[0]
+        self.normalize = normalize
+
+    def convert(self, events):
+        C, H, W = self.voxel_grid.shape
+        with torch.no_grad():
+            self.voxel_grid = torch.zeros((self.input_size), dtype=torch.float, requires_grad=False,
+                                          device=events['p'].device)
+            voxel_grid = self.voxel_grid.clone()
+
+            t_norm = events['t']
+            t_norm = (C - 1) * (t_norm-t_norm[0]) / (t_norm[-1]-t_norm[0])
+
+            x0 = events['x'].int()
+            y0 = events['y'].int()
+            t0 = t_norm.int()
+
+            value = 2*events['p']-1
+
+            for xlim in [x0,x0+1]:
+                for ylim in [y0,y0+1]:
+                    for tlim in [t0,t0+1]:
+
+                        mask = (xlim < W) & (xlim >= 0) & (ylim < H) & (ylim >= 0) & (tlim >= 0) & (tlim < self.nb_channels)
+                        interp_weights = value * (1 - (xlim-events['x']).abs()) * (1 - (ylim-events['y']).abs()) * (1 - (tlim - t_norm).abs())
+
+                        index = H * W * tlim.long() + \
+                                W * ylim.long() + \
+                                xlim.long()
+
+                        voxel_grid.put_(index[mask], interp_weights[mask], accumulate=True)
+
+            if self.normalize:
+                mask = torch.nonzero(voxel_grid, as_tuple=True)
+                if mask[0].size()[0] > 0:
+                    mean = voxel_grid[mask].mean()
+                    std = voxel_grid[mask].std()
+                    if std > 0:
+                        voxel_grid[mask] = (voxel_grid[mask] - mean) / std
+                    else:
+                        voxel_grid[mask] = voxel_grid[mask] - mean
+
+        return voxel_grid
+
+
+class EventStack(EventRepresentation):
+    def __init__(self, input_size: tuple):
+        """
+        :param input_size: (C, H, W)
+        """
+        assert len(input_size) == 3
+        self.input_size = input_size
+        self.event_stack = torch.zeros((input_size), dtype=torch.float, requires_grad=False)
+        self.nb_channels = input_size[0]
+
+    def convert(self, events):
+        C, H, W = self.event_stack.shape
+        with torch.no_grad():
+            self.event_stack = torch.zeros((self.input_size), dtype=torch.float, requires_grad=False,
+                                           device=events['p'].device)
+            event_stack = self.event_stack.clone()
+
+            t = events['t'].cpu().numpy()
+            dt_bin = 1. / self.nb_channels
+            x0 = events['x'].int()
+            y0 = events['y'].int()
+            p0 = 2*events['p'].int()-1
+            t0 = events['t']
+
+            # iterate over bins
+            for i_bin in range(self.nb_channels):
+                t0_bin = i_bin * dt_bin
+                t1_bin = t0_bin + dt_bin
+
+                # mask_t = np.logical_and(time > t0_bin, time <= t1_bin)
+                # x_bin, y_bin, p_bin, t_bin = x[mask_t], y[mask_t], p[mask_t], time[mask_t]
+                idx0 = np.searchsorted(t, t0_bin, side='left')
+                idx1 = np.searchsorted(t, t1_bin, side='right')
+                x_bin = x0[idx0:idx1]
+                y_bin = y0[idx0:idx1]
+                p_bin = p0[idx0:idx1]
+
+                n_events = len(x_bin)
+                for i in range(n_events):
+                    if 0 <= x_bin[i] < W and 0 <= y_bin[i] < H:
+                        event_stack[i_bin, y_bin[i], x_bin[i]] += p_bin[i]
+
+        return event_stack
+
+
+def events_to_time_surface(time_surface, p, t, x, y):
+    t = (t - t[0]).astype('float32')
+    t = (t/t[-1])
+    x = x.astype('float32')
+    y = y.astype('float32')
+    pol = p.astype('float32')
+    event_data_torch = {
+        'p': torch.from_numpy(pol),
+        't': torch.from_numpy(t),
+        'x': torch.from_numpy(x),
+        'y': torch.from_numpy(y),
+    }
+    return time_surface.convert(event_data_torch)
+
+def events_to_time_order_surface(time_order_surface, p, t, x, y, dt):
+    t = t.astype('float32')
+    x = x.astype('float32')
+    y = y.astype('float32')
+    pol = p.astype('float32')
+    event_data_torch = {
+        'p': torch.from_numpy(pol),
+        't': torch.from_numpy(t),
+        'x': torch.from_numpy(x),
+        'y': torch.from_numpy(y),
+    }
+    
+    time_order_surface.convert(event_data_torch, dt)
+
+def events_to_event_stack(event_stack, p, t, x, y, dt):
+    t = (t - t[0]).astype('float32')
+    t = (t/t[-1])
+    x = x.astype('float32')
+    y = y.astype('float32')
+    pol = p.astype('float32')
+    event_data_torch = {
+        'p': torch.from_numpy(pol),
+        't': torch.from_numpy(t),
+        'x': torch.from_numpy(x),
+        'y': torch.from_numpy(y),
+    }
+    return event_stack.convert(event_data_torch)
+
+
+def events_to_voxel_grid(voxel_grid, p, t, x, y):
+    t = (t - t[0]).astype('float32')
+    t = (t/t[-1])
+    x = x.astype('float32')
+    y = y.astype('float32')
+    pol = p.astype('float32')
+    event_data_torch = {
+        'p': torch.from_numpy(pol),
+        't': torch.from_numpy(t),
+        'x': torch.from_numpy(x),
+        'y': torch.from_numpy(y),
+    }
+    return voxel_grid.convert(event_data_torch)