Spaces:
Paused
Paused
| """Dataset fusion engine: temporal sync, stratification, manifest output.""" | |
| from __future__ import annotations | |
| import json | |
| from collections import Counter | |
| from pathlib import Path | |
| from typing import Any, Dict, List, Optional | |
| from sensorflow.adapters.a2d2_adapter import A2D2Adapter | |
| from sensorflow.adapters.alpamayo_adapter import AlpamayoAdapter, DEFAULT_ALPAMAYO_SAMPLES | |
| from sensorflow.adapters.local_adapter import DEFAULT_MAX_FRAMES, LocalSequenceAdapter | |
| from sensorflow.adapters.waymo_adapter import WaymoAdapter | |
| from sensorflow.schemas.taxonomy_axes import assign_taxonomy_axes | |
| from sensorflow.schemas.unified_frame import FusedFrame, UnifiedSequence | |
| class DatasetFusionEngine: | |
| """Ingest, fuse, and stratify multi-vendor sensor sequences.""" | |
| def __init__(self, sync_tolerance_us: int = 50_000): | |
| self.sync_tolerance_us = sync_tolerance_us | |
| self.alpamayo = AlpamayoAdapter() | |
| self.waymo = WaymoAdapter() | |
| self.a2d2 = A2D2Adapter() | |
| self.local = LocalSequenceAdapter() | |
| def ingest( | |
| self, | |
| vendors: List[str], | |
| sequence_id: str, | |
| alpamayo_source: Optional[Dict[str, Any]] = None, | |
| waymo_source: Optional[Dict[str, Any]] = None, | |
| a2d2_source: Optional[Dict[str, Any]] = None, | |
| source_path: Optional[str] = None, | |
| max_frames: Optional[int] = None, | |
| ) -> UnifiedSequence: | |
| sequences: List[UnifiedSequence] = [] | |
| vendors_norm = [v.lower().strip() for v in vendors] | |
| if "local" in vendors_norm: | |
| local_src: Dict[str, Any] = { | |
| "source_path": source_path or "data", | |
| "max_frames": DEFAULT_MAX_FRAMES if max_frames is None else max_frames, | |
| } | |
| sequences.append(self.local.load(local_src, f"{sequence_id}_local")) | |
| if "alpamayo" in vendors_norm: | |
| src = alpamayo_source or DEFAULT_ALPAMAYO_SAMPLES.get("physical_ai", {}) | |
| sequences.append(self.alpamayo.load(src, f"{sequence_id}_alp")) | |
| if "waymo" in vendors_norm: | |
| sequences.append(self.waymo.load(waymo_source or {}, f"{sequence_id}_way")) | |
| if "a2d2" in vendors_norm: | |
| sequences.append(self.a2d2.load(a2d2_source or {}, f"{sequence_id}_a2d2")) | |
| if not sequences: | |
| sequences.append(self.alpamayo.load(DEFAULT_ALPAMAYO_SAMPLES["physical_ai"], sequence_id)) | |
| fused = self._merge_sequences(sequences, sequence_id) | |
| fused = self._stratify(fused) | |
| return fused | |
| def _merge_sequences(self, sequences: List[UnifiedSequence], sequence_id: str) -> UnifiedSequence: | |
| if len(sequences) == 1: | |
| seq = sequences[0] | |
| seq.sequence_id = sequence_id | |
| return seq | |
| all_frames: List[FusedFrame] = [] | |
| for seq in sequences: | |
| all_frames.extend(seq.frames) | |
| all_frames.sort(key=lambda f: f.timestamp_us) | |
| vendors = [s.vendor for s in sequences] | |
| vendor = "mixed" if len(set(vendors)) > 1 else vendors[0] | |
| demo_stub = all(bool(s.taxonomy_manifest.get("demo_stub")) for s in sequences) | |
| return UnifiedSequence( | |
| sequence_id=sequence_id, | |
| vendor=vendor, | |
| frames=all_frames, | |
| calibration=sequences[0].calibration, | |
| taxonomy_manifest={ | |
| "merged_vendors": vendors, | |
| "demo_stub": demo_stub, | |
| "total_frames": len(all_frames), | |
| }, | |
| ) | |
| def _stratify(self, sequence: UnifiedSequence) -> UnifiedSequence: | |
| """Tag frames/objects with six-axis taxonomy and build stratification manifest.""" | |
| strata: Counter = Counter() | |
| for frame in sequence.frames: | |
| speed = frame.ego_pose.speed_kmh | |
| for gt in frame.ground_truth: | |
| gt.taxonomy_axes = assign_taxonomy_axes(gt.class_name, speed_kmh=speed) | |
| key = f"{gt.taxonomy_axes.mode}|{gt.taxonomy_axes.actor}|{gt.taxonomy_axes.infra}" | |
| strata[key] += 1 | |
| sequence.taxonomy_manifest["stratification"] = dict(strata) | |
| sequence.taxonomy_manifest["total_frames"] = len(sequence.frames) | |
| sequence.taxonomy_manifest["total_objects"] = sum(len(f.ground_truth) for f in sequence.frames) | |
| if "demo_stub" not in sequence.taxonomy_manifest: | |
| sequence.taxonomy_manifest["demo_stub"] = sequence.vendor in ("alpamayo", "waymo", "a2d2") | |
| return sequence | |
| def save_manifest(self, sequence: UnifiedSequence, output_dir: Optional[Path] = None) -> Path: | |
| output_dir = output_dir or Path("runs/pipeline") / sequence.sequence_id | |
| manifest_path = output_dir / "manifest.json" | |
| sequence.save(manifest_path) | |
| state_path = Path("runs/pipeline/state.json") | |
| state = self._load_state(state_path) | |
| state[sequence.sequence_id] = { | |
| "ingest_complete": True, | |
| "frames": len(sequence.frames), | |
| "vendor": sequence.vendor, | |
| "demo_stub": bool(sequence.taxonomy_manifest.get("demo_stub")), | |
| "manifest": str(manifest_path), | |
| } | |
| state_path.parent.mkdir(parents=True, exist_ok=True) | |
| with open(state_path, "w") as f: | |
| json.dump(state, f, indent=2) | |
| return manifest_path | |
| def _load_state(path: Path) -> Dict[str, Any]: | |
| if path.exists(): | |
| with open(path) as f: | |
| return json.load(f) | |
| return {} | |
| def get_status(self, sequence_id: str) -> Dict[str, Any]: | |
| state = self._load_state(Path("runs/pipeline/state.json")) | |
| return state.get(sequence_id, {"ingest_complete": False}) | |