diff --git a/pyproject.toml b/pyproject.toml index 849f8db..3b8280c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,7 +27,7 @@ dependencies = [ "joblib", "scipy", "shapely", - "py123d==0.6.0", + "py123d>=0.6.0,<0.8", ] [project.optional-dependencies] diff --git a/src/bin_factory/loader/extractor.py b/src/bin_factory/loader/extractor.py index 61879bc..b2f0e2e 100644 --- a/src/bin_factory/loader/extractor.py +++ b/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 diff --git a/src/bin_factory/loader/load.py b/src/bin_factory/loader/load.py index 64e57d7..2dde5a5 100644 --- a/src/bin_factory/loader/load.py +++ b/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) diff --git a/src/bin_factory/loader/mapping.py b/src/bin_factory/loader/mapping.py index 31cb831..dd22c88 100644 --- a/src/bin_factory/loader/mapping.py +++ b/src/bin_factory/loader/mapping.py @@ -22,6 +22,7 @@ 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 = { @@ -29,7 +30,6 @@ 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 = { diff --git a/src/bin_factory/main.py b/src/bin_factory/main.py index 81e1547..ad45ee1 100644 --- a/src/bin_factory/main.py +++ b/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") diff --git a/src/bin_factory/presets.toml b/src/bin_factory/presets.toml index c449a0c..be6a240 100644 --- a/src/bin_factory/presets.toml +++ b/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 diff --git a/src/bin_factory/schema.py b/src/bin_factory/schema.py index 2e14623..5e470e3 100644 --- a/src/bin_factory/schema.py +++ b/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 diff --git a/src/bin_factory/transforms/traffic_controls.py b/src/bin_factory/transforms/traffic_controls.py index 408ffff..f10f8d7 100644 --- a/src/bin_factory/transforms/traffic_controls.py +++ b/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