ros_contorl/scripts/check_sim_cameras.py

#!/usr/bin/env python3
"""仿真相机的几何验收：把深度反投影到 world，看它是否真的落在该落的平面上。

=============================================================================
为什么需要这个工具（而不是"看图有图像就行"）
=============================================================================
仿真里相机的错误几乎都是**静默**的：话题在、图像有、深度有、TF 也查得到，
而三维点位整体是错的。下面每一类都真实踩过：

  1. 深度单位错了（gz 发 32FC1 米，应用当 16UC1 毫米读）→ 差 1000 倍。
  2. frame_id 盖成了相机**本体**系而不是光学系 → tf2 查得到、不报错，
     但反投影出的点整体转错。
  3. **gz 发的 camera_info 是错的**（恒为 fx=fy=277, cx=160, cy=120，
     即它的默认 60° 视场 + 半分辨率主点），而渲染用的是配置的 87° 视场。
     照抄它反投影，点会偏半米量级。
  4. 相机位姿的 rpy 写错 → 相机**背对**目标，只看得见地面。
     这种情况几何仍然自洽（点确实落在真实地面上），所以只查数值发现不了 ——
     必须比对"该看见什么"，也就是本工具要说的平面。

所以本工具的判据是：**反投影出的点必须命中预期平面**。
"有图像"不算通过，"落在正确的位置"才算。

用法（必须先 source scripts/sim_env.sh，因为需要 ROS 环境）：
    python3 scripts/check_sim_cameras.py
    python3 scripts/check_sim_cameras.py --expect d435:0.40:30 --cameras d435

退出码 0 = 全部通过；非 0 = 有不通过的相机（可直接用于 CI / 启动自检）。
=============================================================================
"""

from __future__ import annotations

import argparse
import math
import sys
import time

import numpy as np

try:
    import rclpy
    from rclpy.node import Node
    from rclpy.qos import qos_profile_sensor_data
    from sensor_msgs.msg import CameraInfo, Image
    from tf2_ros import Buffer, TransformListener
except ImportError as exc:  # pragma: no cover - 环境问题，不是逻辑问题
    print(
        f"✘ 导入 ROS 失败：{exc}\n"
        f"  这个脚本要在仿真环境里跑，先执行：\n"
        f"      source scripts/sim_env.sh && . install_sim/setup.bash",
        file=sys.stderr,
    )
    sys.exit(2)


# 默认期望：哪台相机"必须"看见哪个平面。
#   d435 是手外相机，它的职责就是看见台面（P6 的粗定位要靠它），
#        所以要求 ≥10% 的像素落在台面上。
#   d405 是腕上相机，视野窄且会拍到机械臂自己，台面占比要求放宽到 5%。
#   地面(z=0) 不作为要求：它被拍到多少取决于姿态，不构成"对不对"的判据。
DEFAULT_EXPECT = ("d435:0.40:10", "d405:0.40:5")


def quat_to_R(q) -> np.ndarray:
    """四元数 (x, y, z, w) → 3x3 旋转矩阵。"""
    x, y, z, w = q
    n = x * x + y * y + z * z + w * w
    if n <= 0.0:
        return np.eye(3)
    s = 2.0 / n
    return np.array(
        [
            [1 - s * (y * y + z * z), s * (x * y - z * w), s * (x * z + y * w)],
            [s * (x * y + z * w), 1 - s * (x * x + z * z), s * (y * z - x * w)],
            [s * (x * z - y * w), s * (y * z + x * w), 1 - s * (x * x + y * y)],
        ]
    )


class Grab(Node):
    """抓一帧深度 + 一份内参。"""

    def __init__(self, ns: str) -> None:
        super().__init__(f"check_sim_cameras_{ns}")
        self.ns = ns
        self.img = None
        self.info = None
        self.buf = Buffer()
        self.tl = TransformListener(self.buf, self)
        # QoS 必须与发布端一致：适配器与真机都用 sensor data，
        # 用默认 RELIABLE 订阅会一条都收不到（表现是"话题在但 echo 为空"）。
        self.create_subscription(
            Image,
            f"/{ns}/aligned_depth_to_color/image_raw",
            self._on_img,
            qos_profile_sensor_data,
        )
        self.create_subscription(
            CameraInfo, f"/{ns}/color/camera_info", self._on_info, qos_profile_sensor_data
        )

    def _on_img(self, m) -> None:
        self.img = m

    def _on_info(self, m) -> None:
        self.info = m


def spin_until(node: Node, pred, timeout: float) -> bool:
    t0 = time.time()
    while rclpy.ok() and time.time() - t0 < timeout:
        rclpy.spin_once(node, timeout_sec=0.3)
        if pred():
            return True
    return False


def expected_intrinsics(ns: str):
    """从 cameras.yaml 算出该相机的应有内参。拿不到就返回 None（不阻断检查）。"""
    try:
        from arm_sim_gazebo.sim_cameras import intrinsics, load_cameras
    except ImportError:
        return None
    try:
        for s in load_cameras():
            if s.ros_ns == ns:
                return s, intrinsics(s, s.width or 640, s.height or 480)
    except Exception:
        return None
    return None


def check_one(ns: str, args) -> tuple[bool, list[str]]:
    """检查一台相机。返回 (是否通过, 问题列表)。"""
    problems: list[str] = []
    node = Grab(ns)

    if not spin_until(
        node, lambda: node.img is not None and node.info is not None, args.timeout
    ):
        node.destroy_node()
        return False, [f"拿不到图像或内参（等 {args.timeout:.0f}s）"]

    frame = node.info.header.frame_id
    if not spin_until(
        node, lambda: node.buf.can_transform("world", frame, rclpy.time.Time()), 20.0
    ):
        node.destroy_node()
        return False, [f"TF world ← {frame} 查不到（相机在世上没有位姿）"]

    K = np.array(node.info.k, dtype=np.float64).reshape(3, 3)
    fx, fy, cx, cy = K[0, 0], K[1, 1], K[0, 2], K[1, 2]
    if not (frame.endswith("_optical_frame") or "optical" in frame):
        problems.append(
            f"frame_id 是 {frame}，看着不像光学系（应以 _optical_frame 结尾）。"
            f"盖成本体系时 tf2 查得到、不报错，但反投影出的点会整体转错。"
        )

    mm = np.frombuffer(node.img.data, dtype=np.uint16)
    if mm.size != node.img.height * node.img.width:
        node.destroy_node()
        return False, [
            f"深度图像不是 16UC1 单通道（{node.img.encoding}, "
            f"{len(node.img.data)} 字节 / 期望 {node.img.height * node.img.width * 2}）"
            f"—— gz 的 32FC1 米必须转成 16UC1 毫米，这一步没生效"
        ]
    mm = mm.reshape(node.img.height, node.img.width)

    # 与 cameras.yaml 推导出的应有内参比对：这能直接抓住"gz 的默认内参被透传"
    # 那个坑（它会给出 fx=277, cx=160, cy=120）。
    exp = expected_intrinsics(ns)
    if exp is not None:
        spec, (exp_k, _) = exp
        if abs(exp_k[0] - fx) > 1.0 or abs(exp_k[2] - cx) > 1.0:
            problems.append(
                f"内参不是按 cameras.yaml 算出来的：实测 fx={fx:.1f} cx={cx:.1f}，"
                f"应为 fx={exp_k[0]:.1f} cx={exp_k[2]:.1f}"
                f"（视场 {spec.horizontal_fov}）。\n"
                f"      这正是 gz 透传默认内参（fx=277, cx=160）时的症状，"
                f"会让所有三维点位偏半米量级。"
            )

    valid = (mm > 200) & (mm < 8000)
    vs, us = np.nonzero(valid)
    if len(us) < 500:
        node.destroy_node()
        return False, [f"有效深度像素太少（{len(us)} 个），图像基本是空的"]
    rng = np.random.default_rng(0)
    if len(us) > args.max_pixels:
        sel = rng.choice(len(us), args.max_pixels, replace=False)
        us, vs = us[sel], vs[sel]

    Z = mm[vs, us].astype(np.float64) / 1000.0
    pts = np.stack([(us - cx) * Z / fx, (vs - cy) * Z / fy, Z], axis=1)

    tf = node.buf.lookup_transform("world", frame, rclpy.time.Time())
    R = quat_to_R(
        (
            tf.transform.rotation.x,
            tf.transform.rotation.y,
            tf.transform.rotation.z,
            tf.transform.rotation.w,
        )
    )
    t = np.array(
        [
            tf.transform.translation.x,
            tf.transform.translation.y,
            tf.transform.translation.z,
        ]
    )
    zw = (pts @ R.T + t)[:, 2]
    node.destroy_node()

    print(f"\n[{ns}]  光学系={frame}   fx={fx:.1f} fy={fy:.1f} cx={cx:.1f} cy={cy:.1f}")
    print(f"        world-z 中位={np.median(zw):+.3f} m  范围[{zw.min():+.2f}, {zw.max():+.2f}]")

    # 逐平面命中率
    rates = {}
    for pz, name in args.planes:
        rate = float(np.mean(np.abs(zw - pz) < args.tol)) * 100.0
        rates[pz] = rate
        print(f"        命中 {name}(z={pz:+.2f}) : {rate:5.1f}%")

    # 期望判据
    for want_z, want_min in args.expect.get(ns, []):
        name = dict(args.planes).get(want_z, f"z={want_z}")
        got = rates.get(want_z, 0.0)
        if got < want_min:
            problems.append(
                f"只看见 {name} 的 {got:.1f}%，要求 ≥{want_min:.0f}%。\n"
                f"      相机很可能是**背对**目标（位姿 rpy 写错）—— "
                f"这种错几何自洽、不报错，只会让相机拍不到该拍的东西。"
            )

    total = sum(rates.values())
    if total < 1.0:
        problems.append(f"所有预期平面加起来命中率只有 {total:.1f}%，几乎没有指向场景")

    ok = not problems
    print(f"       → {'✓ 通过' if ok else '✘ 不通过'}")
    for p in problems:
        print(f"       ✘ {p}")
    return ok, problems


def parse_planes(text: str):
    out = []
    for item in text.split(","):
        item = item.strip()
        if not item:
            continue
        if ":" in item:
            z, name = item.split(":", 1)
        else:
            z, name = item, f"z={item}"
        out.append((float(z), name.strip()))
    if not out:
        raise argparse.ArgumentTypeError("--planes 不能为空")
    return out


def parse_expect(items):
    """把 'd435:0.40:10' 解析成 {ns: [(z, min_pct), ...]}。"""
    out: dict[str, list[tuple[float, float]]] = {}
    for it in items:
        parts = it.split(":")
        if len(parts) < 2:
            raise argparse.ArgumentTypeError(f"--expect 格式应为 相机:z[:最低百分比]，收到 {it!r}")
        ns = parts[0]
        z = float(parts[1])
        pct = float(parts[2]) if len(parts) > 2 else 5.0
        out.setdefault(ns, []).append((z, pct))
    return out


def main() -> int:
    ap = argparse.ArgumentParser(
        description="仿真相机几何验收：反投影深度，检查是否命中预期平面。",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    ap.add_argument("--cameras", default="d405,d435", help="要检查的相机命名空间，逗号分隔")
    ap.add_argument(
        "--planes",
        type=parse_planes,
        default=parse_planes("0.0:地面,0.40:台面"),
        help="要统计命中率的平面，格式 z:名称，逗号分隔（默认 0.0:地面,0.40:台面）",
    )
    ap.add_argument(
        "--expect",
        action="append",
        dest="expect_raw",
        default=None,
        help="要求某相机命中某平面，格式 相机:z[:最低百分比]，可重复。"
        f"默认 {list(DEFAULT_EXPECT)}",
    )
    ap.add_argument("--tol", type=float, default=0.02, help="平面命中容差（米），默认 0.02")
    ap.add_argument("--timeout", type=float, default=25.0, help="等图像/内参的秒数")
    ap.add_argument("--max-pixels", type=int, default=40000, help="最多采样多少像素")
    args = ap.parse_args()
    args.expect = parse_expect(args.expect_raw or list(DEFAULT_EXPECT))

    rclpy.init()
    try:
        results = {}
        for ns in [c.strip() for c in args.cameras.split(",") if c.strip()]:
            try:
                results[ns] = check_one(ns, args)
            except KeyboardInterrupt:
                raise
            except Exception as exc:  # 单台相机出错不该拖垮整轮检查
                print(f"\n[{ns}] ✘ 检查时异常：{type(exc).__name__}: {exc}")
                results[ns] = (False, [str(exc)])
    finally:
        rclpy.try_shutdown()

    failed = [ns for ns, (ok, _) in results.items() if not ok]
    print()
    if failed:
        print(f"=== ✘ 不通过：{', '.join(failed)} ===")
        print("    仿真里相机出错几乎都不报错，别只看'有没有图像'。")
        return 1
    print(f"=== ✓ 全部通过：{', '.join(results)} ===")
    return 0


if __name__ == "__main__":
    sys.exit(main())