import os
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文件的读取有无问题。