diff --git a/src/av2/datasets/sensor/av2_sensor_dataloader.py b/src/av2/datasets/sensor/av2_sensor_dataloader.py index bb1d4d4b..4bf52e16 100644 --- a/src/av2/datasets/sensor/av2_sensor_dataloader.py +++ b/src/av2/datasets/sensor/av2_sensor_dataloader.py @@ -46,10 +46,6 @@ def __init__(self, data_dir: Path, labels_dir: Path) -> None: self._labels_dir = labels_dir self._sdb = SynchronizationDB(str(data_dir)) - def get_city_name(self, log_id: str) -> str: - """Return the name of the city where the log of interest was captured.""" - raise NotImplementedError("Method is not available yet.") - def get_log_pinhole_camera(self, log_id: str, cam_name: str) -> PinholeCamera: """Return a PinholeCamera parameterized by sensor pose in vehicle frame, intrinsics, and image dimensions.""" log_dir = self._data_dir / log_id @@ -136,6 +132,28 @@ def get_log_map_dirpath(self, log_id: str) -> Path: """Fetch the path to the directory containing map files for a single vehicle log.""" return self._data_dir / log_id / "map" + def get_city_name(self, log_id: str) -> str: + """Return the name of the city where the log of interest was captured. + + Vector map filenames contain the city name, and have a name in the following format: + `log_map_archive_453e5558-6363-38e3-bf9b-42b5ba0a6f1d____PAO_city_71741.json` + + Args: + log_id: unique ID of vehicle log. + + Returns: + Name of the city where the log of interest was captured. + + Raises: + RuntimeError: If no vector map file is found for the query log ID. + """ + vector_map_fpaths = list(self._data_dir.glob(f"{log_id}/map/log_map_archive*")) + if len(vector_map_fpaths) == 0: + raise RuntimeError(f"Vector map file is missing for {log_id}.") + vector_map_fpath = vector_map_fpaths[0] + log_city_name = vector_map_fpath.name.split("____")[1].split("_")[0] + return log_city_name + def get_log_ids(self) -> List[str]: """Return a list of all vehicle log IDs available at the provided dataroot.""" return sorted([d.name for d in self._data_dir.glob("*") if d.is_dir()])