Adding a few new features like specifying the max number of agents an…

…d if only the ego should be the focus agent, as well as a hotfix bringing back rasterized maps in their original format.

setup.cfg

@@ -1,6 +1,6 @@
 [metadata]
 name = trajdata
-version = 1.0.7
+version = 1.0.8
 author = Boris Ivanovic
 author_email = [email protected]
 description = A unified interface to many trajectory forecasting datasets.

src/trajdata/caching/df_cache.py

@@ -637,13 +637,14 @@ def cache_map(
     @staticmethod
     def cache_map_layers(
         cache_path: Path,
+        vec_map: VectorizedMap,
         map_info: RasterizedMapMetadata,
         layer_fn: Callable[[str], np.ndarray],
         env_name: str,
     ) -> None:
         (
             maps_path,
-            _,
+            vector_map_path,
             raster_map_path,
             raster_metadata_path,
         ) = DataFrameCache.get_map_paths(
@@ -653,10 +654,16 @@ def cache_map_layers(
         # Ensuring the maps directory exists.
         maps_path.mkdir(parents=True, exist_ok=True)
 
+        # Saving the vectorized map data.
+        with open(vector_map_path, "wb") as f:
+            f.write(vec_map.SerializeToString())
+
+        # Saving the rasterized map data.
         disk_data = zarr.open_array(raster_map_path, mode="w", shape=map_info.shape)
         for idx, layer_name in enumerate(map_info.layers):
             disk_data[idx] = layer_fn(layer_name)
 
+        # Saving the rasterized map metadata.
         with open(raster_metadata_path, "wb") as f:
             dill.dump(map_info, f)
 

src/trajdata/caching/scene_cache.py

@@ -140,6 +140,7 @@ def cache_map(
     @staticmethod
     def cache_map_layers(
         cache_path: Path,
+        vec_map: VectorizedMap,
         map_info: RasterizedMapMetadata,
         layer_fn: Callable[[str], np.ndarray],
         env_name: str,

src/trajdata/data_structures/batch_element.py

@@ -1,6 +1,6 @@
 from collections import defaultdict
 from math import sqrt
-from typing import Callable, Dict, List, Optional, Tuple
+from typing import Any, Callable, Dict, List, Optional, Tuple
 
 import numpy as np
 
@@ -26,9 +26,10 @@ def __init__(
         ] = defaultdict(lambda: np.inf),
         incl_robot_future: bool = False,
         incl_map: bool = False,
-        map_params: Optional[Dict[str, int]] = None,
+        map_params: Optional[Dict[str, Any]] = None,
         standardize_data: bool = False,
         standardize_derivatives: bool = False,
+        max_neighbor_num: Optional[int] = None,
     ) -> None:
         self.cache: SceneCache = cache
         self.data_index: int = data_index
@@ -40,6 +41,7 @@ def __init__(
         agent_info: AgentMetadata = scene_time_agent.agent
         self.agent_name: str = agent_info.name
         self.agent_type: AgentType = agent_info.type
+        self.max_neighbor_num = max_neighbor_num
 
         self.curr_agent_state_np: np.ndarray = cache.get_state(
             agent_info.name, self.scene_ts
@@ -156,82 +158,86 @@ def get_agent_future(
         )
         return agent_future_np, agent_extent_future_np
 
-    def get_neighbor_history(
+    # @profile
+    def get_neighbor_data(
         self,
         scene_time: SceneTimeAgent,
         agent_info: AgentMetadata,
-        history_sec: Tuple[Optional[float], Optional[float]],
+        length_sec: Tuple[Optional[float], Optional[float]],
         distance_limit: Callable[[np.ndarray, int], np.ndarray],
+        mode: str,
     ) -> Tuple[int, np.ndarray, List[np.ndarray], List[np.ndarray], np.ndarray]:
-        # The indices of the returned ndarray match the scene_time agents list (including the index of the central agent,
-        # which would have a distance of 0 to itself).
+        # The indices of the returned ndarray match the scene_time agents list
+        # (including the index of the central agent, which would have a distance
+        # of 0 to itself).
         agent_distances: np.ndarray = scene_time.get_agent_distances_to(agent_info)
         agent_idx: int = scene_time.agents.index(agent_info)
+
         neighbor_types: np.ndarray = np.array([a.type.value for a in scene_time.agents])
         nearby_mask: np.ndarray = agent_distances <= distance_limit(
             neighbor_types, agent_info.type
         )
         nearby_mask[agent_idx] = False
 
+        nb_idx = agent_distances.argsort()
         nearby_agents: List[AgentMetadata] = [
-            agent for (idx, agent) in enumerate(scene_time.agents) if nearby_mask[idx]
+            scene_time.agents[idx] for idx in nb_idx if nearby_mask[idx]
         ]
         neighbor_types_np: np.ndarray = neighbor_types[nearby_mask]
 
+        if self.max_neighbor_num is not None:
+            # Pruning nearby_agents and re-creating
+            # neighbor_types_np with the remaining agents.
+            nearby_agents = nearby_agents[: self.max_neighbor_num]
+            neighbor_types_np: np.ndarray = np.array(
+                [a.type.value for a in nearby_agents]
+            )
+
         num_neighbors: int = len(nearby_agents)
-        (
-            neighbor_histories,
-            neighbor_history_extents,
-            neighbor_history_lens_np,
-        ) = self.cache.get_agents_history(self.scene_ts, nearby_agents, history_sec)
+
+        if mode == "history":
+            (
+                neighbor_data,
+                neighbor_extents_data,
+                neighbor_data_lens_np,
+            ) = self.cache.get_agents_history(self.scene_ts, nearby_agents, length_sec)
+        elif mode == "future":
+            (
+                neighbor_data,
+                neighbor_extents_data,
+                neighbor_data_lens_np,
+            ) = self.cache.get_agents_future(self.scene_ts, nearby_agents, length_sec)
+        else:
+            raise ValueError(f"Unknown mode {mode} passed in!")
 
         return (
             num_neighbors,
             neighbor_types_np,
-            neighbor_histories,
-            neighbor_history_extents,
-            neighbor_history_lens_np,
+            neighbor_data,
+            neighbor_extents_data,
+            neighbor_data_lens_np,
         )
 
-    # @profile
-    def get_neighbor_future(
+    def get_neighbor_history(
         self,
         scene_time: SceneTimeAgent,
         agent_info: AgentMetadata,
-        future_sec: Tuple[Optional[float], Optional[float]],
+        history_sec: Tuple[Optional[float], Optional[float]],
         distance_limit: Callable[[np.ndarray, int], np.ndarray],
     ) -> Tuple[int, np.ndarray, List[np.ndarray], List[np.ndarray], np.ndarray]:
-        scene_ts: int = self.scene_ts
-
-        # The indices of the returned ndarray match the scene_time agents list (including the index of the central agent,
-        # which would have a distance of 0 to itself).
-        agent_distances: np.ndarray = scene_time.get_agent_distances_to(agent_info)
-        agent_idx: int = scene_time.agents.index(agent_info)
-
-        neighbor_types: np.ndarray = np.array([a.type.value for a in scene_time.agents])
-        nearby_mask: np.ndarray = agent_distances <= distance_limit(
-            neighbor_types, agent_info.type
+        return self.get_neighbor_data(
+            scene_time, agent_info, history_sec, distance_limit, mode="history"
         )
-        nearby_mask[agent_idx] = False
-
-        nearby_agents: List[AgentMetadata] = [
-            agent for (idx, agent) in enumerate(scene_time.agents) if nearby_mask[idx]
-        ]
-        neighbor_types_np: np.ndarray = neighbor_types[nearby_mask]
-
-        num_neighbors: int = len(nearby_agents)
-        (
-            neighbor_futures,
-            neighbor_future_extents,
-            neighbor_future_lens_np,
-        ) = self.cache.get_agents_future(scene_ts, nearby_agents, future_sec)
 
-        return (
-            num_neighbors,
-            neighbor_types_np,
-            neighbor_futures,
-            neighbor_future_extents,
-            neighbor_future_lens_np,
+    def get_neighbor_future(
+        self,
+        scene_time: SceneTimeAgent,
+        agent_info: AgentMetadata,
+        future_sec: Tuple[Optional[float], Optional[float]],
+        distance_limit: Callable[[np.ndarray, int], np.ndarray],
+    ) -> Tuple[int, np.ndarray, List[np.ndarray], List[np.ndarray], np.ndarray]:
+        return self.get_neighbor_data(
+            scene_time, agent_info, future_sec, distance_limit, mode="future"
         )
 
     def get_robot_current_and_future(
@@ -310,7 +316,7 @@ def __init__(
         ] = defaultdict(lambda: np.inf),
         incl_robot_future: bool = False,
         incl_map: bool = False,
-        map_params: Optional[Dict[str, int]] = None,
+        map_params: Optional[Dict[str, Any]] = None,
         standardize_data: bool = False,
         standardize_derivatives: bool = False,
         max_agent_num: Optional[int] = None,

src/trajdata/dataset.py

@@ -3,7 +3,7 @@
 from functools import partial
 from itertools import chain
 from pathlib import Path
-from typing import Callable, Dict, Final, List, Optional, Set, Tuple, Union
+from typing import Any, Callable, Dict, Final, List, Optional, Set, Tuple, Union
 
 import numpy as np
 from torch.utils.data import DataLoader, Dataset
@@ -61,14 +61,16 @@ def __init__(
         ] = defaultdict(lambda: np.inf),
         incl_robot_future: bool = False,
         incl_map: bool = False,
-        map_params: Optional[Dict[str, int]] = None,
+        map_params: Optional[Dict[str, Any]] = None,
         only_types: Optional[List[AgentType]] = None,
         only_predict: Optional[List[AgentType]] = None,
         no_types: Optional[List[AgentType]] = None,
         standardize_data: bool = True,
         standardize_derivatives: bool = False,
         augmentations: Optional[List[Augmentation]] = None,
         max_agent_num: Optional[int] = None,
+        max_neighbor_num: Optional[int] = None,
+        ego_only: Optional[bool] = False,
         data_dirs: Dict[str, str] = {
             # "nusc_trainval": "~/datasets/nuScenes",
             # "nusc_test": "~/datasets/nuScenes",
@@ -104,14 +106,16 @@ def __init__(
             agent_interaction_distances: (Dict[Tuple[AgentType, AgentType], float]): A dictionary mapping agent-agent interaction distances in meters (determines which agents are included as neighbors to the predicted agent). Defaults to infinity for all types.
             incl_robot_future (bool, optional): Include the ego agent's future trajectory in batches (accordingly, never predict the ego's future). Defaults to False.
             incl_map (bool, optional): Include a local cropping of the rasterized map (if the dataset provides a map) per agent. Defaults to False.
-            map_params (Optional[Dict[str, int]], optional): Local map cropping parameters, must be specified if incl_map is True. Must contain keys {"px_per_m", "map_size_px"} and can optionally contain {"offset_frac_xy"}. Defaults to None.
+            map_params (Optional[Dict[str, Any]], optional): Local map cropping parameters, must be specified if incl_map is True. Must contain keys {"px_per_m", "map_size_px"} and can optionally contain {"offset_frac_xy"}. Defaults to None.
             only_types (Optional[List[AgentType]], optional): Filter out all agents EXCEPT for those of the specified types. Defaults to None.
             only_predict (Optional[List[AgentType]], optional): Only predict the specified types of agents. Importantly, this keeps other agent types in the scene, e.g., as neighbors of the agent to be predicted. Defaults to None.
             no_types (Optional[List[AgentType]], optional): Filter out all agents with the specified types. Defaults to None.
             standardize_data (bool, optional): Standardize all data such that (1) the predicted agent's orientation at the current timestep is 0, (2) all data is made relative to the predicted agent's current position, and (3) the agent's heading value is replaced with its sin, cos values. Defaults to True.
             standardize_derivatives (bool, optional): Make agent velocities and accelerations relative to the agent being predicted. Defaults to False.
             augmentations (Optional[List[Augmentation]], optional): Perform the specified augmentations to the batch or dataset. Defaults to None.
             max_agent_num (int, optional): The maximum number of agents to include in a batch for scene-centric batching.
+            max_neighbor_num (int, optional): The maximum number of neighbors to include in a batch for agent-centric batching.
+            ego_only (bool, optional): If True, only return batches where the ego-agent is the one being predicted.
             data_dirs (Optional[Dict[str, str]], optional): Dictionary mapping dataset names to their directories on disk. Defaults to { "eupeds_eth": "~/datasets/eth_ucy_peds", "eupeds_hotel": "~/datasets/eth_ucy_peds", "eupeds_univ": "~/datasets/eth_ucy_peds", "eupeds_zara1": "~/datasets/eth_ucy_peds", "eupeds_zara2": "~/datasets/eth_ucy_peds", "nusc_mini": "~/datasets/nuScenes", "lyft_sample": "~/datasets/lyft/scenes/sample.zarr", }.
             cache_type (str, optional): What type of cache to use to store preprocessed, cached data on disk. Defaults to "dataframe".
             cache_location (str, optional): Where to store and load preprocessed, cached data. Defaults to "~/.unified_data_cache".
@@ -157,6 +161,8 @@ def __init__(
         self.extras = extras
         self.verbose = verbose
         self.max_agent_num = max_agent_num
+        self.max_neighbor_num = max_neighbor_num
+        self.ego_only = ego_only
 
         # Ensuring scene description queries are all lowercase
         if scene_description_contains is not None:
@@ -178,6 +184,7 @@ def __init__(
             if any(env.name in dataset_tuple for dataset_tuple in matching_datasets):
                 all_data_cached: bool = False
                 all_maps_cached: bool = not env.has_maps or not self.incl_map
+
                 if self.env_cache.env_is_cached(env.name) and not self.rebuild_cache:
                     scenes_list: List[Scene] = self.get_desired_scenes_from_env(
                         matching_datasets, scene_description_contains, env
@@ -224,7 +231,7 @@ def __init__(
                         env.cache_maps(
                             self.cache_path,
                             self.cache_class,
-                            resolution=self.map_params["px_per_m"],
+                            self.map_params,
                         )
 
                     scenes_list: List[SceneMetadata] = self.get_desired_scenes_from_env(
@@ -294,7 +301,10 @@ def get_data_index(
                 history_sec=self.history_sec,
                 future_sec=self.future_sec,
                 desired_dt=self.desired_dt,
+                ego_only=self.ego_only,
             )
+        else:
+            raise ValueError(f"{self.centric}-centric data batches are not supported.")
 
         # data_index is either:
         # [(scene_path, total_index_len, valid_scene_ts)] for scene-centric data, or
@@ -381,6 +391,7 @@ def _get_data_index_agent(
         future_sec: Tuple[Optional[float], Optional[float]],
         desired_dt: Optional[float],
         ret_scene_info: bool = False,
+        ego_only: bool = False,
     ) -> Tuple[Optional[Scene], Path, int, List[Tuple[str, np.ndarray]]]:
         index_elems_len: int = 0
         index_elems: List[Tuple[str, np.ndarray]] = list()
@@ -395,10 +406,15 @@ def _get_data_index_agent(
         )
 
         for agent_info in filtered_agents:
-            # Don't want to predict the ego if we're going to be giving the model its future!
+            # Don't want to predict the ego if we're going to be
+            # giving the model its future!
             if incl_robot_future and agent_info.name == "ego":
                 continue
 
+            if ego_only and agent_info.name != "ego":
+                # We only want to return the ego.
+                continue
+
             valid_ts: Tuple[int, int] = filtering.get_valid_ts(
                 agent_info, scene.dt, history_sec, future_sec
             )
@@ -443,6 +459,8 @@ def get_collate_fn(
                 pad_format=pad_format,
                 batch_augments=batch_augments,
             )
+        else:
+            raise ValueError(f"{self.centric}-centric data batches are not supported.")
 
         return collate_fn
 
@@ -698,6 +716,7 @@ def __getitem__(self, idx: int) -> Union[SceneBatchElement, AgentBatchElement]:
                 self.map_params,
                 self.standardize_data,
                 self.standardize_derivatives,
+                self.max_neighbor_num,
             )
 
         for key, extra_fn in self.extras.items():

src/trajdata/dataset_specific/eth_ucy_peds/eupeds_dataset.py

@@ -1,5 +1,5 @@
 from pathlib import Path
-from typing import Dict, Final, List, Optional, Tuple, Type
+from typing import Any, Dict, Final, List, Optional, Tuple, Type
 
 import numpy as np
 import pandas as pd
@@ -325,7 +325,10 @@ def cache_map(
         pass
 
     def cache_maps(
-        self, cache_path: Path, map_cache_class: Type[SceneCache], resolution: float
+        self,
+        cache_path: Path,
+        map_cache_class: Type[SceneCache],
+        map_params: Dict[str, Any],
     ) -> None:
         """
         No maps in this dataset!