AgentDataset and EgoDataset mapping about l5kit HOT 4 CLOSED

If I've understood correctly you should be able to pull it off by using the get_frame_indices function in the AgentDataset. That gives you the indices in the AgentDataset itself for that given frame. The Ego information for that frame can be retrieved from the EgoDataset by indexing at that frame, something like this:

frame_idx = 0
agent_indices = agent_dt.get_frame_indices(frame_idx)  # this function is slow, but there is a branch with an improved version
agents = [agent_dt[idx] for idx in agent_indices]
ego_info = ego_dt[frame_idx]  # this is the ego for those agents

Note this works on Dataset, not DataLoader (sadly)

Let me know if this helps you!

from l5kit.

Thanks @lucabergamini , will try it out and get back. If possible, can you name which branch you were referring to as the improved one?

from l5kit.

Hey @suryajayaraman,
sure! the branch is this one. You may not need it, it depends on the number of accesses.

from l5kit.

Note, that branch has been merged into master :)
I'm closing this, but feel free to reopen it if you need further assistance

from l5kit.