import os.path
import numpy as np
import torch
import csv
import pandas
import random
from collections import OrderedDict
from .base_video_dataset import BaseVideoDataset
from lib.train.data import jpeg4py_loader
from lib.train.admin import env_settings
import scipy.io as scio

class Coesot(BaseVideoDataset):
def init(self, root=None, image_loader=jpeg4py_loader, split=None, seq_ids=None, data_fraction=None):

root = env_settings().got10k_dir if root is None else root
super().init(‘Coesot’, root, image_loader)

self.sequence_list = self._get_sequence_list()

# seq_id is the index of the folder inside the got10k root path
if split is not None:
if seq_ids is not None:
raise ValueError(‘Cannot set both split_name and seq_ids.’)
if split == ‘train’:
file_path = os.path.join(self.root, ‘train.txt’)
elif split == ‘val’:
file_path = os.path.join(self.root, ‘val.txt’)
else:
raise ValueError(‘Unknown split name’)
seq_ids = pandas.read_csv(file_path, header=None, dtype=np.int64).squeeze(“columns”).values.tolist()
elif seq_ids is None:
seq_ids = list(range(0, len(self.sequence_list)))

self.sequence_list = [self.sequence_list[i] for i in seq_ids]

def get_name(self):
return ‘coesot’

def _get_sequence_list(self):
with open(os.path.join(self.root, ‘list.txt’)) as f:
dir_list = list(csv.reader(f))
dir_list = [dir_name[0] for dir_name in dir_list]
return dir_list

def _read_bb_anno(self, seq_path):
bb_anno_file = os.path.join(seq_path, “groundtruth.txt”)
gt = pandas.read_csv(bb_anno_file, delimiter=‘,’, header=None, dtype=np.float32, na_filter=False, low_memory=False).values
return torch.tensor(gt)

def _get_sequence_path(self, seq_id):
return os.path.join(self.root, self.sequence_list[seq_id], self.sequence_list[seq_id] + “_aps”)


def _get_event_img_sequence_path(self, seq_id):
return os.path.join(self.root, self.sequence_list[seq_id], self.sequence_list[seq_id] + “_graph”)

def _get_grountgruth_path(self, seq_id):
return os.path.join(self.root, self.sequence_list[seq_id])

def get_sequence_info(self, seq_id):
bbox_path = self._get_grountgruth_path(seq_id)
bbox = self._read_bb_anno(bbox_path)

valid = (bbox[:, 2] > 0) & (bbox[:, 3] > 0)
visible = valid.clone().byte()
# return {‘bbox’: bbox, ‘valid’: valid, ‘visible’: visible, ‘visible_ratio’: visible_ratio}
return {‘bbox’: bbox, ‘valid’: valid, ‘visible’: visible, }

def _get_frame_path(self, seq_path, frame_id):
if os.path.exists(os.path.join(seq_path, ‘frame{:04}.png’.format(frame_id))):
return os.path.join(seq_path, ‘frame{:04}.png’.format(frame_id)) # frames start from 0
else:
return os.path.join(seq_path, ‘frame{:04}.bmp’.format(frame_id)) # some image is bmp

def _get_frame(self, seq_path, frame_id):
return self.image_loader(self._get_frame_path(seq_path, frame_id))

def _get_event_sequence_path(self, seq_id): ## get evemts’ frames
return os.path.join(self.root, self.sequence_list[seq_id], self.sequence_list[seq_id] + “_voxel”)

def _get_event_frame(self, seq_path, frame_id):
frame_event_list = []
for f_id in frame_id:
event_frame_file = os.path.join(seq_path, ‘frame{:04}.mat’.format(f_id))
if os.path.getsize(event_frame_file) == 0:
event_features = np.zeros(4096, 19)
# need_data = [np.zeros([4096, 3]), np.zeros([4096, 16])]
else:
mat_data = scio.loadmat(event_frame_file)
# need_data = [mat_data[‘coor’], mat_data[‘features’]]
event_features = np.concatenate((mat_data[‘coor’], mat_data[‘features’]), axis=1) # concat coorelate and features (x,y,z, feauture32/16)
if np.isnan(event_features).any():
event_features = np.zeros(4096, 19)
print(event_frame_file, ‘exist nan value in voxel.’)
frame_event_list.append(event_features)
return frame_event_list

def _get_event_frame_graph(self, seq_path, frame_id):
frame_event_list1 = []
for f_id in frame_id:
event_frame_file = os.path.join(seq_path, ‘frame{:04}.pt’.format(f_id))
if os.path.getsize(event_frame_file) == 0:
event_features = np.zeros(5000, 19)
else:
data = torch.load(event_frame_file)
event_features = torch.cat((data.pos, data.x), axis=1) # concat coorelate and features (x,y,z, feauture32/16)
if np.isnan(event_features).any():
event_features = np.zeros(5000, 19)
print(event_frame_file, ‘exist nan value in voxel.’)
frame_event_list1.append(event_features)
return frame_event_list1

def get_frames(self, seq_id, frame_ids, anno=None):
seq_path = self._get_sequence_path(seq_id)
# obj_meta = self.sequence_meta_info[self.sequence_list[seq_id]]

frame_list = [self._get_frame(seq_path, f_id) for f_id in frame_ids]
seq_event_path = self._get_event_img_sequence_path(seq_id)
# frame_event_img_list = [self._get_frame(seq_event_path, f_id) for f_id in frame_ids]
frame_event_img_list = self._get_event_frame_graph(seq_event_path, frame_ids)
if anno is None:
anno = self.get_sequence_info(seq_id)

anno_frames = {}
for key, value in anno.items():
anno_frames[key] = [value[f_id, …].clone() for f_id in frame_ids]

object_meta = OrderedDict({‘object_class_name’: None,
‘motion_class’: None,
‘major_class’: None,
‘root_class’: None,
‘motion_adverb’: None})

seq_event_path = self._get_event_sequence_path(seq_id)
frame_event_list = self._get_event_frame(seq_event_path, frame_ids)
# seq_event_path = self._get_event_img_sequence_path(seq_id)
# frame_event_list = self._get_event_frame_graph(seq_event_path, frame_ids)
frame_event_list = frame_event_img_list
return frame_list, anno_frames, object_meta, frame_event_list, frame_event_img_list

中，我们期望event模态读取的是graph目录下的.pt文件，其中.pt文件是由voxel生成的，import os
import pdb
import csv
import numpy as np
import random
import cv2
import torch
import pandas as pd
from tqdm import tqdm
import scipy.io as sio
from torch_geometric.data import Data
from spconv.pytorch.utils import PointToVoxel
import math

# CPU usage limit
import os
os.environ[“KMP_BLOCKTIME”] = “0”
os.environ[“OMP_NUM_THREADS”] = “1”
import torch
torch.set_num_threads(1)
import cv2
cv2.setNumThreads(0)
cv2.ocl.setUseOpenCL(False)

def calculate_edges(data, r=5):
# threshold of radius
d = 32
# scaling factor to tune the difference between temporal and spatial resolution
alpha = 1
beta = 1
data_size = data.shape[0]
# max number of edges is 1000000,
edges = np.zeros([100000, 2])
# get t, x,y
points = data[:, 0:3]
row_num = 0
edge_sum=0
for i in range(data_size - 1):
count = 0
distance_matrix = points[i + 1 : data_size + 1, 0:3]
distance_matrix[:, 1:3] = distance_matrix[:, 1:3] - points[i, 1:3]
distance_matrix[:, 0] = distance_matrix[:, 0] - points[i, 0]
distance_matrix = np.square(distance_matrix)
distance_matrix[:, 0] *= alpha
distance_matrix[:, 1:3] *= beta
# calculate the distance of each pair of events
distance = np.sqrt(np.sum(distance_matrix, axis=1))
index = np.where(distance <= r)
# save the edges
if index:
index = index[0].tolist()
for id in index:
edges[row_num, 0] = i
edges[row_num + 1, 1] = i
edges[row_num, 1] = int(id) + i + 1
edges[row_num + 1, 0] = int(id) + i + 1
row_num = row_num + 2
count = count + 1
edge_sum+=2
if count > d:
break
if edge_sum>40000:
break
edges = edges[~np.all(edges == 0, axis=1)]
edges = np.transpose(edges)

return edges

def generate_graph(data):
feature = data[‘features’]

position = data[‘coor’]

edges = calculate_edges(data[‘coor’], 3)
return feature,position,edges


if name == ‘main’:
data_path = ‘/data2/local_userdata/heyuting/datasets/COESOT/train’
save_path = ‘/data2/local_userdata/heyuting/datasets/COESOT/train’
device = torch.device(“cuda:4”)

video_files = os.listdir(data_path)
video_files = [file for file in video_files if file not in (‘list.txt’, ‘train.txt’)]
video_files.sort()
for videoID in range(len(video_files)):
foldName = video_files[videoID]
if foldName not in [‘dvSave-2022_03_21_11_22_27’]:
print("==>> pass it: ", foldName)
continue

label = videoID
# fileLIST = os.listdir(os.path.join(data_path, foldName))
voxel_dir = os.path.join(data_path, foldName, f’{foldName}_voxel2’)
if not os.path.exists(os.path.join(save_path, foldName, foldName+‘_graph’)):
os.mkdir(os.path.join(save_path, foldName, foldName +‘_graph’))
else:
continue
fileLIST = sorted(os.listdir(voxel_dir)) # 确保顺序读取
for file_Name in tqdm(fileLIST):
frame_no = int(file_Name.split(‘.’)[0][-4:])
save_name = f’frame{frame_no:04d}‘
read_path = os.path.join(voxel_dir, file_Name)
data = sio.loadmat(read_path)
feature,position,edges = generate_graph(data)

feature = torch.tensor(feature)[:, :].float()
edge_index = torch.tensor(
np.array(edges).astype(np.int32), dtype=torch.long
)
pos = torch.tensor(np.array(position), dtype=torch.float32)
label_idx = torch.tensor(int(label), dtype=torch.long)
data = Data(
x=feature, edge_index=edge_index, pos=pos, y=label_idx.unsqueeze(0)
)

G_save_path =os.path.join(os.path.join(save_path, foldName, foldName+’_graph’, ‘{}.pt’.format(save_name)))

torch.save(data, G_save_path)，现在我们在终端查看数据：>>> import torch
>>> path=‘/data2/local_userdata/heyuting/datasets/COESOT/train/dvSave-2022_03_21_16_23_44/dvSave-2022_03_21_16_23_44_graph/frame0001.pt’ >>> a=torch.load(path)
>>> print(a)
Data(x=[5000, 16], edge_index=[2, 40056], y=[1], pos=[5000, 3]) ，那么检查代码中，关于.pt文件的读取有无问题。（中文回答），首先直接回答是否有错误。