calibration/tools/export_h32_rscap_station.py

#!/usr/bin/env python3
"""Export one static-station H32 capture into LiDAR frame NPZs.

Supports:

- V2 ``.rscap`` (legacy MSOP-only RawCapture)
- Medulla dlog from ``RSLidarH32_3D_DLogCaptureNet48`` (raw MSOP + DIFOP)

This is an **internal** helper used by ``export_raw_to_combined.py``.
For RTK–LiDAR calibration, prefer the one-shot exporter that writes ``combined/``.

Output frame contract (consumed by ``build_multisensor_npz.py``):

- ``points_raw``: (N, 5) polar ``d_mm, azimuth_deg, altitude_deg, intensity, progression``
- ``unix_time_ns``: H32 MSOP device timestamp (seconds+us → ns)
- ``frame_counter``, ``point_count``, optional host receive stamp

Raw ``.rscap`` / dlog files are never modified.
"""

from __future__ import annotations

import argparse
import csv
import json
import sys
from pathlib import Path
from typing import Any

import numpy as np

ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(ROOT))
sys.path.insert(0, str(ROOT / "rscap_v2"))

from capture_format_v2 import file_summary, read_capture  # noqa: E402
from h32_dlog.dobject import discover_records, resolve_dlog_root  # noqa: E402
from h32_dlog.load_session import load_h32_dlog_lidar  # noqa: E402
from h32_msop import (  # noqa: E402
    iter_h32_frames_polar,
    iter_h32_frames_polar_from_packets,
)


def resolve_lidar_rscap(station_dir: Path, capture_name: str = "h32.rscap") -> Path:
    candidates = [
        station_dir / capture_name,
        station_dir / "h32.rscap",
        station_dir / "lidar.rscap",
    ]
    for path in candidates:
        if path.is_file():
            return path
    raise FileNotFoundError(
        f"no LiDAR .rscap under {station_dir}; tried {[str(p.name) for p in candidates]}"
    )


def try_resolve_dlog_root(station_dir: Path) -> Path | None:
    try:
        return resolve_dlog_root(station_dir)
    except FileNotFoundError:
        return None


def is_h32_raw_dlog_station(
    station_dir: Path,
    *,
    msop_object: str = "frontlidar-msop-raw",
) -> bool:
    root = try_resolve_dlog_root(station_dir)
    if root is None:
        return False
    return len(discover_records(root, msop_object)) > 0


def is_legacy_pointcloud_dlog_station(
    station_dir: Path,
    *,
    msop_object: str = "frontlidar-msop-raw",
) -> bool:
    root = try_resolve_dlog_root(station_dir)
    if root is None:
        return False
    return not is_h32_raw_dlog_station(station_dir, msop_object=msop_object)


def _write_polar_frames(
    *,
    out: Path,
    frames,
    source_label: str,
    compress: bool,
    write_reports: bool,
    resume: bool,
    metadata_extra: dict[str, Any],
) -> dict[str, Any]:
    frames_dir = out / "frames"
    frames_dir.mkdir(parents=True, exist_ok=True)
    saver = np.savez_compressed if compress else np.savez
    manifest_rows: list[dict[str, Any]] = []
    written = 0
    for index, frame in enumerate(frames):
        unix_time_ns = int(round(frame.t_start_s * 1_000_000_000))
        name = f"h32_{index:06d}_{unix_time_ns}_frame{index}.npz"
        destination = frames_dir / name
        if resume and destination.exists():
            continue
        points = np.asarray(frame.points_raw, dtype=np.float32)
        payload = {
            "points_raw": points,
            "frame_counter": np.asarray([index], dtype=np.int32),
            "point_count": np.asarray([points.shape[0]], dtype=np.int32),
            "unix_time_ns": np.asarray([unix_time_ns], dtype=np.int64),
            "device_time_s": np.asarray([frame.t_start_s], dtype=np.float64),
            "device_time_end_s": np.asarray([frame.t_end_s], dtype=np.float64),
            "host_receive_utc_ns": np.asarray([frame.host_receive_utc_ns], dtype=np.int64),
            "source_file_utf8": np.frombuffer(source_label.encode("utf-8"), dtype=np.uint8),
        }
        saver(destination, **payload)
        written += 1
        manifest_rows.append(
            {
                "index": index,
                "output": name,
                "unix_time_ns": unix_time_ns,
                "point_count": int(points.shape[0]),
                "host_receive_utc_ns": int(frame.host_receive_utc_ns),
            }
        )

    metadata: dict[str, Any] = {
        "frames_decoded": len(frames),
        "frames_written": written,
        "frames_dir": str(frames_dir.resolve()),
        "time_basis": "H32 MSOP device timestamp (packet seconds+microseconds)",
        "points_raw_columns": ["d_mm", "azimuth_deg", "altitude_deg", "intensity", "progression"],
        **metadata_extra,
    }
    (out / "metadata.json").write_text(json.dumps(metadata, ensure_ascii=False, indent=2), encoding="utf-8")
    (out / "README.md").write_text(
        "# H32 station export (internal)\n\n"
        f"- source: `{source_label}`\n"
        f"- frames: `{frames_dir}`\n"
        "- Prefer ``tools/export_raw_to_combined.py`` for the full RTK–LiDAR package.\n",
        encoding="utf-8",
    )
    if write_reports:
        reports = out / "reports"
        reports.mkdir(parents=True, exist_ok=True)
        with (reports / "manifest.csv").open("w", encoding="utf-8", newline="") as stream:
            writer = csv.DictWriter(
                stream, fieldnames=list(manifest_rows[0].keys()) if manifest_rows else ["index"]
            )
            writer.writeheader()
            writer.writerows(manifest_rows)
        (reports / "export_summary.json").write_text(
            json.dumps(metadata, ensure_ascii=False, indent=2), encoding="utf-8"
        )
    return metadata


def export_station_h32(
    station: Path,
    out: Path,
    *,
    capture_name: str = "h32.rscap",
    stride: int = 1,
    min_frame_points: int = 100,
    min_range_m: float = 0.3,
    max_range_m: float = 120.0,
    compress: bool = True,
    write_reports: bool = False,
    resume: bool = False,
) -> dict[str, Any]:
    """Decode one station H32 ``.rscap`` into ``out/frames/*.npz``."""

    rscap = (
        station
        if station.is_file() and station.suffix.lower() == ".rscap"
        else resolve_lidar_rscap(station, capture_name)
    )
    capture = read_capture(rscap)
    frames = iter_h32_frames_polar(
        capture,
        min_frame_points=min_frame_points,
        frame_stride=max(1, stride),
        min_range_m=min_range_m,
        max_range_m=max_range_m,
    )
    if not frames:
        raise RuntimeError(f"no H32 frames decoded from {rscap}")
    return _write_polar_frames(
        out=out,
        frames=frames,
        source_label=str(rscap.resolve()),
        compress=compress,
        write_reports=write_reports,
        resume=resume,
        metadata_extra={
            "kind": "h32_rscap",
            "source_rscap": str(rscap.resolve()),
            "capture": file_summary(capture),
            "angle_source": "default_msop_only_vertical_-16_to_16_deg",
        },
    )


def export_station_h32_dlog(
    station: Path,
    out: Path,
    *,
    msop_object: str = "frontlidar-msop-raw",
    difop_object: str = "frontlidar-difop-raw",
    require_difop: bool = True,
    stride: int = 1,
    min_frame_points: int = 100,
    min_range_m: float = 0.3,
    max_range_m: float = 120.0,
    compress: bool = True,
    write_reports: bool = False,
    resume: bool = False,
) -> dict[str, Any]:
    """Decode one station H32 raw-MSOP/DIFOP dlog into ``out/frames/*.npz``."""

    session = load_h32_dlog_lidar(
        station,
        msop_object=msop_object,
        difop_object=difop_object,
        require_difop=require_difop,
    )
    frames = iter_h32_frames_polar_from_packets(
        session.msop_packets,
        host_utc_ticks=session.msop_host_utc_ticks,
        min_frame_points=min_frame_points,
        frame_stride=max(1, stride),
        min_range_m=min_range_m,
        max_range_m=max_range_m,
        vertical_deg=session.vertical_deg,
        horizontal_deg=session.horizontal_deg,
    )
    if not frames:
        raise RuntimeError(f"no H32 frames decoded from dlog {session.dlog_root}")
    return _write_polar_frames(
        out=out,
        frames=frames,
        source_label=str(session.dlog_root.resolve()),
        compress=compress,
        write_reports=write_reports,
        resume=resume,
        metadata_extra={
            "kind": "h32_dlog_raw",
            "source_dlog": str(session.dlog_root.resolve()),
            "msop_object": session.msop_object,
            "difop_object": session.difop_object,
            "msop_packets": len(session.msop_packets),
            "msop_batches": session.msop_batch_count,
            "difop_records": session.difop_record_count,
            "session_id": session.session_id,
            "lidar_ip": session.lidar_ip,
            "angle_source": session.angle_source,
        },
    )


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--station", type=Path, required=True, help="Station directory, .rscap, or dlog root")
    parser.add_argument("--out", type=Path, required=True)
    parser.add_argument("--capture-name", default="h32.rscap")
    parser.add_argument("--msop-object", default="frontlidar-msop-raw")
    parser.add_argument("--difop-object", default="frontlidar-difop-raw")
    parser.add_argument(
        "--require-difop",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="For dlog stations, require valid DIFOP angles (default: true)",
    )
    parser.add_argument("--stride", type=int, default=1)
    parser.add_argument("--min-frame-points", type=int, default=100)
    parser.add_argument("--min-range-m", type=float, default=0.3)
    parser.add_argument("--max-range-m", type=float, default=120.0)
    parser.add_argument("--compress", action="store_true", default=True)
    parser.add_argument("--write-reports", action="store_true")
    parser.add_argument("--resume", action="store_true", help="Skip frames that already exist")
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    station = args.station
    if station.is_file() and station.suffix.lower() == ".rscap":
        metadata = export_station_h32(
            station,
            args.out,
            capture_name=args.capture_name,
            stride=args.stride,
            min_frame_points=args.min_frame_points,
            min_range_m=args.min_range_m,
            max_range_m=args.max_range_m,
            compress=args.compress,
            write_reports=args.write_reports,
            resume=args.resume,
        )
    elif is_h32_raw_dlog_station(station, msop_object=args.msop_object):
        metadata = export_station_h32_dlog(
            station,
            args.out,
            msop_object=args.msop_object,
            difop_object=args.difop_object,
            require_difop=args.require_difop,
            stride=args.stride,
            min_frame_points=args.min_frame_points,
            min_range_m=args.min_range_m,
            max_range_m=args.max_range_m,
            compress=args.compress,
            write_reports=args.write_reports,
            resume=args.resume,
        )
    else:
        metadata = export_station_h32(
            station,
            args.out,
            capture_name=args.capture_name,
            stride=args.stride,
            min_frame_points=args.min_frame_points,
            min_range_m=args.min_range_m,
            max_range_m=args.max_range_m,
            compress=args.compress,
            write_reports=args.write_reports,
            resume=args.resume,
        )
    print(json.dumps(metadata, ensure_ascii=False, indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())