initial commit

2026-10-22 09:07:39 +08:00
commit 265f11aa38
15 changed files with 2924 additions and 0 deletions
@@ -0,0 +1,611 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+session_server_node.py — 板子端：会话式「抓帧 / 回传图像 / 回传点云 / 误差跟踪」
+
+一次运行 = 一个会话。电脑端每一条请求、板子每一条应答都带 session
+（id / client_ip / board_ip / started_ns），会话之间状态隔离。
+
+pynng REQ/REP 协议（默认 tcp://0.0.0.0:5555）：
+
+  hello       握手；返回标定分辨率、是否支持板子端计算
+  capture     等一帧同步的（图像 + 点云 + 位姿），返回图像字节；
+              本会话记住这一帧（点云、位姿），并清掉旧目标
+  click       默认：电脑只发 (u, v)，板子用本地点云算
+              p_cam / p_lidar / p_map，并把它记成目标（点云不出板）
+  cloud       可选（--compute pc）：返回 capture 那一帧的点云
+  set_target  可选（--compute pc）：电脑算好的 p_map 回传
+  status      返回当前 pos / yaw / target / Δ（只给坐标，不给距离）
+  reset       清除本会话的目标
+
+投影优化：
+  click 第一次会建一个「像素格子 -> 点索引」倒排表（每帧一次），
+  之后同一帧的点击只查周围几个格子，不再全量投影，平面拟合也只在
+  局部候选里做。格子大小 --index-cell（默认 32 px）。
+
+用法（板子）：
+  source /opt/ros/jazzy/setup.bash
+  python3 session_server_node.py --port 5555
+
+常用参数：
+  --image-topic /odin1/image/compressed   （默认走压缩图，直接转发，不用编码）
+  --no-compressed                          （强制用原始 /odin1/image，需要 cv2 或 Pillow）
+  --cloud-topic /odin1/cloud_raw
+  --odom-topic  /odin1/odometry
+"""
+import argparse
+import threading
+import time
+import zlib
+from collections import deque
+
+import numpy as np
+import rclpy
+from rclpy.node import Node
+from rclpy.qos import (QoSProfile, ReliabilityPolicy, HistoryPolicy,
+                       qos_profile_sensor_data)
+from sensor_msgs.msg import Image, CompressedImage, PointCloud2
+from nav_msgs.msg import Odometry
+from message_filters import Subscriber, ApproximateTimeSynchronizer
+
+import pynng
+
+import wire
+
+try:
+    import calib
+    CALIB_W = int(calib.IMAGE_WIDTH)
+    CALIB_H = int(calib.IMAGE_HEIGHT)
+except Exception:  # 没放 calib.py 也能跑，只是分辨率用默认值
+    CALIB_W, CALIB_H = 1600, 1296
+
+# 板子端本地算 p_cam/p_lidar/p_map（默认模式）。没有 target_math.py 时
+# 自动降级为"电脑端算 + 点云回传"。
+try:
+    import target_math
+    _HAS_TARGET_MATH = True
+except Exception:
+    target_math = None
+    _HAS_TARGET_MATH = False
+
+# sensor_msgs/PointField.datatype -> numpy dtype
+_PF = {1: np.int8, 2: np.uint8, 3: np.int16, 4: np.uint16,
+       5: np.int32, 6: np.uint32, 7: np.float32, 8: np.float64}
+_CH = {'bgr8': 3, 'rgb8': 3, 'bgra8': 4, 'rgba8': 4, 'mono8': 1}
+
+
+def stamp_ns(h):
+    return int(h.stamp.sec) * 1_000_000_000 + int(h.stamp.nanosec)
+
+
+def unpack_cloud(msg):
+    """按字段名/偏移解析 PointCloud2（兼容 padding 与非对齐字段）。"""
+    names, formats, offsets = [], [], []
+    for f in msg.fields:
+        if f.datatype not in _PF:
+            continue
+        names.append(f.name)
+        formats.append(_PF[f.datatype])
+        offsets.append(f.offset)
+    dt = np.dtype({'names': names, 'formats': formats,
+                   'offsets': offsets, 'itemsize': msg.point_step})
+    arr = np.frombuffer(msg.data, dtype=dt, count=msg.width * msg.height)
+    return arr, names
+
+
+def image_to_array(msg):
+    ch = _CH.get(msg.encoding)
+    if ch is None:
+        return None
+    row = msg.step if msg.step else msg.width * ch
+    buf = np.frombuffer(msg.data, np.uint8).reshape(msg.height, row)
+    img = buf[:, :msg.width * ch].reshape(msg.height, msg.width, ch)
+    return img.copy()
+
+
+def quat_to_yaw(x, y, z, w):
+    return float(np.arctan2(2.0 * (w * z + x * y),
+                            1.0 - 2.0 * (y * y + z * z)))
+
+
+def quat_to_rot(q):
+    x, y, z, w = [float(v) for v in q]
+    return np.array([
+        [1 - 2 * (y * y + z * z), 2 * (x * y - z * w), 2 * (x * z + y * w)],
+        [2 * (x * y + z * w), 1 - 2 * (x * x + z * z), 2 * (y * z - x * w)],
+        [2 * (x * z - y * w), 2 * (y * z + x * w), 1 - 2 * (x * x + y * y)],
+    ], dtype=np.float64)
+
+
+class SessionServer(Node):
+    def __init__(self, args):
+        super().__init__('session_server')
+        self.args = args
+        self._stop = False
+        self._rep = None
+
+        self._lock = threading.Lock()
+        self._cv = threading.Condition(self._lock)
+        self._latest = None
+        self._odom_hist = deque(maxlen=400)
+        self._sessions = {}
+
+        self._setup_subs()
+
+        self.get_logger().info(
+            '会话服务: image=%s cloud=%s odom=%s'
+            % (args.image_topic, args.cloud_topic, args.odom_topic))
+        self.get_logger().info(
+            '标定分辨率 %dx%d，等待 %s 的第一帧…' % (CALIB_W, CALIB_H, args.image_topic))
+
+    # ------------------------------------------------------------------ 订阅
+    def _setup_subs(self):
+        img_type = CompressedImage if self.args.image_topic.rstrip('/').endswith('/compressed') else Image
+        a = Subscriber(self, img_type, self.args.image_topic,
+                       qos_profile=qos_profile_sensor_data)
+        b = Subscriber(self, PointCloud2, self.args.cloud_topic,
+                       qos_profile=qos_profile_sensor_data)
+        self.sync = ApproximateTimeSynchronizer([a, b],
+                                                queue_size=self.args.queue,
+                                                slop=self.args.slop)
+        self.sync.registerCallback(self.on_pair)
+
+        # 里程计：best-effort 订阅既能接 reliable 也能接 best-effort 的发布者
+        odom_qos = QoSProfile(depth=100,
+                              reliability=ReliabilityPolicy.BEST_EFFORT,
+                              history=HistoryPolicy.KEEP_LAST)
+        self.create_subscription(Odometry, self.args.odom_topic,
+                                 self.on_odom, odom_qos)
+
+    def on_odom(self, msg):
+        p = msg.pose.pose.position
+        q = msg.pose.pose.orientation
+        rec = {
+            'stamp_ns': stamp_ns(msg.header),
+            'recv_mono': time.monotonic(),
+            'pos': (float(p.x), float(p.y), float(p.z)),
+            'quat': (float(q.x), float(q.y), float(q.z), float(q.w)),
+            'frame_id': msg.header.frame_id,
+            'child_frame_id': msg.child_frame_id,
+        }
+        with self._lock:
+            self._odom_hist.append(rec)
+
+    def on_pair(self, img_msg, pc_msg):
+        """图像 + 点云时间对齐后，把这一帧放进 latest。"""
+        arr, names = unpack_cloud(pc_msg)
+        if not all(k in names for k in ('x', 'y', 'z')):
+            self.get_logger().warn('点云没有 x/y/z 字段，实际：%s' % names)
+            return
+
+        xyz = np.stack([arr['x'], arr['y'], arr['z']], axis=-1).astype(np.float32)
+        conf = np.ascontiguousarray(arr['confidence']) if 'confidence' in names else None
+        cloud_stamp = stamp_ns(pc_msg.header)
+        img_stamp = stamp_ns(img_msg.header)
+
+        if isinstance(img_msg, CompressedImage):
+            frame = {
+                'kind': 'compressed',
+                'blob': bytes(img_msg.data),
+                'format': str(img_msg.format or 'jpeg'),
+                'width': None,
+                'height': None,
+            }
+        else:
+            img = image_to_array(img_msg)
+            if img is None:
+                self.get_logger().warn('不支持的图像编码：%s' % img_msg.encoding)
+                return
+            frame = {
+                'kind': 'raw',
+                'img': img,
+                'format': None,
+                'width': int(img_msg.width),
+                'height': int(img_msg.height),
+                'encoding': img_msg.encoding,
+            }
+
+        frame.update({
+            'xyz': xyz,
+            'conf': conf,
+            'cloud_stamp_ns': cloud_stamp,
+            'img_stamp_ns': img_stamp,
+            'cloud_frame_id': pc_msg.header.frame_id,
+            'recv_mono': time.monotonic(),
+        })
+
+        with self._cv:
+            self._latest = frame
+            self._cv.notify_all()
+
+        self.get_logger().debug(
+            '同步帧: 图像 %s, 点云 %d 点, 时间差 %.2f ms'
+            % (frame['kind'], xyz.shape[0], (cloud_stamp - img_stamp) / 1e6))
+
+    # ------------------------------------------------------------------ 帧工具
+    def wait_for_frame(self, after_mono, timeout):
+        """等一帧 recv 时间 >= after_mono 的帧；超时返回当前最新帧。"""
+        deadline = time.monotonic() + timeout
+        with self._cv:
+            while True:
+                frame = self._latest
+                if frame is not None and frame['recv_mono'] >= after_mono:
+                    return dict(frame)
+                left = deadline - time.monotonic()
+                if left <= 0:
+                    return dict(frame) if frame is not None else None
+                self._cv.wait(min(left, 0.2))
+
+    def _nearest_odom_locked(self, stamp_ns_value):
+        if not self._odom_hist:
+            return None
+        return min(self._odom_hist, key=lambda r: abs(r['stamp_ns'] - stamp_ns_value))
+
+    def frame_odom(self, frame):
+        with self._lock:
+            od = self._nearest_odom_locked(frame['cloud_stamp_ns'])
+        if od is None:
+            return None
+        return {
+            'stamp_ns': int(od['stamp_ns']),
+            'position': [float(v) for v in od['pos']],
+            'orientation_xyzw': [float(v) for v in od['quat']],
+            'frame_id': od['frame_id'],
+            'child_frame_id': od['child_frame_id'],
+            'delta_ms': (od['stamp_ns'] - frame['cloud_stamp_ns']) / 1e6,
+        }
+
+    def encode_image(self, frame):
+        """返回 (bytes, format)。压缩图直接转发；原图才编码。"""
+        if frame.get('kind') == 'compressed':
+            return frame['blob'], frame.get('format') or 'jpeg'
+
+        img = frame['img']
+        try:
+            import cv2
+            ok, buf = cv2.imencode(
+                '.jpg', img,
+                [int(cv2.IMWRITE_JPEG_QUALITY), int(self.args.jpeg_quality)])
+            if ok:
+                return buf.tobytes(), 'jpeg'
+        except Exception:
+            pass
+
+        try:
+            import io
+            from PIL import Image as PILImage
+            enc = frame.get('encoding', 'bgr8')
+            rgb = img[..., ::-1] if enc.startswith('bgr') else img[..., :3]
+            b = io.BytesIO()
+            PILImage.fromarray(np.ascontiguousarray(rgb)).save(
+                b, format='JPEG', quality=int(self.args.jpeg_quality))
+            return b.getvalue(), 'jpeg'
+        except Exception as exc:
+            raise RuntimeError('原图编码失败，需要 OpenCV 或 Pillow：%s' % exc)
+
+    # ------------------------------------------------------------------ 会话
+    def _prune(self):
+        now = time.monotonic()
+        for sid in list(self._sessions):
+            if now - self._sessions[sid]['last'] > self.args.session_ttl:
+                del self._sessions[sid]
+
+    def handle_hello(self, hdr):
+        sid = hdr['session']['id']
+        with self._lock:
+            self._sessions.setdefault(sid, {
+                'session': hdr['session'], 'frame': None,
+                'target': None, 'target_info': None,
+                'last': time.monotonic(),
+            })
+            self._sessions[sid]['last'] = time.monotonic()
+        return {
+            'ok': True, 'type': 'hello',
+            'node': self.get_name(),
+            'calib_width': CALIB_W, 'calib_height': CALIB_H,
+            'board_can_compute': bool(_HAS_TARGET_MATH),
+            'compute_modes': (['board', 'pc'] if _HAS_TARGET_MATH else ['pc']),
+            'topics': {'image': self.args.image_topic,
+                       'cloud': self.args.cloud_topic,
+                       'odom': self.args.odom_topic},
+        }
+
+    def handle_capture(self, hdr):
+        self._prune()
+        t_req = time.monotonic()
+        frame = self.wait_for_frame(t_req - self.args.max_frame_age,
+                                    self.args.frame_timeout)
+        if frame is None:
+            return {'ok': False, 'type': 'frame',
+                    'error': '还没有收到同步的图像+点云'}, b''
+
+        blob, fmt = self.encode_image(frame)
+        sid = hdr['session']['id']
+        with self._lock:
+            self._sessions[sid] = {
+                'session': hdr['session'], 'frame': frame,
+                'target': None, 'target_info': None,
+                'last': time.monotonic(),
+            }
+
+        rep = {
+            'ok': True, 'type': 'frame',
+            'image_format': fmt,
+            'image_bytes': len(blob),
+            'image_width': frame.get('width'),
+            'image_height': frame.get('height'),
+            'calib_width': CALIB_W, 'calib_height': CALIB_H,
+            'image_stamp_ns': int(frame['img_stamp_ns']),
+            'cloud_stamp_ns': int(frame['cloud_stamp_ns']),
+            'stamp_delta_ms': (frame['cloud_stamp_ns'] - frame['img_stamp_ns']) / 1e6,
+            'n_points': int(frame['xyz'].shape[0]),
+            'board_can_compute': bool(_HAS_TARGET_MATH),
+            'odom': self.frame_odom(frame),
+        }
+        return rep, blob
+
+    def handle_click(self, hdr):
+        """板子端算：像素 -> p_cam / p_lidar / p_map，并把它记成目标。
+
+        这是默认模式：点云不出板，电脑只发 (u, v)。
+        """
+        if not _HAS_TARGET_MATH:
+            return {'ok': False, 'type': 'target',
+                    'error': '板子没装 target_math.py，请电脑端用 --compute pc'}, b''
+
+        sid = hdr['session']['id']
+        with self._lock:
+            s = self._sessions.get(sid)
+            frame = None if s is None else s.get('frame')
+        if frame is None:
+            return {'ok': False, 'type': 'target', 'error': '先 capture'}, b''
+
+        u = float(hdr['u'])
+        v = float(hdr['v'])
+        radius = float(hdr.get('radius', self.args.radius))
+        fit = bool(hdr.get('fit_plane', True))
+
+        # 每帧只建一次索引；同一帧后续点击只查局部格子
+        with self._lock:
+            index = None if s is None else s.get('index')
+        if index is None:
+            index = target_math.TargetIndex(frame['xyz'], conf=frame.get('conf'),
+                                            cell=self.args.index_cell)
+            with self._lock:
+                if s is not None:
+                    s['index'] = index
+
+        res = index.locate(u, v, radius=radius, fit_plane=fit)
+        out = {
+            'ok': True, 'type': 'target',
+            'found': bool(res['found']),
+            'u': u, 'v': v,
+            'n_valid': int(res['n_valid']),
+            'n_cand': int(res['n_cand']),
+            'plane_used': bool(res['plane_used']),
+            'multi_layer': bool(res['multi_layer']),
+        }
+        if not res['found']:
+            return out, b''
+
+        out.update({
+            'p_cam': [float(x) for x in res['p_cam']],
+            'p_lidar': [float(x) for x in res['p_lidar']],
+            'hit_u': float(res['hit_u']),
+            'hit_v': float(res['hit_v']),
+            'px_err': float(res['px_err']),
+            'confidence': (None if res.get('confidence') is None
+                           else float(res['confidence'])),
+        })
+
+        od = self.frame_odom(frame)
+        if od is None:
+            out['p_map'] = None
+            out['error'] = '本帧没有里程计位姿'
+            return out, b''
+
+        R = quat_to_rot(od['orientation_xyzw'])
+        t = np.asarray(od['position'], dtype=np.float64)
+        pm = R @ np.asarray(res['p_lidar'], dtype=np.float64) + t
+
+        out['p_map'] = [float(x) for x in pm]
+        out['odom'] = od
+        with self._lock:
+            s['target'] = pm
+            s['target_info'] = out
+            s['last'] = time.monotonic()
+        return out, b''
+
+    def handle_cloud(self, hdr):
+        sid = hdr['session']['id']
+        with self._lock:
+            s = self._sessions.get(sid)
+            frame = None if s is None else s.get('frame')
+        if frame is None:
+            return {'ok': False, 'type': 'cloud', 'error': '先 capture'}, b''
+
+        xyz = np.asarray(frame['xyz'], dtype='<f4')
+        body = zlib.compress(xyz.tobytes(), 1)
+        rep = {
+            'ok': True, 'type': 'cloud',
+            'encoding': 'zlib_f32',
+            'n_points': int(xyz.shape[0]),
+            'shape': [int(xyz.shape[0]), 3],
+            'raw_bytes': int(xyz.nbytes),
+            'zlib_bytes': int(len(body)),
+            'cloud_stamp_ns': int(frame['cloud_stamp_ns']),
+        }
+        return rep, body
+
+    def handle_set_target(self, hdr):
+        sid = hdr['session']['id']
+        with self._lock:
+            s = self._sessions.get(sid)
+            if s is None or s.get('frame') is None:
+                return {'ok': False, 'type': 'target', 'error': '先 capture'}, b''
+            pm = hdr.get('p_map')
+            if not pm or len(pm) != 3:
+                return {'ok': False, 'type': 'target',
+                        'error': 'p_map 需要 3 个数'}, b''
+            s['target'] = np.asarray([float(pm[0]), float(pm[1]), float(pm[2])],
+                                     dtype=np.float64)
+            s['target_info'] = {k: hdr.get(k) for k in ('u', 'v', 'p_cam', 'p_lidar', 'p_map')}
+            s['last'] = time.monotonic()
+            target = [float(v) for v in s['target']]
+        return {'ok': True, 'type': 'target', 'target': target}, b''
+
+    def handle_status(self, hdr):
+        sid = hdr['session']['id']
+        with self._lock:
+            s = self._sessions.get(sid)
+            od = self._odom_hist[-1] if self._odom_hist else None
+            target = None if s is None else s.get('target')
+
+        out = {'ok': True, 'type': 'status', 'has_target': target is not None}
+        if od is None:
+            out.update({'valid': False, 'error': '还没有里程计'})
+            return out, b''
+
+        age = time.monotonic() - od['recv_mono']
+        pos = np.asarray(od['pos'], dtype=np.float64)
+        out.update({
+            'valid': bool(age <= self.args.odom_timeout),
+            'pos': [float(v) for v in pos],
+            'yaw_deg': float(np.rad2deg(quat_to_yaw(*od['quat']))),
+            'odom_age_s': float(age),
+            'odom_frame': od['frame_id'],
+            'odom_child': od['child_frame_id'],
+        })
+        if target is not None:
+            d = pos - np.asarray(target, dtype=np.float64)
+            out['target'] = [float(v) for v in target]
+            out['delta'] = [float(v) for v in d]   # 只给坐标差，不给距离
+        return out, b''
+
+    def handle_reset(self, hdr):
+        sid = hdr['session']['id']
+        with self._lock:
+            s = self._sessions.get(sid)
+            if s is not None:
+                s['target'] = None
+                s['target_info'] = None
+                s['last'] = time.monotonic()
+        return {'ok': True, 'type': 'reset'}, b''
+
+    def dispatch(self, typ, hdr):
+        if typ == 'hello':
+            return self.handle_hello(hdr), b''
+        if typ == 'capture':
+            return self.handle_capture(hdr)
+        if typ == 'click':
+            return self.handle_click(hdr)
+        if typ == 'cloud':
+            return self.handle_cloud(hdr)
+        if typ == 'set_target':
+            return self.handle_set_target(hdr)
+        if typ == 'status':
+            return self.handle_status(hdr)
+        if typ == 'reset':
+            return self.handle_reset(hdr)
+        return {'ok': False, 'type': typ, 'error': 'unknown type'}, b''
+
+    # ------------------------------------------------------------------ 服务循环
+    def serve(self):
+        addr = 'tcp://0.0.0.0:%d' % self.args.port
+        self._rep = pynng.Rep0(listen=addr)
+        self._rep.recv_max_size = 32 * 1024 * 1024
+        self._rep.recv_timeout = 500
+        self.get_logger().info('pynng REP 监听 %s' % addr)
+
+        while not self._stop:
+            try:
+                data = self._rep.recv()
+            except pynng.Timeout:
+                continue
+            except pynng.Closed:
+                break
+            except Exception as exc:
+                self.get_logger().error('recv 失败: %s' % exc)
+                continue
+
+            hdr = {}
+            try:
+                hdr, _ = wire.unpack(data)
+                typ = hdr.get('type', '?')
+                sess = hdr.get('session') or {}
+                sid = sess.get('id', '?')
+                reply, out_blob = self.dispatch(typ, hdr)
+                reply['session'] = sess          # 原样带回，含 IP
+                reply['server_ns'] = time.time_ns()
+                self._rep.send(wire.pack(reply, out_blob))
+                self.get_logger().info(
+                    '[%s] %-10s from %s -> ok=%s'
+                    % (sid[:8], typ, sess.get('client_ip', '?'), reply.get('ok')))
+            except Exception as exc:
+                try:
+                    self._rep.send(wire.pack({
+                        'ok': False, 'type': hdr.get('type', 'error'),
+                        'error': str(exc),
+                        'session': hdr.get('session') or {},
+                        'server_ns': time.time_ns(),
+                    }))
+                except Exception:
+                    pass
+                self.get_logger().error('处理请求失败: %s' % exc)
+
+    def stop(self):
+        self._stop = True
+        if self._rep is not None:
+            try:
+                self._rep.close()
+            except Exception:
+                pass
+
+
+def parse_args():
+    ap = argparse.ArgumentParser(description='会话式抓帧/定位/误差服务')
+    ap.add_argument('--port', type=int, default=5555)
+    ap.add_argument('--image-topic', default='/odin1/image/compressed')
+    ap.add_argument('--cloud-topic', default='/odin1/cloud_raw')
+    ap.add_argument('--odom-topic', default='/odin1/odometry')
+    ap.add_argument('--no-compressed', action='store_true',
+                    help='强制用原始 /odin1/image（需要 cv2 或 Pillow）')
+    ap.add_argument('--slop', type=float, default=0.08, help='图像/点云允许时间差(秒)')
+    ap.add_argument('--queue', type=int, default=20)
+    ap.add_argument('--max-frame-age', type=float, default=0.35,
+                    help='抓帧时允许复用多旧的帧(秒)')
+    ap.add_argument('--frame-timeout', type=float, default=2.0,
+                    help='等新帧的最长时间(秒)')
+    ap.add_argument('--jpeg-quality', type=int, default=85)
+    ap.add_argument('--radius', type=float, default=10.0,
+                    help='click 请求没带半径时的默认候选像素半径')
+    ap.add_argument('--index-cell', type=int, default=32,
+                    help='像素格子大小(px)，用于每帧的投影索引')
+    ap.add_argument('--odom-timeout', type=float, default=1.0,
+                    help='里程计超过这个时间没更新就标记 invalid(秒)')
+    ap.add_argument('--session-ttl', type=float, default=600.0)
+    args = ap.parse_args()
+    if args.no_compressed:
+        args.image_topic = '/odin1/image'
+    return args
+
+
+def main():
+    args = parse_args()
+    rclpy.init()
+    node = SessionServer(args)
+    th = threading.Thread(target=node.serve, daemon=True)
+    th.start()
+    try:
+        rclpy.spin(node)
+    except KeyboardInterrupt:
+        pass
+    finally:
+        node.stop()
+        node.destroy_node()
+        if rclpy.ok():
+            rclpy.shutdown()
+
+
+if __name__ == '__main__':
+    main()