313 lines
16 KiB
Python
313 lines
16 KiB
Python
#!/usr/bin/env python3
|
|
"""Build one LiDAR-centric NPZ per frame with matched RTK and an IMU window.
|
|
|
|
Inputs are LiDAR frame NPZ files from frontlidar_dlog_export.py and parsed
|
|
RTK/IMU JSONL files from parse_rtk_imu_v2.py. Raw .rscap files remain the
|
|
traceability source; this script never modifies them.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import csv
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
|
|
|
|
GPS_EPOCH_UNIX_NS = 315964800 * 1_000_000_000
|
|
|
|
|
|
def parse_named_path(text: str) -> tuple[str, Path]:
|
|
if "=" not in text:
|
|
raise argparse.ArgumentTypeError("expected NAME=PATH")
|
|
name, raw_path = text.split("=", 1)
|
|
if not name.strip():
|
|
raise argparse.ArgumentTypeError("segment name is empty")
|
|
return name.strip(), Path(raw_path)
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument(
|
|
"--lidar",
|
|
type=parse_named_path,
|
|
action="append",
|
|
required=True,
|
|
metavar="NAME=FRAMES_DIR",
|
|
help="Repeat for each LiDAR segment; directory contains exported *.npz frames.",
|
|
)
|
|
parser.add_argument("--rtk", type=Path, action="append", required=True, help="Parsed rtk.jsonl; repeat per session.")
|
|
parser.add_argument("--imu", type=Path, action="append", required=True, help="Parsed imu.jsonl; repeat per session.")
|
|
parser.add_argument("--out", type=Path, required=True)
|
|
parser.add_argument("--rtk-max-dt-ms", type=float, default=150.0)
|
|
parser.add_argument("--imu-before-ms", type=float, default=100.0)
|
|
parser.add_argument("--imu-after-ms", type=float, default=100.0)
|
|
parser.add_argument("--gps-utc-leap-seconds", type=int, default=18)
|
|
parser.add_argument("--overwrite", action="store_true")
|
|
return parser.parse_args()
|
|
|
|
|
|
def load_jsonl(paths: list[Path]) -> list[dict[str, Any]]:
|
|
rows: list[dict[str, Any]] = []
|
|
for source_index, path in enumerate(paths):
|
|
source_file = str(path.resolve())
|
|
with path.open("r", encoding="utf-8") as stream:
|
|
for line_number, line in enumerate(stream, start=1):
|
|
if not line.strip():
|
|
continue
|
|
row = json.loads(line)
|
|
row["_source_file"] = source_file
|
|
row["_source_index"] = source_index
|
|
row["_source_line"] = line_number
|
|
rows.append(row)
|
|
return rows
|
|
|
|
|
|
def utf8_array(value: Any) -> np.ndarray:
|
|
return np.frombuffer(str(value if value is not None else "").encode("utf-8"), dtype=np.uint8)
|
|
|
|
|
|
def scalar(array: np.ndarray) -> Any:
|
|
return array.reshape(-1)[0].item()
|
|
|
|
|
|
def nearest_index(times: np.ndarray, target: int) -> int:
|
|
if not len(times):
|
|
return -1
|
|
right = int(np.searchsorted(times, target, side="left"))
|
|
candidates = [index for index in (right - 1, right) if 0 <= index < len(times)]
|
|
return min(candidates, key=lambda index: abs(int(times[index]) - target))
|
|
|
|
|
|
def estimate_imu_times(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
|
"""Recover 100 Hz timing inside each serial chunk from device timestamps.
|
|
|
|
A capture chunk has one host receive timestamp but may contain several IMU
|
|
frames. The last frame is anchored to the chunk receive time and earlier
|
|
frames are moved backwards by their device timestamp difference.
|
|
"""
|
|
groups: dict[tuple[int, int], list[dict[str, Any]]] = {}
|
|
for row in rows:
|
|
if not row.get("crc_valid") or row.get("device_timestamp_ms") is None:
|
|
continue
|
|
key = (int(row["_source_index"]), int(row.get("source_chunk_sequence_last", -1)))
|
|
groups.setdefault(key, []).append(row)
|
|
result: list[dict[str, Any]] = []
|
|
for group in groups.values():
|
|
group.sort(key=lambda row: (int(row["device_timestamp_ms"]), int(row["_source_line"])))
|
|
last_device = int(group[-1]["device_timestamp_ms"])
|
|
host_ns = int(group[-1]["host_receive_utc_ns"])
|
|
for row in group:
|
|
delta_ms = (last_device - int(row["device_timestamp_ms"])) & 0xFFFFFFFF
|
|
if delta_ms > 60_000:
|
|
delta_ms = 0
|
|
copied = dict(row)
|
|
copied["estimated_time_ns"] = host_ns - delta_ms * 1_000_000
|
|
result.append(copied)
|
|
result.sort(key=lambda row: int(row["estimated_time_ns"]))
|
|
return result
|
|
|
|
|
|
def gnss_utc_ns(row: dict[str, Any], leap_seconds: int) -> int | None:
|
|
week, tow_ms = row.get("gnss_week"), row.get("gnss_tow_ms")
|
|
if week is None or tow_ms is None:
|
|
return None
|
|
seconds = int(week) * 604800 + float(tow_ms) / 1000.0 - leap_seconds
|
|
return GPS_EPOCH_UNIX_NS + int(round(seconds * 1_000_000_000))
|
|
|
|
|
|
def numeric_array(rows: list[dict[str, Any]], key: str, dtype: Any, default: Any) -> np.ndarray:
|
|
return np.asarray([row.get(key, default) if row.get(key) is not None else default for row in rows], dtype=dtype)
|
|
|
|
|
|
def raw_frame_matrix(rows: list[dict[str, Any]]) -> tuple[np.ndarray, np.ndarray]:
|
|
frames = [bytes.fromhex(str(row.get("raw_frame_hex", ""))) for row in rows]
|
|
lengths = np.asarray([len(frame) for frame in frames], dtype=np.int32)
|
|
width = max(lengths, default=0)
|
|
matrix = np.zeros((len(frames), width), dtype=np.uint8)
|
|
for index, frame in enumerate(frames):
|
|
matrix[index, : len(frame)] = np.frombuffer(frame, dtype=np.uint8)
|
|
return matrix, lengths
|
|
|
|
|
|
def add_rtk(values: dict[str, np.ndarray], prefix: str, row: dict[str, Any] | None, dt_ns: int | None) -> None:
|
|
values[f"{prefix}_valid"] = np.asarray([row is not None], dtype=np.uint8)
|
|
values[f"{prefix}_dt_ns"] = np.asarray([dt_ns or 0], dtype=np.int64)
|
|
values[f"{prefix}_host_receive_utc_ns"] = np.asarray([0], dtype=np.int64)
|
|
values[f"{prefix}_raw_utf8"] = utf8_array("")
|
|
values[f"{prefix}_source_file_utf8"] = utf8_array("")
|
|
values[f"{prefix}_source_raw_file_offset"] = np.asarray([-1], dtype=np.int64)
|
|
values[f"{prefix}_source_raw_byte_length"] = np.asarray([0], dtype=np.int32)
|
|
if row is None:
|
|
return
|
|
values[f"{prefix}_host_receive_utc_ns"] = np.asarray([row.get("host_receive_utc_ns", 0)], dtype=np.int64)
|
|
values[f"{prefix}_raw_utf8"] = utf8_array(row.get("raw_line", ""))
|
|
values[f"{prefix}_source_file_utf8"] = utf8_array(row.get("_source_file", ""))
|
|
values[f"{prefix}_source_raw_file_offset"] = np.asarray([row.get("source_raw_file_offset", -1)], dtype=np.int64)
|
|
values[f"{prefix}_source_raw_byte_length"] = np.asarray([row.get("source_raw_byte_length", 0)], dtype=np.int32)
|
|
|
|
|
|
def initialize_rtk_measurements(values: dict[str, np.ndarray]) -> None:
|
|
for key, dtype, default in (
|
|
("lat_deg", np.float64, np.nan), ("lon_deg", np.float64, np.nan),
|
|
("altitude_m", np.float64, np.nan), ("hdop", np.float64, np.nan),
|
|
("fix_quality", np.int32, -1), ("gga_satellites", np.int32, -1),
|
|
("differential_age_s", np.float64, np.nan),
|
|
("gnss_week", np.int32, -1), ("gnss_tow_ms", np.int64, -1),
|
|
("baseline_length_m", np.float64, np.nan), ("raw_heading_deg", np.float64, np.nan),
|
|
("pitch_deg", np.float64, np.nan), ("heading_stddev_deg", np.float64, np.nan),
|
|
("pitch_stddev_deg", np.float64, np.nan), ("heading_satellites", np.int32, -1),
|
|
("solution_satellites", np.int32, -1),
|
|
):
|
|
values[f"rtk_{key}"] = np.asarray([default], dtype=dtype)
|
|
values["rtk_fixed"] = np.asarray([0], dtype=np.uint8)
|
|
values["rtk_heading_solution_utf8"] = utf8_array("")
|
|
values["rtk_heading_gnss_utc_ns"] = np.asarray([0], dtype=np.int64)
|
|
values["rtk_heading_host_minus_gnss_ns"] = np.asarray([0], dtype=np.int64)
|
|
|
|
def main() -> int:
|
|
args = parse_args()
|
|
if args.out.exists() and any(args.out.iterdir()) and not args.overwrite:
|
|
raise FileExistsError(f"{args.out} is non-empty; pass --overwrite")
|
|
frames_out = args.out / "frames"
|
|
frames_out.mkdir(parents=True, exist_ok=True)
|
|
|
|
rtk_rows = load_jsonl(args.rtk)
|
|
gga = sorted(
|
|
[row for row in rtk_rows if row.get("type") == "GGA" and row.get("checksum_valid") and row.get("lat_deg") is not None],
|
|
key=lambda row: int(row["host_receive_utc_ns"]),
|
|
)
|
|
heading = sorted(
|
|
[row for row in rtk_rows if row.get("type") == "UNIHEADINGA" and row.get("checksum_valid") and row.get("heading_valid")],
|
|
key=lambda row: int(row["host_receive_utc_ns"]),
|
|
)
|
|
imu = estimate_imu_times(load_jsonl(args.imu))
|
|
gga_times = np.asarray([int(row["host_receive_utc_ns"]) for row in gga], dtype=np.int64)
|
|
heading_times = np.asarray([int(row["host_receive_utc_ns"]) for row in heading], dtype=np.int64)
|
|
imu_times = np.asarray([int(row["estimated_time_ns"]) for row in imu], dtype=np.int64)
|
|
|
|
manifest: list[dict[str, Any]] = []
|
|
global_index = 0
|
|
max_rtk_ns = int(args.rtk_max_dt_ms * 1_000_000)
|
|
before_ns = int(args.imu_before_ms * 1_000_000)
|
|
after_ns = int(args.imu_after_ms * 1_000_000)
|
|
|
|
for segment_name, frame_dir in args.lidar:
|
|
frame_paths = sorted(frame_dir.glob("*.npz"))
|
|
if not frame_paths:
|
|
raise FileNotFoundError(f"no NPZ frames under {frame_dir}")
|
|
for segment_index, source in enumerate(frame_paths):
|
|
with np.load(source, allow_pickle=False) as frame:
|
|
values = {key: np.asarray(frame[key]) for key in frame.files}
|
|
lidar_time_ns = int(scalar(values["unix_time_ns"]))
|
|
|
|
gga_index = nearest_index(gga_times, lidar_time_ns)
|
|
heading_index = nearest_index(heading_times, lidar_time_ns)
|
|
gga_row = gga[gga_index] if gga_index >= 0 else None
|
|
heading_row = heading[heading_index] if heading_index >= 0 else None
|
|
gga_dt = int(gga_times[gga_index]) - lidar_time_ns if gga_index >= 0 else None
|
|
heading_dt = int(heading_times[heading_index]) - lidar_time_ns if heading_index >= 0 else None
|
|
gga_ok = gga_row is not None and abs(gga_dt or 0) <= max_rtk_ns
|
|
heading_ok = heading_row is not None and abs(heading_dt or 0) <= max_rtk_ns
|
|
add_rtk(values, "rtk_gga", gga_row if gga_ok else None, gga_dt)
|
|
add_rtk(values, "rtk_heading", heading_row if heading_ok else None, heading_dt)
|
|
initialize_rtk_measurements(values)
|
|
|
|
if gga_ok and gga_row:
|
|
for key, dtype, default in (
|
|
("lat_deg", np.float64, np.nan), ("lon_deg", np.float64, np.nan),
|
|
("altitude_m", np.float64, np.nan), ("hdop", np.float64, np.nan),
|
|
("fix_quality", np.int32, -1), ("gga_satellites", np.int32, -1),
|
|
("differential_age_s", np.float64, np.nan),
|
|
):
|
|
values[f"rtk_{key}"] = np.asarray([gga_row.get(key, default)], dtype=dtype)
|
|
values["rtk_gga_satellites"] = np.asarray([gga_row.get("satellites", -1)], dtype=np.int32)
|
|
values["rtk_fixed"] = np.asarray([int(gga_row.get("fix_quality", -1)) in {4, 5}], dtype=np.uint8)
|
|
if heading_ok and heading_row:
|
|
for key, dtype, default in (
|
|
("gnss_week", np.int32, -1), ("gnss_tow_ms", np.int64, -1),
|
|
("baseline_length_m", np.float64, np.nan), ("raw_heading_deg", np.float64, np.nan),
|
|
("pitch_deg", np.float64, np.nan), ("heading_stddev_deg", np.float64, np.nan),
|
|
("pitch_stddev_deg", np.float64, np.nan),
|
|
("solution_satellites", np.int32, -1),
|
|
):
|
|
values[f"rtk_{key}"] = np.asarray([heading_row.get(key, default)], dtype=dtype)
|
|
values["rtk_heading_satellites"] = np.asarray([heading_row.get("satellites", -1)], dtype=np.int32)
|
|
values["rtk_heading_solution_utf8"] = utf8_array(heading_row.get("heading_solution", ""))
|
|
device_ns = gnss_utc_ns(heading_row, args.gps_utc_leap_seconds)
|
|
values["rtk_heading_gnss_utc_ns"] = np.asarray([device_ns or 0], dtype=np.int64)
|
|
values["rtk_heading_host_minus_gnss_ns"] = np.asarray(
|
|
[int(heading_row["host_receive_utc_ns"]) - device_ns if device_ns is not None else 0], dtype=np.int64
|
|
)
|
|
|
|
left = int(np.searchsorted(imu_times, lidar_time_ns - before_ns, side="left"))
|
|
right = int(np.searchsorted(imu_times, lidar_time_ns + after_ns, side="right"))
|
|
window = imu[left:right]
|
|
values["imu_window_count"] = np.asarray([len(window)], dtype=np.int32)
|
|
values["imu_valid"] = np.asarray([bool(window)], dtype=np.uint8)
|
|
values["imu_time_ns"] = numeric_array(window, "estimated_time_ns", np.int64, 0)
|
|
values["imu_host_receive_utc_ns"] = numeric_array(window, "host_receive_utc_ns", np.int64, 0)
|
|
for key in ("device_timestamp_ms", "pps_sync_stamp_ms", "tag"):
|
|
values[f"imu_{key}"] = numeric_array(window, key, np.int64, -1)
|
|
for key in (
|
|
"temperature_c", "air_pressure_pa", "accel_x_mps2", "accel_y_mps2", "accel_z_mps2",
|
|
"gyro_x_radps", "gyro_y_radps", "gyro_z_radps", "mag_x_ut", "mag_y_ut", "mag_z_ut",
|
|
"roll_deg", "pitch_deg", "yaw_deg", "quaternion_w", "quaternion_x", "quaternion_y", "quaternion_z",
|
|
):
|
|
values[f"imu_{key}"] = numeric_array(window, key, np.float64, np.nan)
|
|
values["imu_source_index"] = numeric_array(window, "_source_index", np.int32, -1)
|
|
values["imu_source_raw_file_offset"] = numeric_array(window, "source_raw_file_offset", np.int64, -1)
|
|
raw_matrix, raw_lengths = raw_frame_matrix(window)
|
|
values["imu_raw_frame_bytes"] = raw_matrix
|
|
values["imu_raw_frame_length"] = raw_lengths
|
|
values["imu_source_files_json_utf8"] = utf8_array(json.dumps([str(path.resolve()) for path in args.imu], ensure_ascii=False))
|
|
values["source_lidar_file_utf8"] = utf8_array(source.resolve())
|
|
values["segment_name_utf8"] = utf8_array(segment_name)
|
|
|
|
output = frames_out / f"{segment_name}_{segment_index:06d}.npz"
|
|
np.savez_compressed(output, **values)
|
|
manifest.append({
|
|
"global_index": global_index,
|
|
"segment": segment_name,
|
|
"segment_index": segment_index,
|
|
"output": str(output.relative_to(args.out)),
|
|
"source_lidar": str(source.resolve()),
|
|
"lidar_time_ns": lidar_time_ns,
|
|
"rtk_gga_dt_ns": gga_dt,
|
|
"rtk_heading_dt_ns": heading_dt,
|
|
"rtk_valid": gga_ok,
|
|
"heading_valid": heading_ok,
|
|
"rtk_fix_quality": gga_row.get("fix_quality") if gga_ok and gga_row else None,
|
|
"rtk_fixed": bool(gga_ok and gga_row and int(gga_row.get("fix_quality", -1)) in {4, 5}),
|
|
"imu_window_count": len(window),
|
|
})
|
|
global_index += 1
|
|
|
|
fields = sorted({key for row in manifest for key in row})
|
|
with (args.out / "manifest.csv").open("w", encoding="utf-8", newline="") as stream:
|
|
writer = csv.DictWriter(stream, fieldnames=fields)
|
|
writer.writeheader()
|
|
writer.writerows(manifest)
|
|
summary = {
|
|
"frames": len(manifest),
|
|
"segments": {name: sum(row["segment"] == name for row in manifest) for name, _ in args.lidar},
|
|
"rtk_valid": sum(bool(row["rtk_valid"]) for row in manifest),
|
|
"heading_valid": sum(bool(row["heading_valid"]) for row in manifest),
|
|
"rtk_fixed": sum(bool(row["rtk_fixed"]) for row in manifest),
|
|
"imu_window_nonempty": sum(int(row["imu_window_count"]) > 0 for row in manifest),
|
|
"rtk_max_dt_ms": args.rtk_max_dt_ms,
|
|
"imu_window_ms": [-args.imu_before_ms, args.imu_after_ms],
|
|
"time_basis": "LiDAR and serial host UTC; RTK GNSS time and IMU device time are retained for clock-model refinement",
|
|
"imu_orientation_warning": "IMU values are in the raw IMU sensor frame; no LiDAR/body extrinsic is applied",
|
|
}
|
|
(args.out / "dataset_summary.json").write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8")
|
|
print(json.dumps(summary, ensure_ascii=False, indent=2))
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|