在import os
import pdb
# import csv
import numpy as np
import cv2
import torch
import pandas as pd
from tqdm import tqdm
import scipy.io
from spconv.pytorch.utils import PointToVoxel
from dv import AedatFile
import numpy as np
import multiprocessing
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

def transform_points_to_voxels(data_dict={}, voxel_generator=None, device=torch.device("cuda:0")):
    """
    将点云转换为voxel,调用spconv的VoxelGeneratorV2
    """
    points = data_dict['points']
    # 将points打乱
    shuffle_idx = np.random.permutation(points.shape[0])
    points = points[shuffle_idx]
    data_dict['points'] = points

    # 使用spconv生成voxel输出
    points = torch.as_tensor(data_dict['points']).to(device)
    voxel_output = voxel_generator(points)

    # 假设一份点云数据是N*4，那么经过pillar生成后会得到三份数据
    # voxels代表了每个生成的voxel数据，维度是[M, 5, 4]
    # coordinates代表了每个生成的voxel所在的zyx轴坐标，维度是[M,3]
    # num_points代表了每个生成的voxel中有多少个有效的点维度是[m,]，因为不满5会被0填充
    voxels, coordinates, num_points = voxel_output
    voxels = voxels.to(device)
    coordinates = coordinates.to(device)
    num_points = num_points.to(device)
    # 选event数量在前5000的voxel  8000 from(4k+,6k+)
    # print(torch.where(num_points>=16)[0].shape)
    if num_points.shape[0] < save_voxel:
        features = voxels[:, :, 3]
        coor = coordinates[:, :]
    else:
        _, voxels_idx = torch.topk(num_points, save_voxel)
        # 将每个voxel的1024个p拼接作为voxel初始特征   16
        features = voxels[voxels_idx][:, :, 3]
        # 前5000个voxel的三维坐标
        coor = coordinates[voxels_idx]
    # 将y.x.t改为t,x,y
    coor[:, [0, 1, 2]] = coor[:, [2, 1, 0]]

    return coor, features


if __name__ == '__main__':
    save_voxel = 5000
    use_mode = 'frame_exposure_time'
    device = torch.device("cuda:0")
    # data_path = r"/home/ioe/xxxxx/EventTracking/COESOT/train"
    data_path = r"/data2/local_userdata/heyuting/datasets/fe108_coesot/test"
    save_path = r"/data2/local_userdata/heyuting/datasets/fe108_coesot/test"
    # save_path = r"/home/ioe/xxxxx/COESOT/train"
    video_files = os.listdir(data_path)
    dvs_img_interval = 1
    voxel_generator = PointToVoxel(
        # 给定每个voxel的长宽高  [0.05, 0.05, 0.1]
        vsize_xyz=[50, 10, 10],  # [0.2, 0.25, 0.16]  # [50, 10, 10]  [50, 35, 26] 因此坐标范围（20,20,20）  (20, 34/35, 26)
        # 给定点云的范围 [  0.  -40.   -3.   70.4  40.    1. ]
        coors_range_xyz=[0, 0, 0, 1000, 345, 259],
        # 给定每个点云的特征维度，这里是x，y，z，r 其中r是激光雷达反射强度       # 346x260  t,x,y
        num_point_features=4,
        # 最多选取多少个voxel，训练16000，推理40000
        max_num_voxels=16000,  # 16000
        # 给定每个pillar中有采样多少个点，不够则补0  因此我将neg voxel改为-1;
        max_num_points_per_voxel=16,  # 1024
        device=device
    )

    for videoID in range(len(video_files)):
        foldName = video_files[videoID]
        print("==>> foldName: ", foldName)
        # if foldName in ['dvSave-2022_01_26_20_01_46', 'dvSave-2022_03_21_11_09_50',
        #                 'dvSave-2022_03_21_16_11_40', 'dvSave-2022_03_21_11_12_27',
        #                 'dvSave-2022_03_21_11_19_37', 'dvSave-2022_02_25_13_21_41',
        #                 'dvSave-2022_02_25_13_31_18', 'dvSave-2022_03_21_09_05_49']:
        #     print("==>> pass it: ", foldName)
        #     continue
        # if 'dvSave-2022_03_21' not in foldName:
        #     continue
        fileLIST = os.listdir(os.path.join(data_path, foldName))
        if not os.path.exists(os.path.join(save_path, foldName, foldName+'_voxel')):
            os.mkdir(os.path.join(save_path, foldName, foldName+'_voxel'))
        else:
            continue
        mat_save = os.path.join(save_path, foldName, foldName+'_voxel/')

        # for FileID in tqdm(range(len(fileLIST))):
        #     videoName = fileLIST[FileID]
        #     (filename, extension) = os.path.splitext(videoName)
        #     if (extension == '.aedat4') :  #  and ('dvSave-2022_03_21' in foldName):
        #         # if os.path.exists(mat_save + '{}.mat'.format(filename)):
        #         #     print("==>> Skip this video ... ")
        #         #     continue
        aedat4_file = 'events.aedat4'
        # print('filename:', aedat4_file)
        read_path = os.path.join(data_path, foldName, aedat4_file)

        # read aeda4;
        frame_all = []
        frame_exposure_time = []
        frame_interval_time = []
        with AedatFile(read_path) as f:
            # print(f.names)
            for frame in f['frames']:
                frame_all.append(frame.image)
                frame_exposure_time.append([frame.timestamp_start_of_exposure,
                                            frame.timestamp_end_of_exposure])  ## [1607928583397102, 1607928583401102]
                frame_interval_time.append([frame.timestamp_start_of_frame,
                                            frame.timestamp_end_of_frame])  ## [1607928583387944, 1607928583410285]
        if use_mode == 'frame_exposure_time':
            frame_timestamp = frame_exposure_time
        elif use_mode == 'frame_interval_time':
            frame_timestamp = frame_interval_time
        frame_num = len(frame_timestamp)

        events = np.hstack([packet for packet in f['events'].numpy()])

        t_all = torch.tensor(events['timestamp']).unsqueeze(1).to(device)
        x_all = torch.tensor(events['x']).unsqueeze(1).to(device)
        y_all = torch.tensor(events['y']).unsqueeze(1).to(device)
        p_all = torch.tensor(events['polarity']).unsqueeze(1).to(device)

        for frame_no in range(0, int(frame_num / dvs_img_interval) - 1):
            start_idx = np.where(events['timestamp'] >= frame_timestamp[frame_no][0])[0][0]
            end_idx = np.where(events['timestamp'] >= frame_timestamp[frame_no][1])[0][0]
            sub_event = events[start_idx:end_idx]
            idx_length = end_idx - start_idx

            t = t_all[start_idx: end_idx]
            if start_idx == end_idx:
                time_length = 0
                t = ((t-t).float() / time_length) * 1000
                scipy.io.savemat(mat_save + 'frame{:0>4d}.mat'.format(frame_no), mdict={'coor': np.zeros([100, 3]),
                                 'features': np.zeros([100, 16])})  # coor: numpy.nan;   features: numpy.nan
                print('empty event frame ', frame_no)
                continue
            else:
                time_length = t[-1] - t[0]
                # rescale the timestampes to start from 0 up to 1000
                t = ((t-t[0]).float() / time_length) * 1000
            all_idx = np.where(sub_event['polarity'] != -1)      # all event
            neg_idx = np.where(sub_event['polarity'] == 0)      # neg event
            t = t[all_idx]
            x = x_all[all_idx]
            y = y_all[all_idx]
            p = p_all[all_idx]
            p[neg_idx] = -1     # negtive voxel change from 0 to -1. because after append 0 operation.
            current_events = torch.cat((t, x, y, p), dim=1)
            # if current_events.shape[0] < 10:   # remove it
            #     continue
            data_dict = {'points': current_events}

            coor, features = transform_points_to_voxels(data_dict=data_dict, voxel_generator=voxel_generator,
                                                        device=device)
            # pdb.set_trace()
            coor = coor.cpu().numpy()
            features = features.cpu().numpy()
            # print('coor', coor)
            scipy.io.savemat(mat_save + 'frame{:0>4d}.mat'.format(frame_no), mdict={'coor': coor, 'features': features})   # coor: Nx(t,x,y);   features:Nx32 or Nx10024中for frame_no in range(0, int(frame_num / dvs_img_interval) - 1):
            start_idx = np.where(events['timestamp'] >= frame_timestamp[frame_no][0])[0][0]
            end_idx = np.where(events['timestamp'] >= frame_timestamp[frame_no][1])[0][0]
            sub_event = events[start_idx:end_idx]
            idx_length = end_idx - start_idx

            t = t_all[start_idx: end_idx]
            if start_idx == end_idx:
                time_length = 0
                t = ((t-t).float() / time_length) * 1000
                scipy.io.savemat(mat_save + 'frame{:0>4d}.mat'.format(frame_no), mdict={'coor': np.zeros([100, 3]),
                                 'features': np.zeros([100, 16])})  # coor: numpy.nan;   features: numpy.nan
                print('empty event frame ', frame_no)
                continue ，生成的文件是 'frame{:0>4d}.mat'.format(frame_no)形式，但是在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
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/test'
    save_path = '/data2/local_userdata/heyuting/datasets/test'
    device = torch.device("cuda:0")

    class_dir = os.listdir(data_path)
    class_dir.sort()
    # label_dict = {lbl: i for i, lbl in enumerate(class_dir)}
    for classID in range(len(class_dir)):
        foldName = class_dir[classID]
        label = classID
        fileLIST = os.listdir(os.path.join(data_path, foldName))
        voxel_dir = os.path.join(data_path, foldName, f'{foldName}_voxel')
        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     
        # if not os.path.exists(os.path.join(save_path,'v2g_8_r3_base', cls_name)):
        #     os.makedirs(os.path.join(save_path,'v2g_8_r3_base',  cls_name))
        for FileID in tqdm(range(len(fileLIST))):
            file_Name = fileLIST[FileID]
            save_name = file_Name.split('.')[0]
            if not os.path.exists(os.path.join(save_path, foldName, foldName+'_graph')):
                os.makedirs(os.path.join(save_path, foldName, foldName+'_graph'))
            if os.path.exists(os.path.join(save_path, foldName, foldName+'_graph', '{}.pt'.format(save_name))):
                continue
            read_path = os.path.join(voxel_dir, f'{save_name}.mat')
            # read_path = os.path.join(data_path, foldName,foldName+'_voxel','{}.mat'.format(save_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)中读取到的.mat形式不对，需要修改。