Big update for v1.0.7, introducing a standardized map format, speedup…

…s, new features, and more!

.gitignore

@@ -1,5 +1,5 @@
 public/
-.vscode/launch.json
+.vscode/
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

README.md

@@ -83,8 +83,8 @@ Currently, the dataloader supports interfacing with the following datasets:
 
 | Dataset | ID | Splits | Add'l Tags | Description | dt | Maps |
 |---------|----|--------|------------|-------------|----|------|
-| nuScenes Train/Val | `nusc_trainval` | `train`, `val` | `boston`, `singapore` | nuScenes' training/validation splits (700/150 scenes) | 0.5s (2Hz) | :white_check_mark: |
-| nuScenes Test | `nusc_test` | `test` | `boston`, `singapore` | nuScenes' test split (150 scenes) | 0.5s (2Hz) | :white_check_mark: |
+| nuScenes Train/TrainVal/Val | `nusc_trainval` | `train`, `train_val`, `val` | `boston`, `singapore` | nuScenes prediction challenge training/validation/test splits (500/200/150 scenes) | 0.5s (2Hz) | :white_check_mark: |
+| nuScenes Test | `nusc_test` | `test` | `boston`, `singapore` | nuScenes' test split, no annotations (150 scenes) | 0.5s (2Hz) | :white_check_mark: |
 | nuScenes Mini | `nusc_mini` | `mini_train`, `mini_val` | `boston`, `singapore` | nuScenes mini training/validation splits (8/2 scenes) | 0.5s (2Hz) | :white_check_mark: |
 | Lyft Level 5 Train | `lyft_train` | `train` | `palo_alto` | Lyft Level 5 training data - part 1/2 (8.4 GB) | 0.1s (10Hz) | :white_check_mark: |
 | Lyft Level 5 Train Full | `lyft_train_full` | `train` | `palo_alto` | Lyft Level 5 training data - part 2/2 (70 GB) | 0.1s (10Hz) | :white_check_mark: |

examples/batch_example.py

@@ -17,7 +17,7 @@ def main():
         desired_dt=0.1,
         history_sec=(3.2, 3.2),
         future_sec=(4.8, 4.8),
-        only_types=[AgentType.VEHICLE],
+        only_predict=[AgentType.VEHICLE],
         agent_interaction_distances=defaultdict(lambda: 30.0),
         incl_robot_future=False,
         incl_map=True,

examples/custom_batch_data.py

@@ -0,0 +1,112 @@
+"""
+This is an example of how to extend a batch to include custom data
+"""
+
+from collections import defaultdict
+from functools import partial
+from typing import Tuple, Union
+
+import numpy as np
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+from trajdata import AgentBatch, AgentType, UnifiedDataset
+from trajdata.augmentation import NoiseHistories
+from trajdata.data_structures.batch_element import AgentBatchElement, SceneBatchElement
+from trajdata.visualization.vis import plot_agent_batch
+
+
+def custom_random_data(
+    batch_elem: Union[AgentBatchElement, SceneBatchElement]
+) -> np.ndarray:
+    # create new data to add to each batch element
+    return np.random.random((10, 10))
+
+
+def custom_goal_location(
+    batch_elem: Union[AgentBatchElement, SceneBatchElement]
+) -> np.ndarray:
+    # simply access existing element attributes
+    return batch_elem.agent_future_np[:, :2]
+
+
+def custom_min_distance_from_others(
+    batch_elem: Union[AgentBatchElement, SceneBatchElement]
+) -> np.ndarray:
+    # ... or more complicated calculations
+    current_ego_loc = batch_elem.agent_history_np[-1, :2]
+    all_distances = [
+        np.linalg.norm(current_ego_loc - veh[-1, :2])
+        for veh in batch_elem.neighbor_histories
+    ]
+
+    if not len(all_distances):
+        return np.inf
+    else:
+        return np.min(all_distances)
+
+
+def custom_distances_squared(
+    batch_elem: Union[AgentBatchElement, SceneBatchElement]
+) -> np.ndarray:
+    # we can chain extras together if needed
+    return batch_elem.extras["min_distance"] ** 2
+
+
+def custom_raster(
+    batch_elem: Union[AgentBatchElement, SceneBatchElement],
+    raster_size: Tuple[int, ...],
+) -> np.ndarray:
+    # draw a custom raster
+    img = np.zeros(raster_size)
+
+    # ...
+    return img
+
+
+def main():
+    dataset = UnifiedDataset(
+        desired_data=["nusc_mini-mini_train"],
+        centric="agent",
+        desired_dt=0.1,
+        history_sec=(3.2, 3.2),
+        future_sec=(4.8, 4.8),
+        only_types=[AgentType.VEHICLE],
+        agent_interaction_distances=defaultdict(lambda: 30.0),
+        incl_robot_future=False,
+        incl_map=True,
+        map_params={"px_per_m": 2, "map_size_px": 224, "offset_frac_xy": (-0.5, 0.0)},
+        num_workers=0,
+        verbose=True,
+        data_dirs={  # Remember to change this to match your filesystem!
+            "nusc_mini": "~/datasets/nuScenes",
+        },
+        extras={  # a dictionary that contains functions that generate our custom data. Can be any function and has access to the batch element.
+            "random_data": custom_random_data,
+            "goal_location": custom_goal_location,
+            "min_distance": custom_min_distance_from_others,
+            "min_distance_sq": custom_distances_squared,  # in Python >= 3.7 dictionaries are guaranteed to maintain order => you can use previously computed keys
+            "raster": partial(custom_raster, raster_size=(100, 100)),
+        },
+    )
+
+    print(f"# Data Samples: {len(dataset):,}")
+
+    dataloader = DataLoader(
+        dataset,
+        batch_size=4,
+        shuffle=True,
+        collate_fn=dataset.get_collate_fn(),
+        num_workers=4,
+    )
+
+    batch: AgentBatch
+    for batch in tqdm(dataloader):
+        assert "random_data" in batch.extras
+        assert "goal_location" in batch.extras
+        assert "min_distance" in batch.extras
+        assert "raster" in batch.extras
+
+
+if __name__ == "__main__":
+    main()

setup.cfg

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

src/trajdata/augmentation/noise_histories.py

@@ -2,21 +2,49 @@
 
 from trajdata.augmentation.augmentation import BatchAugmentation
 from trajdata.data_structures.batch import AgentBatch, SceneBatch
+from trajdata.utils.arr_utils import PadDirection, mask_up_to
 
 
 class NoiseHistories(BatchAugmentation):
-    def __init__(self, mean: float = 0.0, stddev: float = 0.1) -> None:
+    def __init__(
+        self,
+        mean: float = 0.0,
+        stddev: float = 0.1,
+    ) -> None:
         self.mean = mean
         self.stddev = stddev
 
     def apply_agent(self, agent_batch: AgentBatch) -> None:
-        agent_batch.agent_hist[..., :-1, :] += torch.normal(
-            self.mean, self.stddev, size=agent_batch.agent_hist[..., :-1, :].shape
+        agent_hist_noise = torch.normal(
+            self.mean, self.stddev, size=agent_batch.agent_hist.shape
         )
-        agent_batch.neigh_hist[..., :-1, :] += torch.normal(
-            self.mean, self.stddev, size=agent_batch.neigh_hist[..., :-1, :].shape
+        neigh_hist_noise = torch.normal(
+            self.mean, self.stddev, size=agent_batch.neigh_hist.shape
         )
 
+        if agent_batch.history_pad_dir == PadDirection.BEFORE:
+            agent_hist_noise[..., -1, :] = 0
+            neigh_hist_noise[..., -1, :] = 0
+        else:
+            len_mask = ~mask_up_to(
+                agent_batch.agent_hist_len,
+                delta=-1,
+                max_len=agent_batch.agent_hist.shape[1],
+            ).unsqueeze(-1)
+            agent_hist_noise[len_mask.expand(-1, -1, agent_hist_noise.shape[-1])] = 0
+
+            len_mask = ~mask_up_to(
+                agent_batch.neigh_hist_len,
+                delta=-1,
+                max_len=agent_batch.neigh_hist.shape[2],
+            ).unsqueeze(-1)
+            neigh_hist_noise[
+                len_mask.expand(-1, -1, -1, neigh_hist_noise.shape[-1])
+            ] = 0
+
+        agent_batch.agent_hist += agent_hist_noise
+        agent_batch.neigh_hist += neigh_hist_noise
+
     def apply_scene(self, scene_batch: SceneBatch) -> None:
         scene_batch.agent_hist[..., :-1, :] += torch.normal(
             self.mean, self.stddev, size=scene_batch.agent_hist[..., :-1, :].shape