#!/usr/bin/env python3 """Prepare one static LiDAR frame and one yaw-only RTK body pose per NPZ segment.""" from __future__ import annotations import argparse import csv import json import math import re import shutil from pathlib import Path from typing import Any import numpy as np POSE_FIELDS = ["time", "x", "y", "z", "qx", "qy", "qz", "qw"] def natural_key(value: str) -> list[Any]: return [int(part) if part.isdigit() else part.lower() for part in re.split(r"(\d+)", value)] def truth(value: Any) -> bool: return str(value).strip().lower() in {"1", "true", "yes", "y"} def circular_mean_deg(values: np.ndarray) -> float: radians = np.deg2rad(values) return float(np.rad2deg(math.atan2(np.mean(np.sin(radians)), np.mean(np.cos(radians)))) % 360.0) def circular_std_deg(values: np.ndarray) -> float: radians = np.deg2rad(values) resultant = max(math.hypot(np.mean(np.cos(radians)), np.mean(np.sin(radians))), 1e-12) return float(np.rad2deg(math.sqrt(-2.0 * math.log(resultant)))) def geodetic_to_ecef(lat_deg: float, lon_deg: float, height_m: float) -> np.ndarray: a, e2 = 6378137.0, 6.69437999014e-3 lat, lon = math.radians(lat_deg), math.radians(lon_deg) sin_lat, cos_lat, sin_lon, cos_lon = math.sin(lat), math.cos(lat), math.sin(lon), math.cos(lon) n = a / math.sqrt(1.0 - e2 * sin_lat * sin_lat) return np.array([(n + height_m) * cos_lat * cos_lon, (n + height_m) * cos_lat * sin_lon, (n * (1.0 - e2) + height_m) * sin_lat], dtype=float) def ecef_to_enu(ecef: np.ndarray, origin: np.ndarray, lat_deg: float, lon_deg: float) -> np.ndarray: lat, lon = math.radians(lat_deg), math.radians(lon_deg) slat, clat, slon, clon = math.sin(lat), math.cos(lat), math.sin(lon), math.cos(lon) rotation = np.array([[-slon, clon, 0.0], [-slat * clon, -slat * slon, clat], [clat * clon, clat * slon, slat]], dtype=float) return rotation @ (ecef - origin) def yaw_rotation(yaw: float) -> np.ndarray: c, s = math.cos(yaw), math.sin(yaw) return np.array([[c, -s, 0.0], [s, c, 0.0], [0.0, 0.0, 1.0]]) def scalar(data: np.lib.npyio.NpzFile, name: str) -> float: return float(np.asarray(data[name]).reshape(-1)[0]) def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--combined-root", type=Path, required=True) parser.add_argument("--output", type=Path, required=True) parser.add_argument("--pose-name", default="rear_gga_raw_rear_to_front") parser.add_argument("--heading-offset-deg", type=float, required=True) parser.add_argument("--antenna-lever", type=float, nargs=3, required=True, metavar=("X", "Y", "Z")) parser.add_argument("--accepted-fixes", type=int, nargs="+", default=[4, 5]) parser.add_argument("--heading-std-limit-deg", type=float, default=0.5) parser.add_argument("--min-stations", type=int, default=30) parser.add_argument("--expected-stations", type=int, default=0) parser.add_argument("--overwrite", action="store_true") return parser.parse_args() def main() -> int: args = parse_args() manifest_path = args.combined_root / "manifest.csv" with manifest_path.open("r", encoding="utf-8-sig", newline="") as stream: rows = list(csv.DictReader(stream)) required = {"segment", "output", "lidar_time_ns", "rtk_valid", "heading_valid", "rtk_fix_quality"} if not rows or not required.issubset(rows[0]): raise ValueError(f"{manifest_path} is empty or lacks {sorted(required)}") groups: dict[str, list[dict[str, str]]] = {} for row in rows: groups.setdefault(row["segment"], []).append(row) selected, summaries, rejected = [], [], [] accepted_fixes = set(args.accepted_fixes) for segment in sorted(groups, key=natural_key): group = sorted(groups[segment], key=lambda row: int(row["lidar_time_ns"])) good = [row for row in group if truth(row["rtk_valid"]) and truth(row["heading_valid"]) and int(row["rtk_fix_quality"]) in accepted_fixes] if not good: rejected.append({"station": segment, "reason": "no associated fixed RTK position and valid heading"}) continue samples = [] for row in good: path = args.combined_root / Path(row["output"]) with np.load(path, allow_pickle=False) as data: samples.append((scalar(data, "rtk_lat_deg"), scalar(data, "rtk_lon_deg"), scalar(data, "rtk_altitude_m"), scalar(data, "rtk_raw_heading_deg"), scalar(data, "rtk_pitch_deg"), scalar(data, "rtk_heading_stddev_deg"))) values = np.asarray(samples, dtype=float) heading_std = circular_std_deg(values[:, 3]) if heading_std > args.heading_std_limit_deg: rejected.append({"station": segment, "reason": f"heading std {heading_std:.4f} deg exceeds limit"}) continue frame = good[len(good) // 2] source = args.combined_root / Path(frame["output"]) selected.append({"station": segment, "source": source, "time": int(frame["lidar_time_ns"]) / 1e9, "lat": float(np.mean(values[:, 0])), "lon": float(np.mean(values[:, 1])), "alt": float(np.mean(values[:, 2])), "heading": circular_mean_deg(values[:, 3])}) summaries.append({"station": segment, "frames": len(group), "valid_fixed_frames": len(good), "heading_mean_deg": circular_mean_deg(values[:, 3]), "heading_circular_std_deg": heading_std, "rtk_pitch_mean_deg": float(np.mean(values[:, 4])), "reported_heading_std_mean_deg": float(np.nanmean(values[:, 5])), "altitude_std_m": float(np.std(values[:, 2])), "selected_source": str(source)}) if args.expected_stations and len(selected) != args.expected_stations: raise RuntimeError(f"expected {args.expected_stations} usable stations, got {len(selected)}; rejected={rejected}") if len(selected) < args.min_stations: raise RuntimeError(f"need at least {args.min_stations} usable stations, got {len(selected)}; rejected={rejected}") if args.output.exists() and any(args.output.iterdir()) and not args.overwrite: raise FileExistsError(f"{args.output} is non-empty; pass --overwrite") frames = args.output / "frames_all" frames.mkdir(parents=True, exist_ok=True) origin = selected[0] origin_ecef = geodetic_to_ecef(origin["lat"], origin["lon"], origin["alt"]) lever = np.asarray(args.antenna_lever, dtype=float) pose_rows = [] for index, item in enumerate(selected, 1): destination = frames / f"station_{index:02d}.npz" shutil.copy2(item["source"], destination) antenna = ecef_to_enu(geodetic_to_ecef(item["lat"], item["lon"], item["alt"]), origin_ecef, origin["lat"], origin["lon"]) corrected_heading = (item["heading"] + args.heading_offset_deg) % 360.0 yaw = math.radians(90.0 - corrected_heading) body = antenna - yaw_rotation(yaw) @ lever pose_rows.append(dict(zip(POSE_FIELDS, [item["time"], *body, 0.0, 0.0, math.sin(yaw / 2.0), math.cos(yaw / 2.0)]))) summaries[index - 1].update({"sequence": index, "prepared_frame": destination.name, "corrected_heading_deg": corrected_heading}) pose_path = args.output / f"body_poses_{args.pose_name}.csv" with pose_path.open("w", encoding="utf-8", newline="") as stream: writer = csv.DictWriter(stream, fieldnames=POSE_FIELDS); writer.writeheader(); writer.writerows(pose_rows) with (args.output / "station_summary.csv").open("w", encoding="utf-8", newline="") as stream: fields = sorted({key for row in summaries for key in row}) writer = csv.DictWriter(stream, fieldnames=fields); writer.writeheader(); writer.writerows(summaries) document = {"source_combined_root": str(args.combined_root.resolve()), "station_count": len(selected), "rejected": rejected, "pose_csv": pose_path.name, "selection_policy": "middle LiDAR frame among fixed-position and valid-heading associations", "body_pose_configuration": {"raw_heading_offset_deg": args.heading_offset_deg, "antenna_lever_body_m": args.antenna_lever, "orientation_model": "yaw-only, identical to the previous calibration workflow"}, "stations": [{"sequence": i + 1, "source_station": item["station"], "source_frame": str(item["source"]), "prepared_frame": f"station_{i + 1:02d}.npz"} for i, item in enumerate(selected)]} (args.output / "manifest.json").write_text(json.dumps(document, ensure_ascii=False, indent=2), encoding="utf-8") print(json.dumps({"prepared": str(args.output.resolve()), "stations": len(selected), "rejected": rejected, "pose_csv": pose_path.name}, ensure_ascii=False, indent=2)) return 0 if __name__ == "__main__": raise SystemExit(main())