Function parse_node_centric

tbsim/utils/trajdata_utils.py:346–475

(batch: dict, overwrite_nan=True)

    return d

def parse_node_centric(batch: dict, overwrite_nan=True):
    maybe_pad_neighbor(batch)
    fut_pos, fut_yaw, _, fut_mask = trajdata2posyawspeed(batch["agent_fut"], nan_to_zero=overwrite_nan)
    hist_pos, hist_yaw, hist_speed, hist_mask = trajdata2posyawspeed(batch["agent_hist"], nan_to_zero=overwrite_nan)
    curr_speed = hist_speed[..., -1]
    curr_state = batch["curr_agent_state"]
    assert isinstance(curr_state,StateTensor) or isinstance(curr_state,StateArray)
    h1, h2 = curr_state[:, -1], curr_state.heading[...,0]
    p1, p2 = curr_state[:, :2], curr_state.position
    assert torch.all(h1[~torch.isnan(h1)] == h2[~torch.isnan(h2)])
    assert torch.all(p1[~torch.isnan(p1)] == p2[~torch.isnan(p2)])
    curr_yaw = curr_state.heading[...,0]
    curr_pos = curr_state.position

    # convert nuscenes types to l5kit types
    agent_type = batch["agent_type"]
    agent_type = convert_nusc_type_to_lyft_type(agent_type)
    
    # mask out invalid extents
    agent_hist_extent = batch["agent_hist_extent"]
    agent_hist_extent[torch.isnan(agent_hist_extent)] = 0.
    neigh_indices = batch["neigh_indices"]
    neigh_hist_pos, neigh_hist_yaw, neigh_hist_speed, neigh_hist_mask = trajdata2posyawspeed(batch["neigh_hist"], nan_to_zero=overwrite_nan)
    neigh_fut_pos, neigh_fut_yaw, _, neigh_fut_mask = trajdata2posyawspeed(batch["neigh_fut"], nan_to_zero=overwrite_nan)
    neigh_curr_speed = neigh_hist_speed[..., -1]
    neigh_types = batch["neigh_types"]
    # convert nuscenes types to l5kit types
    neigh_types = convert_nusc_type_to_lyft_type(neigh_types)
    # mask out invalid extents
    neigh_hist_extents = batch["neigh_hist_extents"]
    neigh_hist_extents[torch.isnan(neigh_hist_extents)] = 0.

    world_from_agents = torch.inverse(batch["agents_from_world_tf"])

    raster_cfg = BATCH_RASTER_CFG
    map_res = 1.0 / raster_cfg["pixel_size"] # convert to pixels/meter
    h = w = raster_cfg["raster_size"]
    ego_cent = raster_cfg["ego_center"]

    raster_from_agent = torch.Tensor([
            [map_res, 0, ((1.0 + ego_cent[0])/2.0) * w],
            [0, map_res, ((1.0 + ego_cent[1])/2.0) * h],
            [0, 0, 1]
    ]).to(curr_state.device)
    
    bsize = batch["agents_from_world_tf"].shape[0]
    agent_from_raster = torch.inverse(raster_from_agent)
    raster_from_agent = TensorUtils.unsqueeze_expand_at(raster_from_agent, size=bsize, dim=0)
    agent_from_raster = TensorUtils.unsqueeze_expand_at(agent_from_raster, size=bsize, dim=0)
    raster_from_world = torch.bmm(raster_from_agent, batch["agents_from_world_tf"])

    all_hist_pos = torch.cat((hist_pos[:, None], neigh_hist_pos.to(hist_pos.device)), dim=1)
    all_hist_yaw = torch.cat((hist_yaw[:, None], neigh_hist_yaw.to(hist_pos.device)), dim=1)
    all_hist_mask = torch.cat((hist_mask[:, None], neigh_hist_mask.to(hist_pos.device)), dim=1)

    maps_rasterize_in = batch["maps"]
    if maps_rasterize_in is None and BATCH_RASTER_CFG["include_hist"]:
        maps_rasterize_in = torch.empty((bsize, 0, h, w)).to(all_hist_pos.device)