Adding support for the INTERACTION dataset and SDD pedestrian dataset!

.gitignore

@@ -1,5 +1,8 @@
 public/
 .vscode/
+*.html
+*.mp4
+*.avi
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

DATASETS.md

@@ -96,6 +96,56 @@ It should look like this after downloading:
 
 **Note**: Not all the dataset parts need to be downloaded, only the necessary `.zarr` files need to be downloaded (e.g., `sample.zarr` for the small sample dataset).
 
+## INTERACTION Dataset
+Nothing special needs to be done for the INTERACTION Dataset, simply download it as per [the instructions on the dataset website](http://interaction-dataset.com/).
+
+It should look like this after downloading:
+```
+/path/to/interaction_single/
+            ├── maps/
+            │   ├── DR_CHN_Merging_ZS0.osm
+            |   ├── DR_CHN_Merging_ZS0.osm_xy
+            |   └── ...
+            ├── test_conditional-single-agent/
+            │   ├── DR_CHN_Merging_ZS0_obs.csv
+            |   ├── DR_CHN_Merging_ZS2_obs.csv
+            |   └── ...
+            └── test_single-agent/
+            │   ├── DR_CHN_Merging_ZS0_obs.csv
+            |   ├── DR_CHN_Merging_ZS2_obs.csv
+            |   └── ...
+            └── train/
+            │   ├── DR_CHN_Merging_ZS0_train.csv
+            |   ├── DR_CHN_Merging_ZS2_train.csv
+            |   └── ...
+            └── val/
+                ├── DR_CHN_Merging_ZS0_val.csv
+                ├── DR_CHN_Merging_ZS2_val.csv
+                └── ...
+
+/path/to/interaction_multi/
+            ├── maps/
+            │   ├── DR_CHN_Merging_ZS0.osm
+            |   ├── DR_CHN_Merging_ZS0.osm_xy
+            |   └── ...
+            ├── test_conditional-multi-agent/
+            │   ├── DR_CHN_Merging_ZS0_obs.csv
+            |   ├── DR_CHN_Merging_ZS2_obs.csv
+            |   └── ...
+            └── test_multi-agent/
+            │   ├── DR_CHN_Merging_ZS0_obs.csv
+            |   ├── DR_CHN_Merging_ZS2_obs.csv
+            |   └── ...
+            └── train/
+            │   ├── DR_CHN_Merging_ZS0_train.csv
+            |   ├── DR_CHN_Merging_ZS2_train.csv
+            |   └── ...
+            └── val/
+                ├── DR_CHN_Merging_ZS0_val.csv
+                ├── DR_CHN_Merging_ZS2_val.csv
+                └── ...
+```
+
 ## ETH/UCY Pedestrians
 The raw data can be found in many places online, ranging from [research projects' data download scripts](https://github.com/agrimgupta92/sgan/blob/master/scripts/download_data.sh) to [copies of the original data itself](https://github.com/StanfordASL/Trajectron-plus-plus/tree/master/experiments/pedestrians/raw/raw/all_data) on GitHub. In this data loader, we assume the data was sourced from the latter.
 
@@ -111,3 +161,30 @@ It should look like this after downloading:
             ├── students003.txt
             └── uni_examples.txt
 ```
+
+## Stanford Drone Dataset
+The raw data can be found in many places online, the easiest is probably [this space-optimized version](https://www.kaggle.com/datasets/aryashah2k/stanford-drone-dataset) on Kaggle.
+
+It should look like this after downloading:
+```
+/path/to/sdd/
+            ├── bookstore/
+            |   ├── video0
+            |       ├── annotations.txt
+            |       └── reference.jpg
+            |   ├── video1
+            |       ├── annotations.txt
+            |       └── reference.jpg
+            |   └── ...
+            ├── coupa/
+            |   ├── video0
+            |       ├── annotations.txt
+            |       └── reference.jpg
+            |   ├── video1
+            |       ├── annotations.txt
+            |       └── reference.jpg
+            |   └── ...
+            └── ...
+```
+
+**Note**: Only the annotations need to be downloaded (not the videos).

README.md

@@ -13,7 +13,7 @@ The easiest way to install trajdata is through PyPI with
 pip install trajdata
 ```
 
-In case you would also like to use datasets such as nuScenes, Lyft Level 5, or Waymo Open Motion Dataset (which require their own devkits to access raw data), the following will also install the respective devkits.
+In case you would also like to use datasets such as nuScenes, Lyft Level 5, or Waymo Open Motion Dataset (which require their own devkits to access raw data or additional package dependencies), the following will also install the respective devkits and/or package dependencies.
 ```sh
 # For nuScenes
 pip install "trajdata[nusc]"
@@ -24,10 +24,13 @@ pip install "trajdata[lyft]"
 # For Waymo
 pip install "trajdata[waymo]"
 
+# For INTERACTION
+pip install "trajdata[interaction]"
+
 # All
-pip install "trajdata[nusc,lyft,waymo]"
+pip install "trajdata[nusc,lyft,waymo,interaction]"
 ```
-Then, download the raw datasets (nuScenes, Lyft Level 5, ETH/UCY, etc) in case you do not already have them. For more information about how to structure dataset folders/files, please see [`DATASETS.md`](./DATASETS.md).
+Then, download the raw datasets (nuScenes, Lyft Level 5, ETH/UCY, etc.) in case you do not already have them. For more information about how to structure dataset folders/files, please see [`DATASETS.md`](./DATASETS.md).
 
 ### Package Developer Installation
 
@@ -84,7 +87,7 @@ For more information on all of the possible `UnifiedDataset` constructor argumen
 ## Supported Datasets
 Currently, the dataloader supports interfacing with the following datasets:
 
-| Dataset | ID | Splits | Add'l Tags | Description | dt | Maps |
+| Dataset | ID | Splits | Locations | Description | dt | Maps |
 |---------|----|--------|------------|-------------|----|------|
 | 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: |
@@ -97,11 +100,14 @@ Currently, the dataloader supports interfacing with the following datasets:
 | 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: |
 | Lyft Level 5 Validation | `lyft_val` | `val` | `palo_alto` | Lyft Level 5 validation data (8.2 GB) | 0.1s (10Hz) | :white_check_mark: |
 | Lyft Level 5 Sample | `lyft_sample` | `mini_train`, `mini_val` | `palo_alto` | Lyft Level 5 sample data (100 scenes, randomly split 80/20 for training/validation) | 0.1s (10Hz) | :white_check_mark: |
+| INTERACTION Dataset Single-Agent | `interaction_single` | `train`, `val`, `test`, `test_conditional` | `usa`, `china`, `germany`, `bulgaria` | Single-agent split of the INTERACTION Dataset (where the goal is to predict one target agents' future motion) | 0.1s (10Hz) | :white_check_mark: |
+| INTERACTION Dataset Multi-Agent | `interaction_multi` | `train`, `val`, `test`, `test_conditional` | `usa`, `china`, `germany`, `bulgaria` | Multi-agent split of the INTERACTION Dataset (where the goal is to jointly predict multiple agents' future motion) | 0.1s (10Hz) | :white_check_mark: |
 | ETH - Univ | `eupeds_eth` | `train`, `val`, `train_loo`, `val_loo`, `test_loo` | `zurich` | The ETH (University) scene from the ETH BIWI Walking Pedestrians dataset | 0.4s (2.5Hz) | |
 | ETH - Hotel | `eupeds_hotel` | `train`, `val`, `train_loo`, `val_loo`, `test_loo` | `zurich` | The Hotel scene from the ETH BIWI Walking Pedestrians dataset | 0.4s (2.5Hz) | |
 | UCY - Univ | `eupeds_univ` | `train`, `val`, `train_loo`, `val_loo`, `test_loo` | `cyprus` | The University scene from the UCY Pedestrians dataset | 0.4s (2.5Hz) | |
 | UCY - Zara1 | `eupeds_zara1` | `train`, `val`, `train_loo`, `val_loo`, `test_loo` | `cyprus` | The Zara1 scene from the UCY Pedestrians dataset | 0.4s (2.5Hz) | |
 | UCY - Zara2 | `eupeds_zara2` | `train`, `val`, `train_loo`, `val_loo`, `test_loo` | `cyprus` | The Zara2 scene from the UCY Pedestrians dataset | 0.4s (2.5Hz) | |
+| Stanford Drone Dataset | `sdd` | `train`, `val`, `test` | `stanford` | Stanford Drone Dataset (60 scenes, randomly split 42/9/9 (70%/15%/15%) for training/validation/test) | 0.0333...s (30Hz) | |
 
 ### Adding New Datasets
 The code that interfaces the original datasets (dealing with their unique formats) can be found in `src/trajdata/dataset_specific`.

setup.cfg

@@ -1,6 +1,6 @@
 [metadata]
 name = trajdata
-version = 1.3.1
+version = 1.3.2
 author = Boris Ivanovic
 author_email = [email protected]
 description = A unified interface to many trajectory forecasting datasets.
@@ -52,3 +52,5 @@ waymo =
     tensorflow==2.11.0
     waymo-open-dataset-tf-2-11-0
     intervaltree
+interaction =
+    lanelet2==1.2.1

src/trajdata/caching/env_cache.py

@@ -48,6 +48,20 @@ def save_scene(self, scene: Scene) -> Path:
 
         return scene_file
 
+    @staticmethod
+    def save_scene_with_path(base_path: Path, scene: Scene) -> Path:
+        scene_file: Path = EnvCache.scene_metadata_path(
+            base_path, scene.env_name, scene.name, scene.dt
+        )
+
+        scene_cache_dir: Path = scene_file.parent
+        scene_cache_dir.mkdir(parents=True, exist_ok=True)
+
+        with open(scene_file, "wb") as f:
+            dill.dump(scene, f)
+
+        return scene_file
+
     def load_env_scenes_list(self, env_name: str) -> List[NamedTuple]:
         env_cache_dir: Path = self.path / env_name
         with open(env_cache_dir / "scenes_list.dill", "rb") as f:

src/trajdata/dataset.py

@@ -1,4 +1,5 @@
 import gc
+import random
 import time
 from collections import defaultdict
 from functools import partial
@@ -704,7 +705,9 @@ def get_desired_scenes_from_env(
         env: RawDataset,
     ) -> Union[List[Scene], List[SceneMetadata]]:
         scenes_list: Union[List[Scene], List[SceneMetadata]] = list()
-        for scene_tag in scene_tags:
+        for scene_tag in tqdm(
+            scene_tags, desc=f"Getting Scenes from {env.name}", disable=not self.verbose
+        ):
             if env.name in scene_tag:
                 scenes_list += env.get_matching_scenes(
                     scene_tag,
@@ -742,6 +745,10 @@ def preprocess_scene_data(
                 for scene_info in scenes_list
                 if self.envs_dict[scene_info.env_name].parallelizable
             ]
+
+            # Fixing the seed for random suffling (for debugging and reproducibility).
+            shuffle_rng = random.Random(123)
+            shuffle_rng.shuffle(parallel_scenes)
         else:
             serial_scenes = scenes_list
             parallel_scenes = list()

src/trajdata/dataset_specific/__init__.py

@@ -1,2 +1 @@
 from .raw_dataset import RawDataset
-from .scene_records import EUPedsRecord, LyftSceneRecord, NuscSceneRecord

src/trajdata/dataset_specific/interaction/__init__.py

@@ -0,0 +1 @@
+from .interaction_dataset import InteractionDataset