NuRec support (for running AlpaSim alongside PufferDrive) + generic fixes

pyproject.toml

@@ -27,7 +27,7 @@ dependencies = [
   "joblib",
   "scipy",
   "shapely",
-  "py123d==0.6.0",
+  "py123d>=0.6.0,<0.8",
 ]
 
 [project.optional-dependencies]

src/bin_factory/loader/extractor.py

@@ -75,6 +75,7 @@ def extract_scenario(
         all_objects, labels, tokens_to_object_id = _extract_objects(scene_api, centroid, ego_states)
         for obj in all_objects.values():
             _fill_missing_velocities(obj, dt)
+        dropped_labels: dict[str, int] = {}
         for oid, obj in all_objects.items():
             label = labels[oid]
             if label in mapping.AGENT_TYPE_MAP:
@@ -83,6 +84,16 @@ def extract_scenario(
             elif label in mapping.OBJECT_TYPE_MAP:
                 obj.type = mapping.OBJECT_TYPE_MAP[label]
                 objects[oid] = obj
+            else:
+                dropped_labels[label.name] = dropped_labels.get(label.name, 0) + 1
+        if dropped_labels:
+            log.warning(
+                "Scenario %s: dropped %d track(s) whose labels are in neither AGENT_TYPE_MAP "
+                "nor OBJECT_TYPE_MAP, so they are absent from the bin: %s",
+                metadata.id,
+                sum(dropped_labels.values()),
+                dropped_labels,
+            )
 
         _extract_prediction_targets(scene_api, tokens_to_object_id, agents, metadata)
 
@@ -99,7 +110,7 @@ def extract_scenario(
         objects=objects,
         metadata=metadata,
     )
-    extras = schema.ExtractionExtras(traffic_lights=traffic_lights, stop_zones=stop_zones)
+    extras = schema.ExtractionExtras(traffic_lights=traffic_lights, stop_zones=stop_zones, centroid=centroid.tolist())
     return scenario, extras
 
 

src/bin_factory/loader/load.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import os
 import pathlib
 
 from py123d import api as py123d_api
@@ -21,6 +22,7 @@ def discover_scenes(
     map_only: bool = False,
 ) -> list:
     data_root = pathlib.Path(py123d_data_root)
+    os.environ["PY123D_DATA_ROOT"] = str(data_root)
     if map_only:
         return _discover_maps(data_root / "maps", datasets, num_scenes)
 

src/bin_factory/loader/mapping.py

@@ -22,14 +22,14 @@
     detections.DefaultBoxDetectionLabel.TWO_WHEELER: puffer_types.AgentType.CYCLIST,
     detections.DefaultBoxDetectionLabel.ANIMAL: puffer_types.AgentType.OTHER,
     detections.DefaultBoxDetectionLabel.OTHER: puffer_types.AgentType.OTHER,
+    detections.DefaultBoxDetectionLabel.GENERIC_OBJECT: puffer_types.AgentType.OTHER,
 }
 
 OBJECT_TYPE_MAP = {
     detections.DefaultBoxDetectionLabel.TRAFFIC_SIGN: puffer_types.ObjectType.TRAFFIC_SIGN,
     detections.DefaultBoxDetectionLabel.TRAFFIC_CONE: puffer_types.ObjectType.TRAFFIC_CONE,
     detections.DefaultBoxDetectionLabel.TRAFFIC_LIGHT: puffer_types.ObjectType.TRAFFIC_LIGHT,
     detections.DefaultBoxDetectionLabel.BARRIER: puffer_types.ObjectType.BARRIER,
-    detections.DefaultBoxDetectionLabel.GENERIC_OBJECT: puffer_types.ObjectType.GENERIC_OBJECT,
 }
 
 LANE_TYPE_MAP = {

src/bin_factory/main.py

@@ -1,4 +1,5 @@
 import argparse
+import contextlib
 import json
 import os
 import pathlib
@@ -94,6 +95,13 @@ def build_parser() -> argparse.ArgumentParser:
         action="store_true",
         help="Reverse road-edge polyline order (Waymo convention) for nuplan/carla/opendrive",
     )
+    parser.add_argument(
+        "--emit_metadata",
+        action="store_true",
+        help="Write metadata.jsonl in the output directory, a row per converted scenario "
+        "carrying the recentring centroid, for consumers that map bin coordinates back to "
+        "the source frame",
+    )
     parser.add_argument(
         "--log_level",
         default="INFO",
@@ -172,8 +180,8 @@ def _worker_fn(py123d_data: Any, output_dir: pathlib.Path, config: argparse.Name
     try:
         identity["scenario_id"] = _scenario_identity(py123d_data, config.scenario_id_field)
         bind(dataset=dataset, scenario=identity["scenario_id"])
-        _convert_one(py123d_data, output_dir, config)
-        return {"ok": True, **identity, "error": ""}
+        metadata = _convert_one(py123d_data, output_dir, config)
+        return {"ok": True, **identity, "error": "", "metadata": metadata}
     except loader.ValidationError as ve:
         log.error("validation error: %s", ve)
         return {"ok": False, **identity, "error": str(ve)}
@@ -184,7 +192,7 @@ def _worker_fn(py123d_data: Any, output_dir: pathlib.Path, config: argparse.Name
         unbind(tokens)
 
 
-def _convert_one(py123d_data: Any, output_dir: pathlib.Path, config: argparse.Namespace) -> None:
+def _convert_one(py123d_data: Any, output_dir: pathlib.Path, config: argparse.Namespace) -> dict | None:
     # 1. Load and convert 123D scenario to PufferDrive format
     scenario, extras = loader.extract_scenario(py123d_data, config.scenario_id_field)
 
@@ -218,6 +226,20 @@ def _convert_one(py123d_data: Any, output_dir: pathlib.Path, config: argparse.Na
     temporary_path.write_bytes(binary_data)
     temporary_path.replace(output_path)
 
+    # 6. Per-scenario metadata, returned for the parent to collect: consumers (AlpaSim state
+    # sync, scene manifests) need it to match maps between formats and map bin coords back to
+    # the source frame after centering.
+    if not config.emit_metadata:
+        return None
+    return {
+        "bin_path": output_path.name,
+        "scenario_id": scenario.metadata.id,
+        "dataset": scenario.metadata.dataset,
+        "scenario_length": scenario.metadata.scenario_length,
+        "dt": scenario.metadata.dt,
+        "centroid": extras.centroid,
+    }
+
 
 def _validate_args(args: argparse.Namespace, parser: argparse.ArgumentParser) -> tuple[argparse.Namespace, str]:
     for attr in ("datasets", "split_types", "split_names", "log_names", "scene_uuids"):
@@ -320,6 +342,11 @@ def main() -> int:
 
     with (
         failures_path.open("w", encoding="utf-8") as failure_handle,
+        (
+            (output_dir / "metadata.jsonl").open("w", encoding="utf-8")
+            if args.emit_metadata
+            else contextlib.nullcontext()
+        ) as metadata_handle,
         tqdm.tqdm(total=len(scenes)) as pbar,
         joblib.Parallel(
             n_jobs=args.workers,
@@ -335,6 +362,9 @@ def main() -> int:
             ):
                 if result["ok"]:
                     succeeded += 1
+                    if metadata_handle is not None:
+                        metadata_handle.write(json.dumps(result["metadata"]) + "\n")
+                        metadata_handle.flush()
                 else:
                     failed += 1
                     failure_handle.write(json.dumps(result) + "\n")

src/bin_factory/presets.toml

@@ -33,6 +33,12 @@ datasets = ["opendrive"]
 reverse_road_edges = true
 scenario_id_field = "location"
 
+[nurec]
+datasets = ["nurec"]
+scenario_id_field = "location"
+emit_metadata = true
+# No reverse_road_edges: NuRec boundaries already come with drivable-side-left.
+
 [wod-motion]
 datasets = ["wod-motion"]
 invalid_agent_overlap = true

src/bin_factory/schema.py

@@ -104,6 +104,7 @@ class StopZone:
 class ExtractionExtras:
     traffic_lights: dict = dataclasses.field(default_factory=dict)
     stop_zones: list = dataclasses.field(default_factory=list)
+    centroid: list | None = None
 
 
 @dataclasses.dataclass

src/bin_factory/transforms/traffic_controls.py

@@ -55,9 +55,6 @@ def process_traffic_controls(scenario: schema.PufferScenario, extras: schema.Ext
 
     for element_data in extras.stop_zones:
         stop_zone_type = element_data.type
-        if stop_zone_type == puffer_types.TCType.TRAFFIC_LIGHT and scenario_length > 0:
-            continue  # Skip stop zones if traffic lights are already defined, to avoid duplicates
-
         controlled_lanes = [lid for lid in element_data.controlled_lanes if lid not in covered_lanes]
         if not controlled_lanes:
             continue