data.py NGSIM数据处理的理解
2022-05-30 本文已影响0人
吃醋不吃辣的雷儿
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
