"""Normalize Audi A2D2 samples into UnifiedSequence.""" from __future__ import annotations import json from pathlib import Path from typing import Any, Dict, List import numpy as np from sensorflow.adapters.base import VendorAdapter from sensorflow.adapters.stub_images import stub_camera_path, write_stub_camera_png from sensorflow.schemas.taxonomy_axes import assign_taxonomy_axes from sensorflow.schemas.unified_frame import ( CameraView, EgoPose, FusedFrame, GroundTruthObject, LidarData, UnifiedSequence, ) # Audi A2D2 class names used in the public dataset catalog. A2D2_CLASSES = ( "car", "pedestrian", "truck", "bicycle", "traffic_sign", "signal", ) def _default_a2d2_frames() -> List[Dict[str, Any]]: """Built-in shard when A2D2 HDF5/PNG sources are unavailable.""" return [ { "frame_id": "a2d2_0000", "timestamp_us": 0, "lidar_points": 2200, "ego_speed_kmh": 52.0, "image_path": "https://images.unsplash.com/photo-1449965408869-eaa3f722e40d?w=800", "labels": [ {"instance_id": "a2d2_1", "class_name": "car", "center_x": 12.0, "center_y": 1.5, "center_z": 0.6, "length": 4.4, "width": 1.8, "height": 1.5, "heading": 0.05}, {"instance_id": "a2d2_2", "class_name": "pedestrian", "center_x": 8.0, "center_y": -2.0, "center_z": 0.9, "length": 0.6, "width": 0.5, "height": 1.7, "heading": -0.3}, {"instance_id": "a2d2_3", "class_name": "traffic_sign", "center_x": 18.0, "center_y": 3.0, "center_z": 1.8, "length": 0.3, "width": 0.8, "height": 1.2, "heading": 0.0}, ], }, { "frame_id": "a2d2_0001", "timestamp_us": 100_000, "lidar_points": 2150, "ego_speed_kmh": 51.2, "image_path": "https://images.unsplash.com/photo-1492144534655-ae79c964c9d7?w=800", "labels": [ {"instance_id": "a2d2_1", "class_name": "car", "center_x": 11.5, "center_y": 1.4, "center_z": 0.6, "length": 4.4, "width": 1.8, "height": 1.5, "heading": 0.04}, {"instance_id": "a2d2_2", "class_name": "pedestrian", "center_x": 7.6, "center_y": -1.8, "center_z": 0.9, "length": 0.6, "width": 0.5, "height": 1.7, "heading": -0.2}, {"instance_id": "a2d2_4", "class_name": "bicycle", "center_x": 14.0, "center_y": -3.0, "center_z": 0.7, "length": 1.6, "width": 0.5, "height": 1.2, "heading": 1.1}, ], }, { "frame_id": "a2d2_0002", "timestamp_us": 200_000, "lidar_points": 2300, "ego_speed_kmh": 49.8, "image_path": "https://images.unsplash.com/photo-1503376780353-7e6692767b70?w=800", "labels": [ {"instance_id": "a2d2_1", "class_name": "car", "center_x": 11.0, "center_y": 1.3, "center_z": 0.6, "length": 4.4, "width": 1.8, "height": 1.5, "heading": 0.02}, {"instance_id": "a2d2_5", "class_name": "truck", "center_x": 22.0, "center_y": 0.5, "center_z": 1.2, "length": 8.5, "width": 2.5, "height": 3.2, "heading": 0.0}, {"instance_id": "a2d2_6", "class_name": "signal", "center_x": 16.0, "center_y": 2.5, "center_z": 3.0, "length": 0.4, "width": 0.4, "height": 1.0, "heading": 0.0}, ], }, ] def _write_synthetic_lidar(path: Path, num_points: int, labels: List[Dict]) -> None: path.parent.mkdir(parents=True, exist_ok=True) points = [] for lbl in labels: cx, cy, cz = lbl["center_x"], lbl["center_y"], lbl["center_z"] for _ in range(max(20, num_points // max(len(labels), 1) // 10)): points.append([ cx + np.random.randn() * 0.2, cy + np.random.randn() * 0.2, cz + np.random.randn() * 0.1, ]) if not points: points = np.random.randn(num_points, 3).astype(np.float32) else: points = np.array(points, dtype=np.float32) points.tofile(path) DEFAULT_A2D2_SAMPLES = { "ingolstadt": { "dataset": "audi/a2d2", "frames": _default_a2d2_frames(), }, } class A2D2Adapter(VendorAdapter): """Audi A2D2 adapter — stub frames when HDF5/PNG lake is unavailable.""" def load(self, source: Dict[str, Any], sequence_id: str) -> UnifiedSequence: from sensorflow.adapters.vendor_media import ( frames_from_media_root, media_available, resolve_vendor_root, ) source = source or {} root = resolve_vendor_root(source, "source_path", "root", "path") if root is not None and media_available(root): frames = frames_from_media_root( sequence_id=sequence_id, vendor="a2d2", root=root, max_frames=source.get("max_frames"), ) return UnifiedSequence( sequence_id=sequence_id, vendor="a2d2", frames=frames, taxonomy_manifest={ "source_dataset": source.get("dataset", "audi/a2d2"), "source_path": str(root.resolve()), "demo_stub": False, "total_frames": len(frames), }, ) shard_dir = Path("runs/pipeline/a2d2_shards") shard_path = shard_dir / f"{source.get('shard_id', 'default')}.json" if shard_path.exists(): with open(shard_path) as f: frames_data = json.load(f) elif source.get("frames"): frames_data = source["frames"] else: frames_data = DEFAULT_A2D2_SAMPLES["ingolstadt"]["frames"] shard_dir.mkdir(parents=True, exist_ok=True) with open(shard_path, "w") as f: json.dump(frames_data, f, indent=2) frames: List[FusedFrame] = [] lidar_dir = Path("runs/pipeline") / sequence_id / "lidar" for shard in frames_data: frame_id = shard["frame_id"] lidar_path = lidar_dir / f"{frame_id}_lidar.bin" labels = shard.get("labels", []) num_points = shard.get("lidar_points", 2048) _write_synthetic_lidar(lidar_path, num_points, labels) gt = [] speed = shard.get("ego_speed_kmh", 50.0) for lbl in labels: class_name = lbl.get("class_name", "car") if class_name not in A2D2_CLASSES: class_name = "car" axes = assign_taxonomy_axes(class_name, speed_kmh=speed) gt.append(GroundTruthObject( instance_id=lbl["instance_id"], class_name=class_name, confidence=1.0, bbox_3d=[ lbl["center_x"], lbl["center_y"], lbl["center_z"], lbl["length"], lbl["width"], lbl["height"], lbl["heading"], ], taxonomy_axes=axes, )) image_path = shard.get("image_path") or "" if not image_path or str(image_path).startswith("http"): cam_file = stub_camera_path(sequence_id, frame_id, "front") write_stub_camera_png( cam_file, seed=sum(ord(c) for c in frame_id) % 991, label=f"a2d2:{frame_id}", ) image_path = str(cam_file) frames.append(FusedFrame( frame_id=frame_id, timestamp_us=shard.get("timestamp_us", 0), lidar=LidarData(path=str(lidar_path), num_points=num_points), cameras={"front": CameraView(image_path=image_path)}, ego_pose=EgoPose(speed_kmh=speed), ground_truth=gt, )) demo_stub = bool(source.get("demo_stub", True)) if not source.get("frames") else bool( source.get("demo_stub", False) ) if shard_path.exists() and not source.get("frames") and len(frames_data) <= 3: demo_stub = True return UnifiedSequence( sequence_id=sequence_id, vendor="a2d2", frames=frames, taxonomy_manifest={ "source_dataset": source.get("dataset", "audi/a2d2"), "shard": str(shard_path), "demo_stub": demo_stub, "stub_note": ( "Built-in Audi A2D2 sample (~3 frames), not a full AV video lake. " "Set a2d2_root to a real A2D2 PNG/HDF5 export folder when available." ) if demo_stub else None, "total_frames": len(frames), }, ) @classmethod def load_from_file(cls, path: Path, sequence_id: str) -> UnifiedSequence: with open(path) as f: return cls().load(json.load(f), sequence_id)