data.py NGSIM数据处理的理解

1.读取NGSIM dataset prepocess后的mat文件

打开指定路径下的mat文件，分别读取其中的traj和tracks保存至self.Data及self.Tracks，若文件不存在则报错
变量初始化

### Class for the highway trajectory datasets (NISIM, HighD, etc.)
class highwayTrajDataset(Dataset):

    def __init__(self, path, t_h=30, t_f=50, d_s=2,
                 enc_size=64, targ_enc_size=112, grid_size=(25, 5), fit_plan_traj=False, fit_plan_further_ds=1):
        if not os.path.exists(path):
            raise RuntimeError("{} not exists!!".format(path))
        if path.endswith('.mat'):
            f = h5py.File(path, 'r')
            f_tracks = f['tracks']
            track_cols, track_rows = f_tracks.shape
            self.Data = np.transpose(f['traj'])
            self.Tracks = []
            for i in range(track_rows):
                self.Tracks.append([  np.transpose(f[f_tracks[j][i]][:]) for j in range(track_cols)  ])
        else:
            raise RuntimeError("Path should be end with '.mat' for file or '/' for folder")

        # If torch version >= 1.2.0
        if int(torch.__version__[0])>=1 and int(torch.__version__[2])>=2:
            self.mask_num_type = torch.bool
        else:
            self.mask_num_type = torch.uint8

        self.t_h = t_h  # length of track history.
        self.t_f = t_f  # length of predicted trajectory.
        self.d_s = d_s  # downsampling rate of all trajectories to be processed.
        self.enc_size = enc_size
        self.targ_enc_size = targ_enc_size
        self.hist_len = self.t_h // self.d_s + 1  # data length of the history trajectory
        self.fut_len = self.t_f // self.d_s       # data length of the future  trajectory
        self.plan_len = self.t_f // self.d_s      # data length of the planning  trajectory

        self.fit_plan_traj = fit_plan_traj              # Fitting the future planned trajectory in testing/evaluation.
        self.further_ds_plan = fit_plan_further_ds      # Further downsampling to restrict the planning info

        self.cell_length = 8
        self.cell_width = 7
        self.grid_size = grid_size                # size of social context grid
        self.grid_cells = grid_size[0] * grid_size[1]
        self.grid_length = self.cell_length * grid_size[0]
        self.grid_width = self.cell_width * grid_size[1]

2.信息检索

定义一些信息检索常用的方法
itsDsId
given item index return dataset id 数据集id
itsPlanVehId
given item index return vehicle id 车辆id
itsTime
given item index return frame time 时刻
itsLocation
given item index return location(LocalX LocalY) 车辆位置
itsPlanVehBehavior
given item index return plan vehicle behavior(turn left or right or go straight, brake or normal) 左转右转直行 刹车or正常行驶
itsPlanVehSize
given item index return vehicle size(length, width) 车长车宽
itsPlanVehDynamic
given item index return vehicle dynamic(speed and acc) 车辆速度和加速度
itsCentGrid
given item index return central grid 以plan vehicle为中心 长方形邻域每个grid内的vehicle的id(不存在则为0)
itsTargVehsId
given item index return target vehicle id 上述central grid去0得到邻域车辆id
itsNbrVehsId
given item index return neighbor vehicle id 邻域车辆的邻域车辆id(all existing neighbor vehicles id)
itsTargsCentLoc
given item index return target central location 邻域车辆的位置(all existing target vehicles location)
itsAllAroundSizes
given item index return all around sizes 返回plan vehicle size, target vehicle size, neighbor vehicle size

## Functions of retrieving information according to the item's index
    def itsDsId(self, idx):
        return self.Data[idx, 0].astype(int)

    def itsPlanVehId(self, idx):
        return self.Data[idx, 1].astype(int)

    def itsTime(self, idx):
        return self.Data[idx, 2]

    def itsLocation(self, idx):
        return self.Data[idx, 3:5]

    def itsPlanVehBehavior(self, idx):
        return int(self.Data[idx, 6] + (self.Data[idx, 7] - 1) * 3)

    def itsPlanVehSize(self, idx):
        return self.Data[idx, 8:10]

    def itsPlanVehDynamic(self, idx):
        planVel, planAcc = self.getDynamic(self.itsDsId(idx), self.itsPlanVehId(idx), self.itsTime(idx))
        return planVel, planAcc

    def itsCentGrid(self, idx):
        return self.Data[idx, 13:].astype(int)

    def itsTargVehsId(self, idx):
        centGrid = self.itsCentGrid(idx)
        targVehsId = centGrid[np.nonzero(centGrid)]
        return targVehsId

    def itsNbrVehsId(self, idx):
        dsId = self.itsDsId(idx)
        planVehId = self.itsPlanVehId(idx)
        targVehsId = self.itsTargVehsId(idx)
        t = self.itsTime(idx)
        nbrVehsId = np.array([], dtype=np.int64)
        for target in targVehsId:
            subGrid = self.getGrid(dsId, target, t)
            subIds = subGrid[np.nonzero(subGrid)]
            for i in subIds:
                if i==planVehId or any(i==targVehsId) or any(i==nbrVehsId):
                    continue
                else:
                    nbrVehsId = np.append(nbrVehsId, i)
        return nbrVehsId

    def itsTargsCentLoc(self, idx):
        dsId = self.itsDsId(idx)
        t = self.itsTime(idx)
        centGrid = self.itsCentGrid(idx)
        targsCenterLoc = np.empty((0,2), dtype=np.float32)
        for target in centGrid:
            if target:
                targsCenterLoc = np.vstack([targsCenterLoc, self.getLocation(dsId, target, t)])
        return torch.from_numpy(targsCenterLoc)

    def itsAllAroundSizes(self, idx):
        dsId = self.itsDsId(idx)
        centGrid = self.itsCentGrid(idx)
        t = self.itsTime(idx)
        planVehSize = []
        targVehSizes = []
        nbsVehSizes = []
        planVehSize.append(self.getSize(dsId, self.itsPlanVehId(idx)))
        for i, target in enumerate(centGrid):
            if target:
                targVehSizes.append(self.getSize(dsId, target))
                targVehGrid = self.getGrid(dsId, target, t)
                for targetNb in targVehGrid:
                    if targetNb:
                        nbsVehSizes.append(self.getSize(dsId, targetNb))
        return np.asarray(planVehSize), np.asarray(targVehSizes), np.asarray(nbsVehSizes)

3.轨迹检索

给定item index，即哪个数据集中的哪条数据，获得该vehicle历史和将来的轨迹、targets vehicle历史和将来的轨迹、neighbor vehicle历史和将来的轨迹。
absPlanTraj
given item index return history and future vehicle trajectory
absTargsTraj
given item index return target vehicles history and future trajectory
absNbrsTraj
given item index return neighbor vehicles history and future trajectory
其实这三种方法获取轨迹的实现方式都差不多，以absPlanTraj为例。首先根据item index找到对应的dataset id, vehicle id, frame time，得到该车在相应数据集相应时间的列索引colIndex以及该车的全部track vehTrack，之后根据时间索引往前回溯30帧得到历史轨迹(X Y坐标)planHis，往后前进50帧得到将来轨迹(X Y坐标)planFut，返回[planHis, planFut]。

## Functions for retrieving trajectory data with absolute coordinate, mainly used for visualization
    def itsAllGroundTruthTrajs(self, idx):
        return [self.absPlanTraj(idx), self.absTargsTraj(idx), self.absNbrsTraj(idx)]

    def absPlanTraj(self, idx):
        dsId = self.itsDsId(idx)
        planVeh = self.itsPlanVehId(idx)
        t = self.itsTime(idx)
        colIndex = np.where(self.Tracks[dsId - 1][planVeh - 1][0, :] == t)[0][0]
        vehTrack = self.Tracks[dsId - 1][planVeh - 1].transpose()
        planHis = vehTrack[np.maximum(0, colIndex - self.t_h): (colIndex + 1): self.d_s, 1:3]
        planFut = vehTrack[(colIndex + self.d_s): (colIndex + self.t_f + 1): self.d_s, 1:3]
        return [planHis, planFut]

    def absTargsTraj(self, idx):
        dsId = self.itsDsId(idx)
        targVehs = self.itsTargVehsId(idx)
        t = self.itsTime(idx)
        targHisList, targFutList = [], []
        for targVeh in targVehs:
            colIndex = np.where(self.Tracks[dsId - 1][targVeh - 1][0, :] == t)[0][0]
            vehTrack = self.Tracks[dsId - 1][targVeh - 1].transpose()
            targHis = vehTrack[np.maximum(0, colIndex - self.t_h): (colIndex + 1): self.d_s, 1:3]
            targFut = vehTrack[(colIndex + self.d_s): (colIndex + self.t_f + 1): self.d_s, 1:3]
            targHisList.append(targHis)
            targFutList.append(targFut)
        return [targHisList, targFutList]

    def absNbrsTraj(self, idx):
        dsId = self.itsDsId(idx)
        nbrVehs = self.itsNbrVehsId(idx)
        t = self.itsTime(idx)
        nbrHisList, nbrFutList = [], []
        for nbrVeh in nbrVehs:
            colIndex = np.where(self.Tracks[dsId - 1][nbrVeh - 1][0, :] == t)[0][0]
            vehTrack = self.Tracks[dsId - 1][nbrVeh - 1].transpose()
            targHis = vehTrack[np.maximum(0, colIndex - self.t_h): (colIndex + 1): self.d_s, 1:3]
            nbrHisList.append(targHis)
        return [nbrHisList, nbrFutList]

    def batchTargetVehsInfo(self, idxs):
        count = 0
        dsIds = np.zeros(len(idxs)*self.grid_cells, dtype=int)
        vehIds = np.zeros(len(idxs)*self.grid_cells, dtype=int)
        for idx in idxs:
            dsId = self.itsDsId(idx)
            targets = self.itsCentGrid(idx)
            targetsIndex = np.nonzero(targets)
            for index in targetsIndex[0]:
                dsIds[count] = dsId
                vehIds[count] = targets[index]
                count += 1
        return [dsIds[:count], vehIds[:count]]

4.位置检索

定义位置、行为、动力学等检索方法
getTracksCol
given dataset id, vehicle id, frame time return column index in track
getLocation
given dataset id, vehicle id, frame time return vehicle location(X Y)
getManeuver
given dataset id, vehicle id, frame time return vehicle maneuver
getGrid
given dataset id, vehicle id, frame time return target vehicles id
getDynamic
given dataset id, vehicle id, frame time return vehicle's dynamic (velocity & acceleration)
getSize
given dataset id, vehicle id, frame time return vehicle's size (length & width)

## Avoid searching the correspond column for too many times.
    def getTracksCol(self, dsId, vehId, t):
        return np.where(self.Tracks[dsId - 1][vehId - 1][0, :] == t)[0][0]

    ## Get the vehicle's location from tracks
    def getLocation(self, dsId, vehId, t):
        colIndex = np.where(self.Tracks[dsId - 1][vehId - 1][0, :] == t)[0][0]
        location = self.getLocationByCol(dsId, vehId, colIndex)
        return location
    def getLocationByCol(self, dsId, vehId, colIndex):
        return self.Tracks[dsId - 1][vehId - 1][1:3, colIndex].transpose()

    ## Get the vehicle's maneuver given dataset id, vehicle id and time point t.
    def getManeuver(self, dsId, vehId, t):
        colIndex = np.where(self.Tracks[dsId - 1][vehId - 1][0, :] == t)[0][0]
        lat_lon_maneuvers = self.getManeuverByCol(dsId, vehId, colIndex)
        return lat_lon_maneuvers
    def getManeuverByCol(self, dsId, vehId, colIndex):
        return self.Tracks[dsId - 1][vehId - 1][4:6, colIndex].astype(int)

    ## Get the vehicle's nearby neighbours
    def getGrid(self, dsId, vehId, t):
        colIndex = np.where(self.Tracks[dsId - 1][vehId - 1][0, :] == t)[0][0]
        grid = self.getGridByCol(dsId, vehId, colIndex)
        return grid
    def getGridByCol(self, dsId, vehId, colIndex):
        return self.Tracks[dsId - 1][vehId - 1][11:, colIndex].astype(int)

    ## Get the vehicle's dynamic (velocity & acceleration) given dataset id, vehicle id and time point t.
    def getDynamic(self, dsId, vehId, t):
        colIndex = np.where(self.Tracks[dsId - 1][vehId - 1][0, :] == t)[0][0]
        vel_acc = self.getDynamicByCol(dsId, vehId, colIndex)
        return vel_acc
    def getDynamicByCol(self, dsId, vehId, colIndex):
        return self.Tracks[dsId - 1][vehId - 1][9:11, colIndex]

    ## Get the vehicle's size (length & width) given dataset id and vehicle id
    def getSize(self, dsId, vehId):
        length_width = self.Tracks[dsId - 1][vehId - 1][6:8, 0]
        return length_width

5.辅助函数

getHistory
Helper function to get track history
given dataset id, vehicle id, reference vehicle id, frame time return track history
getPlanFuture
Helper function to get track future
given dataset id, vehicle id, reference vehicle id, frame time return track future

## Helper function to get track history
    def getHistory(self, dsId, vehId, refVehId, t, wholePeriod=False):
        if vehId == 0:
            # if return empty, it denotes there's no vehicle in that grid.
            return np.empty([0, 2])
        else:
            vehColIndex = np.where(self.Tracks[dsId - 1][vehId - 1][0, :] == t)[0][0]
            refColIndex = np.where(self.Tracks[dsId - 1][refVehId - 1][0, :] == t)[0][0]
            vehTrack = self.Tracks[dsId - 1][vehId - 1][1:3].transpose()
            refTrack = self.Tracks[dsId - 1][refVehId - 1][1:3].transpose()
            # Use the sequence of trajectory or just the last instance as the refPos
            if wholePeriod:
                refStpt = np.maximum(0, refColIndex - self.t_h)
                refEnpt = refColIndex + 1
                refPos = refTrack[refStpt:refEnpt:self.d_s, :]
            else:
                refPos = np.tile(refTrack[refColIndex, :], (self.hist_len, 1))
            stpt = np.maximum(0, vehColIndex - self.t_h)
            enpt = vehColIndex + 1
            vehPos = vehTrack[stpt:enpt:self.d_s, :]
            if len(vehPos) < self.hist_len:
                histPart = vehPos - refPos[-len(vehPos)::]
                paddingPart = np.tile(histPart[0, :], (self.hist_len - len(vehPos), 1))
                hist = np.concatenate((paddingPart, histPart), axis=0)
                return hist
            else:
                hist = vehPos - refPos
                return hist

    ## Helper function to get track future
    def getFuture(self, dsId, vehId, t):
        colIndex = np.where(self.Tracks[dsId - 1][vehId - 1][0, :] == t)[0][0]
        futTraj = self.getFutureByCol(dsId, vehId, colIndex)
        return futTraj
    def getFutureByCol(self, dsId, vehId, colIndex):
        vehTrack = self.Tracks[dsId - 1][vehId - 1].transpose()
        refPos = self.Tracks[dsId - 1][vehId - 1][1:3, colIndex].transpose()
        stpt = colIndex + self.d_s
        enpt = np.minimum(len(vehTrack), colIndex + self.t_f + 1)
        futTraj = vehTrack[stpt:enpt:self.d_s, 1:3] - refPos
        return futTraj

    def getPlanFuture(self, dsId, planId, refVehId, t):
        # Traj of the reference veh
        refColIndex = np.where(self.Tracks[dsId - 1][refVehId - 1][0, :] == t)[0][0]
        refPos = self.Tracks[dsId - 1][refVehId - 1][1:3, refColIndex].transpose()
        # Traj of the planned veh
        planColIndex = np.where(self.Tracks[dsId - 1][planId - 1][0, :] == t)[0][0]
        stpt = planColIndex
        enpt = planColIndex + self.t_f + 1
        planGroundTrue = self.Tracks[dsId - 1][planId - 1][1:3, stpt:enpt:self.d_s].transpose()
        planFut = planGroundTrue.copy()
        # Fitting the downsampled waypoints as the planned trajectory in testing and evaluation.
        if self.fit_plan_traj:
            wayPoint        = np.arange(0, self.t_f + self.d_s, self.d_s)
            wayPoint_to_fit = np.arange(0, self.t_f + 1, self.d_s * self.further_ds_plan)
            planFut_to_fit = planFut[::self.further_ds_plan, ]
            laterParam = fitting_traj_by_qs(wayPoint_to_fit, planFut_to_fit[:, 0])
            longiParam = fitting_traj_by_qs(wayPoint_to_fit, planFut_to_fit[:, 1])
            planFut[:, 0] = quintic_spline(wayPoint, *laterParam)
            planFut[:, 1] = quintic_spline(wayPoint, *longiParam)
        revPlanFut = np.flip(planFut[1:,] - refPos, axis=0).copy()
        return revPlanFut

    def __getitem__(self, idx):
        dsId = self.itsDsId(idx)
        centVehId = self.itsPlanVehId(idx)
        t = self.itsTime(idx)
        centGrid = self.itsCentGrid(idx)
        planGridLocs = []
        targsHists = []
        targsFuts = []
        targsLonEnc = []
        targsLatEnc = []
        nbsHists = []
        planFuts = []
        targsVehs = np.zeros(self.grid_cells)
        for id, target in enumerate(centGrid):
            if target:
                targetColumn = self.getTracksCol(dsId, target, t)
                # Get the grid of each neighbour vehicle.
                grid = self.getGridByCol(dsId, target, targetColumn)
                # Targets history and future
                targsVehs[id] = target
                targsHists.append(self.getHistory(dsId, target, target, t))
                targsFuts.append(self.getFutureByCol(dsId, target, targetColumn))
                # Targets maneuvers
                latMan, lonMan = self.getManeuverByCol(dsId, target, targetColumn)
                lat_enc = np.zeros([3])
                lon_enc = np.zeros([2])
                lat_enc[latMan - 1] = 1
                lon_enc[lonMan - 1] = 1
                targsLatEnc.append(lat_enc)
                targsLonEnc.append(lon_enc)
                # Neighbours history
                nbsHists.append([self.getHistory(dsId, i, target, t, wholePeriod=True) for i in grid])
                # PlanVeh future
                planGridLocs.append(np.where(grid == centVehId)[0][0])
                planFuts.append(self.getPlanFuture(dsId, centVehId, target, t))
        return planFuts, nbsHists, \
               targsHists, targsFuts, targsLonEnc, targsLatEnc, \
               centGrid, planGridLocs, idx


    ## Collate function for dataloader
    def collate_fn(self, samples):
        targs_batch_size = 0
        nbs_batch_size = 0
        for _, nbsHists, targsHists, _, _, _, _, _, _ in samples:
            targs_batch_size += len(targsHists)
            nbs_number = [sum([len(nbs) > 0 for nbs in sub_nbsHist]) for sub_nbsHist in nbsHists]
            nbs_batch_size += sum(nbs_number)
        # Initialize all things
        nbsHist_batch   = torch.zeros(self.hist_len,  nbs_batch_size,   2)
        targsHist_batch = torch.zeros(self.hist_len,  targs_batch_size, 2)
        targsFut_batch  = torch.zeros(self.fut_len,   targs_batch_size, 2)
        lat_enc_batch   = torch.zeros(targs_batch_size, 3)
        lon_enc_batch   = torch.zeros(targs_batch_size, 2)
        planFut_batch   = torch.zeros(self.plan_len,   targs_batch_size, 2)
        idxs = []
        pos = [0, 0]
        # Fill 1 on those grid locations with neighbour
        nbsMask_batch      = torch.zeros(targs_batch_size, self.grid_size[1], self.grid_size[0], self.enc_size, dtype=self.mask_num_type)
        planMask_batch     = torch.zeros(targs_batch_size, self.grid_size[1], self.grid_size[0], self.enc_size, dtype=self.mask_num_type)
        targsEncMask_batch = torch.zeros(len(samples),     self.grid_size[1], self.grid_size[0], self.targ_enc_size, dtype=self.mask_num_type)
        targsFutMask_batch = torch.zeros(self.fut_len,     targs_batch_size,  2)
        targetCount = 0
        nbCount = 0
        for i, (planFuts, nbsHists, targsHists, targsFuts, targsLonEnc, targsLatEnc, centGrid, planGridLocs, idx) in enumerate(samples):
            idxs.append(idx)
            centGridIndex = centGrid.nonzero()[0]
            for j in range(len(targsFuts)):
                targsHist_batch[0:len(targsHists[j]), targetCount, 0] = torch.from_numpy(targsHists[j][:, 0])
                targsHist_batch[0:len(targsHists[j]), targetCount, 1] = torch.from_numpy(targsHists[j][:, 1])
                targsFut_batch[0:len(targsFuts[j]), targetCount, 0] = torch.from_numpy(targsFuts[j][:, 0])
                targsFut_batch[0:len(targsFuts[j]), targetCount, 1] = torch.from_numpy(targsFuts[j][:, 1])
                targsFutMask_batch[0:len(targsFuts[j]), targetCount, :] = 1
                pos[0] = centGridIndex[j] % self.grid_size[0]
                pos[1] = centGridIndex[j] // self.grid_size[0]
                targsEncMask_batch[i, pos[1], pos[0], :] = torch.ones(self.targ_enc_size).byte()
                lat_enc_batch[targetCount, :] = torch.from_numpy(targsLatEnc[j])
                lon_enc_batch[targetCount, :] = torch.from_numpy(targsLonEnc[j])
                planFut_batch[0:len(planFuts[j]), targetCount, 0] = torch.from_numpy(planFuts[j][:, 0])
                planFut_batch[0:len(planFuts[j]), targetCount, 1] = torch.from_numpy(planFuts[j][:, 1])
                # Set up neighbor, neighbor sequence length, and mask batches:
                for index, nbHist in enumerate(nbsHists[j]):
                    if len(nbHist) != 0:
                        nbsHist_batch[0:len(nbHist), nbCount, 0] = torch.from_numpy(nbHist[:, 0])
                        nbsHist_batch[0:len(nbHist), nbCount, 1] = torch.from_numpy(nbHist[:, 1])
                        pos[0] = index % self.grid_size[0]
                        pos[1] = index // self.grid_size[0]
                        nbsMask_batch[targetCount, pos[1], pos[0], :] = torch.ones(self.enc_size).byte()
                        nbCount += 1
                        if index == planGridLocs[j]:
                            planMask_batch[targetCount, pos[1], pos[0], :] = torch.ones(self.enc_size).byte()
                targetCount += 1
        return nbsHist_batch, nbsMask_batch, \
               planFut_batch, planMask_batch, \
               targsHist_batch, targsEncMask_batch, \
               targsFut_batch, targsFutMask_batch, lat_enc_batch, lon_enc_batch, idxs