234 lines
11 KiB
Python
234 lines
11 KiB
Python
#!/usr/bin/env python3
|
|
"""Prepare one static LiDAR frame and one RTK reference 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
|
|
|
|
from rtk_attitude import (
|
|
heading_to_enu_yaw,
|
|
parse_pitch_roll_from_heading_raw,
|
|
rotation_to_quat_xyzw,
|
|
rtk_body_rotation,
|
|
)
|
|
|
|
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 scalar(data: np.lib.npyio.NpzFile, name: str, default: float | None = None) -> float:
|
|
if name not in data.files:
|
|
if default is None:
|
|
raise KeyError(name)
|
|
return float(default)
|
|
return float(np.asarray(data[name]).reshape(-1)[0])
|
|
|
|
|
|
def frame_pitch_roll(data: np.lib.npyio.NpzFile) -> tuple[float, float]:
|
|
pitch = scalar(data, "rtk_pitch_deg", math.nan)
|
|
roll = scalar(data, "rtk_roll_deg", math.nan)
|
|
if math.isfinite(pitch) and math.isfinite(roll):
|
|
return pitch, roll
|
|
raw = None
|
|
if "rtk_heading_raw_utf8" in data.files:
|
|
raw = bytes(np.asarray(data["rtk_heading_raw_utf8"]).reshape(-1))
|
|
parsed_pitch, parsed_roll = parse_pitch_roll_from_heading_raw(raw)
|
|
if not math.isfinite(pitch):
|
|
pitch = float(parsed_pitch) if parsed_pitch is not None else 0.0
|
|
if not math.isfinite(roll):
|
|
roll = float(parsed_roll) if parsed_roll is not None else 0.0
|
|
return pitch, roll
|
|
|
|
|
|
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="rtk_gga_raw_heading")
|
|
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(
|
|
"--orientation-model",
|
|
choices=("heading_pitch_roll", "yaw_only"),
|
|
default="heading_pitch_roll",
|
|
help="heading_pitch_roll uses GNHPR/UNIHEADINGA pitch+roll; yaw_only forces roll=pitch=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:
|
|
pitch, roll = frame_pitch_roll(data)
|
|
samples.append((scalar(data, "rtk_lat_deg"), scalar(data, "rtk_lon_deg"),
|
|
scalar(data, "rtk_altitude_m"), scalar(data, "rtk_raw_heading_deg"),
|
|
pitch, roll, scalar(data, "rtk_heading_stddev_deg", math.nan)))
|
|
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"])
|
|
reported_std = values[:, 6]
|
|
reported_std_mean = float(np.nanmean(reported_std)) if np.isfinite(reported_std).any() else None
|
|
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]),
|
|
"pitch": float(np.mean(values[:, 4])), "roll": float(np.mean(values[:, 5])),
|
|
})
|
|
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])),
|
|
"rtk_roll_mean_deg": float(np.mean(values[:, 5])),
|
|
"reported_heading_std_mean_deg": reported_std_mean,
|
|
"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)
|
|
use_attitude = args.orientation_model == "heading_pitch_roll"
|
|
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, yaw = heading_to_enu_yaw(item["heading"], args.heading_offset_deg)
|
|
pitch = float(item["pitch"]) if use_attitude else 0.0
|
|
roll = float(item["roll"]) if use_attitude else 0.0
|
|
rotation = rtk_body_rotation(
|
|
item["heading"], args.heading_offset_deg, pitch_deg=pitch, roll_deg=roll
|
|
)
|
|
reference_position = antenna - rotation @ lever
|
|
quat = rotation_to_quat_xyzw(rotation)
|
|
pose_rows.append(dict(zip(POSE_FIELDS, [item["time"], *reference_position, *quat])))
|
|
summaries[index - 1].update({
|
|
"sequence": index, "prepared_frame": destination.name,
|
|
"corrected_heading_deg": corrected_heading,
|
|
"pose_yaw_enu_deg": math.degrees(yaw),
|
|
"pose_pitch_deg": pitch, "pose_roll_deg": roll,
|
|
})
|
|
pose_path = args.output / f"reference_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",
|
|
"reference_pose_configuration": {
|
|
"raw_heading_offset_deg": args.heading_offset_deg,
|
|
"antenna_lever_body_m": args.antenna_lever,
|
|
"heading_offset_semantics": (
|
|
"added to clockwise-from-north GNHPR heading before ENU yaw conversion"
|
|
),
|
|
"orientation_model": args.orientation_model,
|
|
"orientation_composition": (
|
|
"R_W_body = Rz(yaw_raw) Ry(-pitch) Rx(roll) Rz(-heading_offset); "
|
|
"yaw_raw from rawHeading, pitch/roll stay in baseline frame"
|
|
),
|
|
"pitch_roll_note": (
|
|
"pitch/roll come from dual-antenna GNHPR/UNIHEADINGA (baseline elevation / reported roll). "
|
|
"This is not a fused IMU vehicle attitude; G90 roll is often ~0."
|
|
),
|
|
},
|
|
"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,
|
|
"orientation_model": args.orientation_model}, ensure_ascii=False, indent=2))
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|