Add EgoSpeak prediction alignment for floor evaluation

2026-09-10 11:21:09 +08:00
parent 4b7d2c0c20
commit 1e0aea0956
3 changed files with 206 additions and 0 deletions
@@ -0,0 +1,140 @@
+"""Convert saved EgoSpeak MiniROADA forecasts for floor-event evaluation.
+
+EgoSpeak MiniROADA saves one array per session as
+``[source_time, future_horizon, class]``.  The repository floor evaluator uses
+the LSTR-compatible target-aligned layout ``[target_time, class, horizon]``
+plus an explicit causal-valid mask.  This script performs only that alignment;
+it does not change scores, targets, or class order.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+
+
+NUM_CLASSES = 3
+NUM_HORIZONS = 10
+
+
+def normalize_session_name(value: str | int) -> str:
+    text = str(value)
+    suffix = text.removeprefix("Session_")
+    if not suffix.isdigit():
+        raise ValueError(f"Invalid EasyCom session: {value}")
+    return f"Session_{int(suffix)}"
+
+
+def align_miniroad_prediction(
+    raw_prediction: np.ndarray,
+    target_steps: int,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Align MiniROADA source-indexed forecasts without padding them as valid."""
+
+    raw_prediction = np.asarray(raw_prediction)
+    expected = (target_steps - NUM_HORIZONS, NUM_HORIZONS, NUM_CLASSES)
+    if raw_prediction.shape != expected:
+        raise ValueError(
+            "Expected MiniROADA prediction "
+            f"{expected} for target length {target_steps}, got "
+            f"{raw_prediction.shape}"
+        )
+    if not np.isfinite(raw_prediction).all():
+        raise ValueError("MiniROADA prediction contains NaN/Inf")
+    if (raw_prediction < -1e-6).any() or (raw_prediction > 1.0 + 1e-6).any():
+        raise ValueError("MiniROADA prediction is outside [0,1]")
+
+    raw_prediction = np.asarray(raw_prediction, dtype=np.float32)
+    aligned = np.zeros(
+        (target_steps, NUM_CLASSES, NUM_HORIZONS), dtype=np.float32
+    )
+    valid = np.zeros((target_steps, NUM_HORIZONS), dtype=bool)
+    source_rows = np.arange(len(raw_prediction), dtype=np.int64)
+    for horizon_index in range(NUM_HORIZONS):
+        target_rows = source_rows + horizon_index + 1
+        aligned[target_rows, :, horizon_index] = raw_prediction[
+            :, horizon_index, :
+        ]
+        valid[target_rows, horizon_index] = True
+    return aligned, valid
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--prediction-root",
+        type=Path,
+        required=True,
+        help="Directory containing raw MiniROADA Session_X.npy predictions.",
+    )
+    parser.add_argument("--target-root", type=Path, required=True)
+    parser.add_argument("--output-dir", type=Path, required=True)
+    parser.add_argument("--sessions", nargs="+", required=True)
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    sessions = [normalize_session_name(value) for value in args.sessions]
+    if len(sessions) != len(set(sessions)):
+        raise ValueError("Duplicate sessions")
+
+    converted: dict[str, tuple[np.ndarray, np.ndarray]] = {}
+    sources: dict[str, str] = {}
+    for session in sessions:
+        target_path = args.target_root / f"{session}.npy"
+        prediction_path = args.prediction_root / f"{session}.npy"
+        if not target_path.is_file():
+            raise FileNotFoundError(target_path)
+        if not prediction_path.is_file():
+            raise FileNotFoundError(prediction_path)
+        target = np.load(target_path, allow_pickle=False)
+        if target.ndim != 2 or target.shape[1] != NUM_CLASSES:
+            raise ValueError(
+                f"Expected target [T,{NUM_CLASSES}], got {target.shape}: "
+                f"{target_path}"
+            )
+        raw_prediction = np.load(prediction_path, allow_pickle=False)
+        converted[session] = align_miniroad_prediction(
+            raw_prediction, len(target)
+        )
+        sources[session] = str(prediction_path.expanduser().resolve())
+
+    prediction_dir = args.output_dir / "pred"
+    valid_dir = args.output_dir / "valid"
+    destinations = [
+        directory / f"{session}.npy"
+        for directory in (prediction_dir, valid_dir)
+        for session in sessions
+    ]
+    destinations.append(args.output_dir / "conversion.json")
+    existing = [path for path in destinations if path.exists()]
+    if existing:
+        raise FileExistsError(
+            "Refusing to overwrite converted output: "
+            + ", ".join(str(path) for path in existing)
+        )
+
+    prediction_dir.mkdir(parents=True, exist_ok=True)
+    valid_dir.mkdir(parents=True, exist_ok=True)
+    for session, (prediction, valid) in converted.items():
+        np.save(prediction_dir / f"{session}.npy", prediction)
+        np.save(valid_dir / f"{session}.npy", valid)
+    metadata = {
+        "protocol": "egospeak-miniroada-source-to-target-alignment-v1",
+        "source_layout": "[source_time, future_horizon, class]",
+        "output_layout": "[target_time, class, future_horizon]",
+        "sessions": sessions,
+        "sources": sources,
+    }
+    (args.output_dir / "conversion.json").write_text(
+        json.dumps(metadata, indent=2) + "\n"
+    )
+    print(f"Converted {len(sessions)} session(s): {args.output_dir}")
+
+
+if __name__ == "__main__":
+    main()