in multiple_futures_prediction/dataset_ngsim.py [0:0]
def __getitem__(self, idx_: int) -> Tuple[ List, List, Dict, Union[None, np.ndarray] ] :
idx = self.ind_random[idx_]
dsId_1b = self.D[idx, 0].astype(int)
vehId_1b = self.D[idx, 1].astype(int)
dsId_0b = dsId_1b-1
vehId_0b = vehId_1b-1
t = self.D[idx, 2]
grid = self.D[idx,8:] #1-based
ids = {} #0-based keys
leafs = [vehId_0b]
for _ in range( self.nbr_search_depth ):
new_leafs = []
for id0b in leafs:
nbr_dict = self.find_neighbors(dsId_0b, id0b, t)
if len(nbr_dict) > 0:
ids.update( nbr_dict )
if len(nbr_dict[id0b]) > 0:
nbr_id0b = list( zip(*nbr_dict[id0b]))[0]
new_leafs.extend ( nbr_id0b )
leafs = np.unique( new_leafs )
leafs = [x for x in leafs if x not in ids]
sorted_keys = sorted(ids.keys()) # e.g. [1, 3, 4, 5, ... , 74]
id_map = {key: value for (key, value) in zip(sorted_keys, np.arange(len(sorted_keys)).tolist()) }
#obj id to index within a batch
sz = len(ids)
assert sz > 0
hist = []
fut = []
neighbors: Dict[int,List] = {} # key is batch ind, followed by a list of (batch_ind, ego_id, grid/nbr ind)
for ind, vehId0b in enumerate(sorted_keys):
hist.append( self.getHistory0b(vehId0b,t,dsId_0b) ) # no normalization
fut.append( self.getFuture(vehId0b+1,t,dsId_0b+1 ) ) #subtract off ref pos
neighbors[ind] = []
for v_id, nbr_ind in ids[ vehId0b ]:
if v_id not in id_map:
k2 = -1
else:
k2 = id_map[v_id]
neighbors[ind].append( (k2, v_id, nbr_ind) )
if self.use_context:
x_range_ft = np.array([-15, 15])
y_range_ft = np.array([-30, 300])
pad = int(np.ceil(300/self.maps[dsId_1b-1]['ft_per_pixel'])) # max of all ranges
if not 'im_color' in self.maps[dsId_1b-1]:
im_big = np.pad(self.maps[dsId_1b-1]['im'], ((pad, pad), (pad, pad)), 'constant', constant_values= 0.0)
self.maps[dsId_1b-1]['im_color'] = (im_big[np.newaxis,...].repeat(3, axis=0)*255.0).astype(np.uint8)
im = self.maps[dsId_1b-1]['im_color']
height, width = im.shape[1:]
ref_pos = self.D[idx,3:5]
im_x, im_y = self.convert_pt( ref_pos, dsId_1b-1 )
im_x += pad
im_y += pad
x_range = (x_range_ft/self.maps[dsId_1b-1]['ft_per_pixel']).astype(int)
y_range = (y_range_ft/self.maps[dsId_1b-1]['ft_per_pixel']).astype(int)
x_range[0] = np.maximum( 0, x_range[0]+im_x )
x_range[1] = np.minimum( width-1, x_range[1]+im_x )
y_range[0] = np.maximum( 0, y_range[0]+im_y )
y_range[1] = np.minimum( height-1, y_range[1]+im_y )
im_crop = np.ascontiguousarray(im[:, y_range[0]:y_range[1], x_range[0]:x_range[1]].transpose((1,2,0)))
im_crop[:,:,[0, 1]] = 0
for _, other in neighbors.items():
if len(other) == 0:
continue
for k in range( len(other)-1 ):
x1, y1 = self.convert_pt( other[k]+ref_pos, dsId_1b-1 )
x2, y2 = self.convert_pt( other[k+1]+ref_pos, dsId_1b-1 )
x1+=pad; y1+=pad; x2+=pad; y2+=pad
cv2.line(im_crop,(x1-x_range[0],y1-y_range[0]),(x2-x_range[0],y2-y_range[0]), (255, 0, 0), 2 )
x, y = self.convert_pt( other[-1]+ref_pos, dsId_1b-1 )
x+=pad; y+=pad
cv2.circle(im_crop, (x-x_range[0], y-y_range[0]), 4, (255, 0, 0), -1)
for k in range( len(hist)-1 ):
x1, y1 = self.convert_pt( hist[k]+ref_pos, dsId_1b-1 )
x2, y2 = self.convert_pt( hist[k+1]+ref_pos, dsId_1b-1 )
x1+=pad; y1+=pad; x2+=pad; y2+=pad
cv2.line(im_crop,(x1-x_range[0],y1-y_range[0]),(x2-x_range[0],y2-y_range[0]), (0, 255, 0), 3 )
cv2.circle(im_crop, (im_x-x_range[0], im_y-y_range[0]), 5, (0, 255, 0), -1)
assert im_crop.shape == (660,60,3)
im_crop = im_crop.transpose((2,0,1))
else:
im_crop = None
return hist, fut, neighbors, im_crop # neighbors is a list of all vehicles in the batch