Initial streaming support.

conda/environment.yml

@@ -16,3 +16,4 @@ dependencies:
   - pyproj
   - rich
   - scipy
+  - universal_pathlib

integration_tests/verify_tbv_download.py

@@ -140,7 +140,6 @@ def verify_log_map(data_root: Path, log_id: str) -> None:
     # every log should have one and only one raster height map. (Note: season is stripped from uuid here).
     ground_height_raster_fpaths = list(log_map_dirpath.glob("*_ground_height_surface____*.npy"))
     assert len(ground_height_raster_fpaths) == 1
-    ground_height_raster_fpath = ground_height_raster_fpaths[0]
 
     # every log should have a Sim(2) mapping from raster grid coordinates to city coordinates.
     Sim2_fpaths = list(log_map_dirpath.glob("*___img_Sim2_city.json"))
@@ -181,10 +180,10 @@ def verify_log_map(data_root: Path, log_id: str) -> None:
         assert left_lane_boundary.ndim == 2 and left_lane_boundary.shape[1] == 3
 
     # load every pedestrian crossing
-    pcs = avm.get_scenario_ped_crossings()
+    avm.get_scenario_ped_crossings()
 
     # load every drivable area
-    das = avm.get_scenario_vector_drivable_areas()
+    avm.get_scenario_vector_drivable_areas()
 
 
 def verify_logs_using_dataloader(data_root: Path, log_ids: List[str]) -> None:

setup.cfg

@@ -31,6 +31,7 @@ install_requires =
     pyproj
     rich
     scipy
+    universal_pathlib
 
 package_dir=
     =src

src/av2/evaluation/detection/eval.py

@@ -51,8 +51,12 @@
     in addition to the mean statistics average across all classes, and P refers to the number of included statistics,
     e.g. AP, ATE, ASE, AOE, CDS by default.
 """
+import itertools
 import logging
-from multiprocessing import get_context
+import multiprocessing as mp
+import warnings
+from math import inf
+from statistics import mean
 from typing import Dict, Final, List, Optional, Tuple
 
 import numpy as np
@@ -72,6 +76,8 @@
 from av2.utils.io import TimestampedCitySE3EgoPoses
 from av2.utils.typing import NDArrayBool, NDArrayFloat
 
+warnings.filterwarnings("ignore", module="google")
+
 TP_ERROR_COLUMNS: Final[Tuple[str, ...]] = tuple(x.value for x in TruePositiveErrorNames)
 DTS_COLUMN_NAMES: Final[Tuple[str, ...]] = tuple(ORDERED_CUBOID_COL_NAMES) + ("score",)
 GTS_COLUMN_NAMES: Final[Tuple[str, ...]] = tuple(ORDERED_CUBOID_COL_NAMES) + ("num_interior_pts",)
@@ -161,7 +167,7 @@ def evaluate(
         args_list.append(args)
 
     logger.info("Starting evaluation ...")
-    with get_context("spawn").Pool(processes=n_jobs) as p:
+    with mp.get_context("spawn").Pool(processes=n_jobs) as p:
         outputs: Optional[List[Tuple[NDArrayFloat, NDArrayFloat]]] = p.starmap(accumulate, args_list)
 
     if outputs is None:
@@ -192,7 +198,7 @@ def summarize_metrics(
         dts: (N,14) Table of detections.
         gts: (M,15) Table of ground truth annotations.
         cfg: Detection configuration.
-
+        
     Returns:
         The summary metrics.
     """

src/av2/evaluation/detection/utils.py

@@ -13,11 +13,12 @@
 import logging
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Dict, List, Optional, Tuple
+from typing import Dict, List, Optional, Tuple, Union
 
 import numpy as np
 from joblib import Parallel, delayed
 from scipy.spatial.distance import cdist
+from upath import UPath
 
 from av2.evaluation.detection.constants import (
     MAX_NORMALIZED_ASE,
@@ -63,7 +64,7 @@ class DetectionCfg:
     affinity_thresholds_m: Tuple[float, ...] = (0.5, 1.0, 2.0, 4.0)
     affinity_type: AffinityType = AffinityType.CENTER
     categories: Tuple[str, ...] = tuple(x.value for x in CompetitionCategories)
-    dataset_dir: Optional[Path] = None
+    dataset_dir: Optional[Union[Path, UPath]] = None
     eval_only_roi_instances: bool = True
     filter_metric: FilterMetricType = FilterMetricType.EUCLIDEAN
     max_num_dts_per_category: int = 100
@@ -315,7 +316,7 @@ def compute_average_precision(
     # Evaluate precision at different recalls.
     precision_interpolated: NDArrayFloat = np.interp(recall_interpolated, recall, precision, right=0)
 
-    average_precision: float = np.mean(precision_interpolated)
+    average_precision: float = float(np.mean(precision_interpolated))
     return average_precision, precision_interpolated
 
 
@@ -437,7 +438,7 @@ def compute_evaluated_gts_mask(
 
 
 def load_mapped_avm_and_egoposes(
-    log_ids: List[str], dataset_dir: Path
+    log_ids: List[str], dataset_dir: Union[Path, UPath]
 ) -> Tuple[Dict[str, ArgoverseStaticMap], Dict[str, TimestampedCitySE3EgoPoses]]:
     """Load the maps and egoposes for each log in the dataset directory.
 
@@ -452,9 +453,10 @@ def load_mapped_avm_and_egoposes(
         RuntimeError: If the process for loading maps and timestamped egoposes fails.
     """
     log_id_to_timestamped_poses = {log_id: read_city_SE3_ego(dataset_dir / log_id) for log_id in log_ids}
-    avms: Optional[List[ArgoverseStaticMap]] = Parallel(n_jobs=-1, backend="threading", verbose=1)(
+    avms: Optional[List[ArgoverseStaticMap]] = Parallel(n_jobs=-1, backend="threading")(
         delayed(ArgoverseStaticMap.from_map_dir)(dataset_dir / log_id / "map", build_raster=True) for log_id in log_ids
     )
+
     if avms is None:
         raise RuntimeError("Map and egopose loading has failed!")
     log_id_to_avm = {log_ids[i]: avm for i, avm in enumerate(avms)}

src/av2/geometry/sim2.py

@@ -14,8 +14,10 @@
 import numbers
 from dataclasses import dataclass
 from pathlib import Path
+from typing import Union
 
 import numpy as np
+from upath import UPath
 
 import av2.utils.io as io_utils
 from av2.utils.helpers import assert_np_array_shape
@@ -187,9 +189,9 @@ def save_as_json(self, save_fpath: Path) -> None:
         io_utils.save_json_dict(save_fpath, dict_for_serialization)
 
     @classmethod
-    def from_json(cls, json_fpath: Path) -> Sim2:
+    def from_json(cls, json_fpath: Union[Path, UPath]) -> Sim2:
         """Generate class inst. from a JSON file containing Sim(2) parameters as flattened matrices (row-major)."""
-        with open(json_fpath, "r") as f:
+        with json_fpath.open("r") as f:
             json_data = json.load(f)
 
         R: NDArrayFloat = np.array(json_data["R"]).reshape(2, 2)

src/av2/map/map_api.py

@@ -20,16 +20,18 @@
 from pathlib import Path
 from typing import Dict, Final, List, Optional, Tuple, Union
 
+import fsspec
 import numpy as np
+from upath import UPath
 
 import av2.geometry.interpolate as interp_utils
 import av2.utils.dilation_utils as dilation_utils
-import av2.utils.io as io_utils
 import av2.utils.raster as raster_utils
 from av2.geometry.sim2 import Sim2
 from av2.map.drivable_area import DrivableArea
 from av2.map.lane_segment import LaneSegment
 from av2.map.pedestrian_crossing import PedestrianCrossing
+from av2.utils import io
 from av2.utils.typing import NDArrayBool, NDArrayByte, NDArrayFloat, NDArrayInt
 
 # 1 meter resolution is insufficient for the online-generated drivable area and ROI raster grids
@@ -105,7 +107,7 @@ class GroundHeightLayer(RasterMapLayer):
     """
 
     @classmethod
-    def from_file(cls, log_map_dirpath: Path) -> GroundHeightLayer:
+    def from_file(cls, log_map_dirpath: Union[Path, UPath]) -> GroundHeightLayer:
         """Load ground height values (w/ values at 30 cm resolution) from .npy file, and associated Sim(2) mapping.
 
         Note: ground height values are stored on disk as a float16 2d-array, but cast to float32 once loaded for
@@ -130,7 +132,8 @@ def from_file(cls, log_map_dirpath: Path) -> GroundHeightLayer:
             raise RuntimeError("Sim(2) mapping from city to image coordinates is missing")
 
         # load the file with rasterized values
-        ground_height_array: NDArrayFloat = np.load(ground_height_npy_fpaths[0])
+        with ground_height_npy_fpaths[0].open("rb") as f:
+            ground_height_array: NDArrayFloat = np.load(f)
 
         array_Sim2_city = Sim2.from_json(Sim2_json_fpaths[0])
 
@@ -311,7 +314,7 @@ class ArgoverseStaticMap:
     raster_ground_height_layer: Optional[GroundHeightLayer]
 
     @classmethod
-    def from_json(cls, static_map_path: Path) -> ArgoverseStaticMap:
+    def from_json(cls, static_map_path: Union[Path, UPath]) -> ArgoverseStaticMap:
         """Instantiate an Argoverse static map object (without raster data) from a JSON file containing map data.
 
         Args:
@@ -322,7 +325,7 @@ def from_json(cls, static_map_path: Path) -> ArgoverseStaticMap:
             An Argoverse HD map.
         """
         log_id = static_map_path.stem.split("log_map_archive_")[1]
-        vector_data = io_utils.read_json_file(static_map_path)
+        vector_data = io.read_json_file(static_map_path)
 
         vector_drivable_areas = {da["id"]: DrivableArea.from_dict(da) for da in vector_data["drivable_areas"].values()}
         vector_lane_segments = {ls["id"]: LaneSegment.from_dict(ls) for ls in vector_data["lane_segments"].values()}
@@ -346,7 +349,7 @@ def from_json(cls, static_map_path: Path) -> ArgoverseStaticMap:
         )
 
     @classmethod
-    def from_map_dir(cls, log_map_dirpath: Path, build_raster: bool = False) -> ArgoverseStaticMap:
+    def from_map_dir(cls, log_map_dirpath: Union[Path, UPath], build_raster: bool = False) -> ArgoverseStaticMap:
         """Instantiate an Argoverse map object from data stored within a map data directory.
 
         Note: the ground height surface file and associated coordinate mapping is not provided for the

src/av2/structures/cuboid.py

@@ -318,7 +318,7 @@ def project_to_cam(
 
         # Sort by z-order to respect visibility in the scene.
         # i.e, closer objects cuboids should be drawn on top of farther objects.
-        z_orders: NDArrayFloat = np.argsort(-z_buffer)
+        z_orders: NDArrayInt = np.argsort(-z_buffer)
 
         cuboids_vertices_cam = cuboids_vertices_cam[z_orders]
         front_face_indices = [0, 1, 2, 3, 0]

src/av2/utils/io.py

@@ -10,6 +10,7 @@
 import numpy as np
 import pandas as pd
 from pyarrow import feather
+from upath import UPath
 
 import av2.geometry.geometry as geometry_utils
 from av2.geometry.se3 import SE3
@@ -22,7 +23,7 @@
 SensorPosesMapping = Dict[str, SE3]
 
 
-def read_feather(path: Path, columns: Optional[Tuple[str, ...]] = None) -> pd.DataFrame:
+def read_feather(path: Union[Path, UPath], columns: Optional[Tuple[str, ...]] = None) -> pd.DataFrame:
     """Read Apache Feather data from a .feather file.
 
     AV2 uses .feather to serialize much of its data. This function handles the deserialization
@@ -36,7 +37,8 @@ def read_feather(path: Path, columns: Optional[Tuple[str, ...]] = None) -> pd.Da
     Returns:
         (N,len(columns)) Apache Feather data represented as a `pandas` DataFrame.
     """
-    data: pd.DataFrame = feather.read_feather(path, columns=columns)
+    with path.open("rb") as f:
+        data: pd.DataFrame = feather.read_feather(f, columns=columns)
     return data
 
 
@@ -116,7 +118,7 @@ def read_ego_SE3_sensor(log_dir: Path) -> SensorPosesMapping:
     return sensor_name_to_pose
 
 
-def read_city_SE3_ego(log_dir: Path) -> TimestampedCitySE3EgoPoses:
+def read_city_SE3_ego(log_dir: Union[Path, UPath]) -> TimestampedCitySE3EgoPoses:
     """Read the egovehicle poses in the city reference frame.
 
     The egovehicle city pose defines an SE3 transformation from the egovehicle reference frame to the city ref. frame.
@@ -148,7 +150,7 @@ def read_city_SE3_ego(log_dir: Path) -> TimestampedCitySE3EgoPoses:
     Returns:
         Mapping from egovehicle time (in nanoseconds) to egovehicle pose in the city reference frame.
     """
-    city_SE3_ego_path = Path(log_dir, "city_SE3_egovehicle.feather")
+    city_SE3_ego_path = log_dir / "city_SE3_egovehicle.feather"
     city_SE3_ego = read_feather(city_SE3_ego_path)
 
     quat_wxyz = city_SE3_ego.loc[:, ["qw", "qx", "qy", "qz"]].to_numpy()
@@ -206,7 +208,7 @@ def write_img(img_path: Path, img: NDArrayByte, channel_order: str = "RGB") -> N
     cv2.imwrite(str(img_path), img)
 
 
-def read_json_file(fpath: Path) -> Dict[str, Any]:
+def read_json_file(fpath: Union[Path, UPath]) -> Dict[str, Any]:
     """Load dictionary from JSON file.
 
     Args:
@@ -215,7 +217,7 @@ def read_json_file(fpath: Path) -> Dict[str, Any]:
     Returns:
         Deserialized Python dictionary.
     """
-    with open(fpath, "rb") as f:
+    with fpath.open("rb") as f:
         data: Dict[str, Any] = json.load(f)
         return data
 

tests/datasets/motion_forecasting/eval/test_metrics.py

@@ -183,7 +183,7 @@ def test_compute_brier_ade(
     ],
 )
 def test_compute_brier_ade_data_validation(
-    forecast_probabilities: NDArrayFloat, normalize: bool, expectation: AbstractContextManager # type: ignore
+    forecast_probabilities: NDArrayFloat, normalize: bool, expectation: AbstractContextManager  # type: ignore
 ) -> None:
     """Test that test_compute_brier_ade raises the correct errors when inputs are invalid.
 
@@ -247,7 +247,7 @@ def test_compute_brier_fde(
     ],
 )
 def test_compute_brier_fde_data_validation(
-    forecast_probabilities: NDArrayFloat, normalize: bool, expectation: AbstractContextManager # type: ignore
+    forecast_probabilities: NDArrayFloat, normalize: bool, expectation: AbstractContextManager  # type: ignore
 ) -> None:
     """Test that test_compute_brier_fde raises the correct errors when inputs are invalid.
 

tests/datasets/motion_forecasting/eval/test_submission.py

@@ -39,7 +39,7 @@
     ids=["valid", "wrong_shape_trajectory", "mismatched_trajectory_probability_shape"],
 )
 def test_challenge_submission_data_validation(
-    test_submission_dict: Dict[str, ScenarioPredictions], expectation: AbstractContextManager # type: ignore
+    test_submission_dict: Dict[str, ScenarioPredictions], expectation: AbstractContextManager  # type: ignore
 ) -> None:
     """Test that validation of submitted trajectories works as expected during challenge submission initialization.