flyguard/flyguard_ros2_node.py

@@ -1,216 +1,235 @@
-"""
+"""ROS 2 adapter for the FlyGuard lidar-obstacle detector.
-ROS 2 узел для биоинспирированного пайплайна FlyGuard (ЛЦТ-2026 Кейс 05).
+
-Связывает входные данные 3D-лидара с ядром FlyGuard и публикует детекции/маркеры.
+The processing core deliberately has no ROS dependency. This node adapts
+``sensor_msgs/msg/PointCloud2`` to the core's ``flyguard.cdr.PointCloud2``
+contract and publishes final decisions for consumers and RViz.
 """
 
+from __future__ import annotations
+
+import math
 import sys
 from pathlib import Path
+
 import numpy as np
 
-# Автоматически добавляем корень репозитория в sys.path
 ROOT_DIR = Path(__file__).resolve().parent.parent
 if str(ROOT_DIR) not in sys.path:
     sys.path.insert(0, str(ROOT_DIR))
 
 import rclpy
+from rclpy.executors import ExternalShutdownException
 from rclpy.node import Node
+from rclpy.qos import qos_profile_sensor_data
 
-# Импорт сообщений ROS 2
-from sensor_msgs.msg import PointCloud2
-from std_msgs.msg import String
-from vision_msgs.msg import Detection3DArray, Detection3D, BoundingBox3D
-from visualization_msgs.msg import MarkerArray, Marker
 from geometry_msgs.msg import Vector3
-
+from sensor_msgs.msg import PointCloud2 as RosPointCloud2
 import sensor_msgs_py.point_cloud2 as pc2
+from std_msgs.msg import String
+from vision_msgs.msg import BoundingBox3D, Detection3D, Detection3DArray
+from visualization_msgs.msg import Marker, MarkerArray
 
-# Импорт ядра FlyGuard
+from flyguard.cdr import PointCloud2 as CorePointCloud2
-try:
+from flyguard.export import export_frame
-    from flyguard.pipeline import FlyGuard, Params
+from flyguard.mbon_readout import MbonReadout
 from flyguard.mushroom_body import MushroomBody
-except ImportError as e:
+from flyguard.pipeline import FlyGuard, Params
-    FlyGuard = None
-    Params = None
-    MushroomBody = None
 
 
 class FlyGuardNode(Node):
-    def __init__(self):
+    """Consume lidar clouds and publish final FlyGuard decisions."""
-        super().__init__('flyguard_node')
 
-        # --- Объявление параметров ---
+    def __init__(self) -> None:
-        self.declare_parameter('lidar_topic', '/pandar_points')
+        super().__init__("flyguard_node")
-        self.declare_parameter('frame_id', 'pandar')
+        self.declare_parameter("lidar_topic", "/pandar_points")
-        self.declare_parameter('memory_path', '')
+        # Empty means preserve the PointCloud2 header frame for downstream use.
-        self.declare_parameter('fov_deg', 30.0)
+        self.declare_parameter("frame_id", "")
+        self.declare_parameter("memory_path", "artifacts/mushroom_body.npz")
+        self.declare_parameter("mbon_path", "artifacts/mbon_readout.npz")
+        self.declare_parameter("fov_deg", 30.0)
 
-        self.lidar_topic = self.get_parameter('lidar_topic').get_parameter_value().string_value
+        self.lidar_topic = self.get_parameter("lidar_topic").value
-        self.frame_id = self.get_parameter('frame_id').get_parameter_value().string_value
+        self.frame_id_override = self.get_parameter("frame_id").value
-        memory_path = self.get_parameter('memory_path').get_parameter_value().string_value
+        memory_path = self.get_parameter("memory_path").value
-        fov_deg = self.get_parameter('fov_deg').get_parameter_value().double_value
+        mbon_path = self.get_parameter("mbon_path").value
+        fov_deg = self.get_parameter("fov_deg").value
 
-        # --- Инициализация ядра FlyGuard ---
+        memory = self._load_model(memory_path, "memory", MushroomBody.load)
-        if FlyGuard is not None:
+        readout = self._load_model(mbon_path, "MBON readout", MbonReadout.load)
-            self.get_logger().info("Инициализация ядра FlyGuard...")
+        self.fg = FlyGuard(Params(fov_deg=float(fov_deg)), memory=memory,
-            memory = MushroomBody.load(memory_path) if memory_path else None
+                           readout=readout)
-            params = Params(fov_deg=fov_deg)
-            self.fg = FlyGuard(params, memory=memory)
-        else:
-            self.get_logger().error("Модуль flyguard.pipeline не найден!")
-            self.fg = None
 
-        # --- Подписчики и Публикаторы ---
+        # Pandar and most lidar drivers offer BEST_EFFORT sensor QoS. A
+        # default RELIABLE subscriber is incompatible and would receive no data.
         self.sub_cloud = self.create_subscription(
-            PointCloud2,
+            RosPointCloud2, self.lidar_topic, self.pointcloud_callback,
-            self.lidar_topic,
+            qos_profile_sensor_data,
-            self.pointcloud_callback,
+        )
-            10
+        self.pub_threat = self.create_publisher(String, "flyguard/threat_level", 10)
+        self.pub_boxes = self.create_publisher(
+            Detection3DArray, "flyguard/bounding_boxes", 10)
+        self.pub_markers = self.create_publisher(MarkerArray, "flyguard/markers", 10)
+
+        self.received_frames = 0
+        self.processed_frames = 0
+        self.get_logger().info(
+            "FlyGuard started: input=%s, frame=%s, memory=%s, mbon=%s"
+            % (self.lidar_topic, self.frame_id_override or "<input header>",
+               memory_path or "disabled", mbon_path or "disabled")
         )
 
-        self.pub_threat = self.create_publisher(String, 'flyguard/threat_level', 10)
+    def _load_model(self, value: str, label: str, loader):
-        self.pub_boxes = self.create_publisher(Detection3DArray, 'flyguard/bounding_boxes', 10)
+        """Load an optional model and fail early for a supplied bad path."""
-        self.pub_markers = self.create_publisher(MarkerArray, 'flyguard/markers', 10)
+        if not value:
+            self.get_logger().warning(f"{label} is disabled (empty path)")
+            return None
+        path = Path(value)
+        if not path.is_absolute():
+            path = ROOT_DIR / path
+        if not path.is_file():
+            raise FileNotFoundError(
+                f"{label} file does not exist: {path}. Pass an empty parameter "
+                "only when the model is intentionally disabled."
+            )
+        self.get_logger().info(f"Loading {label} from {path}")
+        return loader(path)
 
-        self.get_logger().info(f"FlyGuard ROS 2 Node запущен! Топик: {self.lidar_topic}")
+    @staticmethod
+    def _to_core_cloud(msg: RosPointCloud2) -> CorePointCloud2:
+        """Adapt ROS metadata and retain the complete raw scan ordering."""
+        # Do not skip NaNs: removing no-return points changes the point order
+        # and prevents the core from recognising an organised Pandar scan.
+        points = np.asarray(pc2.read_points(msg, skip_nans=False))
+        names = points.dtype.names or ()
+        missing = {"x", "y", "z", "intensity"}.difference(names)
+        if missing:
+            raise ValueError("PointCloud2 is missing required fields: %s" %
+                             ", ".join(sorted(missing)))
+        expected = int(msg.height) * int(msg.width)
+        if points.size != expected:
+            raise ValueError(
+                f"PointCloud2 has {points.size} decoded points, expected {expected}")
+        stamp = float(msg.header.stamp.sec) + float(msg.header.stamp.nanosec) * 1e-9
+        return CorePointCloud2(
+            stamp=stamp,
+            frame_id=msg.header.frame_id,
+            height=int(msg.height),
+            width=int(msg.width),
+            point_step=int(msg.point_step),
+            is_dense=bool(msg.is_dense),
+            points=points,
+        )
 
-    def pointcloud_callback(self, msg: PointCloud2):
+    def pointcloud_callback(self, msg: RosPointCloud2) -> None:
-        """Обработка каждого кадра лидара (10 Гц)"""
+        self.received_frames += 1
-        start_time = self.get_clock().now()
+        started = self.get_clock().now()
-
+        try:
-        # 1. Читаем точки из PointCloud2 в structured numpy array
+            cloud = self._to_core_cloud(msg)
-        # (сохраняет стандартную структуру PointCloud)
+            if cloud.n_points == 0:
-        pc_data = pc2.read_points(msg, skip_nans=True)
+                self.get_logger().warning("Ignoring an empty PointCloud2 frame")
-        if pc_data.size == 0:
+                return
+            result = self.fg.process(cloud)
+        except Exception as exc:
+            # A malformed frame must not kill the ROS executor.
+            self.get_logger().error(f"FlyGuard failed to process cloud: {exc}")
             return
 
-        # 2. Вызываем ядро FlyGuard
+        # The first calib_frames scans intentionally produce no result.
-        if self.fg is None:
+        if result is None:
             return
 
-        res = self.fg.process(pc_data)
+        self.processed_frames += 1
-        if res is None:
+        decision = result.decision
-            return
+        frame_id = self.frame_id_override or msg.header.frame_id
+        exported = export_frame(decision, result.plane, result.corridor,
+                                stamp=result.stamp)
+        threat_level = ("EMERGENCY" if decision.emergency else
+                        "WARNING" if decision.detected else "CLEAR")
 
-        d = res.decision
+        self.pub_threat.publish(String(data=threat_level))
-        tracks = getattr(self.fg.cx, 'tracks', [])
+        self.publish_detections(exported.boxes, msg.header.stamp, frame_id)
+        self.publish_rviz_markers(exported.boxes, msg.header.stamp, frame_id)
 
-        # 3. Определяем уровень угрозы
+        elapsed_ms = (self.get_clock().now() - started).nanoseconds / 1e6
-        if d.detected:
+        self.get_logger().debug(
-            threat_level = "EMERGENCY" if getattr(d, 'emergency', False) else "WARNING"
+            "frame=%d processed=%d %.1f ms: %s, %d object(s)"
+            % (self.received_frames, self.processed_frames, elapsed_ms,
+               threat_level, len(exported.boxes))
+        )
+
+    def publish_detections(self, boxes, stamp, frame_id: str) -> None:
+        message = Detection3DArray()
+        message.header.stamp = stamp
+        message.header.frame_id = frame_id
+        for box in boxes:
+            detection = Detection3D()
+            detection.header = message.header
+            detection.bbox = BoundingBox3D()
+            detection.bbox.center.position.x = box.x
+            detection.bbox.center.position.y = box.y
+            detection.bbox.center.position.z = box.z
+            detection.bbox.center.orientation.z = math.sin(box.yaw * 0.5)
+            detection.bbox.center.orientation.w = math.cos(box.yaw * 0.5)
+            detection.bbox.size = Vector3(x=box.dx, y=box.dy, z=box.dz)
+            message.detections.append(detection)
+        self.pub_boxes.publish(message)
+
+    def publish_rviz_markers(self, boxes, stamp, frame_id: str) -> None:
+        message = MarkerArray()
+        clear = Marker()
+        clear.header.stamp = stamp
+        clear.header.frame_id = frame_id
+        clear.action = Marker.DELETEALL
+        message.markers.append(clear)
+
+        for box in boxes:
+            cube = Marker()
+            cube.header.stamp = stamp
+            cube.header.frame_id = frame_id
+            cube.ns = "flyguard_boxes"
+            cube.id = int(box.track_id)
+            cube.type = Marker.CUBE
+            cube.action = Marker.ADD
+            cube.pose.position.x, cube.pose.position.y, cube.pose.position.z = box.x, box.y, box.z
+            cube.pose.orientation.z = math.sin(box.yaw * 0.5)
+            cube.pose.orientation.w = math.cos(box.yaw * 0.5)
+            cube.scale = Vector3(x=max(box.dx, 0.01), y=max(box.dy, 0.01), z=max(box.dz, 0.01))
+            if box.threat_level.name == "EMERGENCY":
+                cube.color.r, cube.color.g, cube.color.b, cube.color.a = 1.0, 0.0, 0.0, 0.75
             else:
-            threat_level = "CLEAR"
+                cube.color.r, cube.color.g, cube.color.b, cube.color.a = 1.0, 0.85, 0.0, 0.65
+            message.markers.append(cube)
 
-        # 4. Публикация статуса
+            label = Marker()
-        threat_msg = String()
+            label.header.stamp = stamp
-        threat_msg.data = threat_level
+            label.header.frame_id = frame_id
-        self.pub_threat.publish(threat_msg)
+            label.ns = "flyguard_labels"
-
+            label.id = 100000 + int(box.track_id)
-        # 5. Публикация маркеров и рамок
+            label.type = Marker.TEXT_VIEW_FACING
-        stamp = msg.header.stamp
+            label.action = Marker.ADD
-        self.publish_detections(tracks, stamp)
+            label.pose.position.x, label.pose.position.y = box.x, box.y
-        self.publish_rviz_markers(tracks, threat_level, d, stamp)
+            label.pose.position.z = box.z + box.dz * 0.5 + 0.35
-
+            label.pose.orientation.w = 1.0
-        elapsed_ms = (self.get_clock().now() - start_time).nanoseconds / 1e6
+            label.scale.z = 0.40
-        self.get_logger().debug(f"Кадр за {elapsed_ms:.1f} мс | Статус: {threat_level} | Треков: {len(tracks)}")
+            label.color.r = label.color.g = label.color.b = label.color.a = 1.0
-
+            label.text = "ID:%d | %.1f m | conf:%.2f" % (
-    def publish_detections(self, tracks, stamp):
+                box.track_id, box.distance_along_track, box.confidence)
-        """Формирование vision_msgs/Detection3DArray"""
+            message.markers.append(label)
-        det_array = Detection3DArray()
+        self.pub_markers.publish(message)
-        det_array.header.stamp = stamp
-        det_array.header.frame_id = self.frame_id
-
-        for trk in tracks:
-            det = Detection3D()
-            det.header = det_array.header
-
-            bbox = BoundingBox3D()
-            bbox.center.position.x = float(getattr(trk, 'x', 0.0))
-            bbox.center.position.y = float(getattr(trk, 'y', 0.0))
-            bbox.center.position.z = float(getattr(trk, 'z', 0.0))
-
-            bbox.size.x = float(getattr(trk, 'dx', 0.5))
-            bbox.size.y = float(getattr(trk, 'dy', 0.5))
-            bbox.size.z = float(getattr(trk, 'dz', 1.5))
-
-            det.bbox = bbox
-            det_array.detections.append(det)
-
-        self.pub_boxes.publish(det_array)
-
-    def publish_rviz_markers(self, tracks, threat_level: str, decision, stamp):
-        """Генерация маркеров для RViz2"""
-        marker_array = MarkerArray()
-
-        delete_all = Marker()
-        delete_all.action = Marker.DELETEALL
-        marker_array.markers.append(delete_all)
-
-        for i, trk in enumerate(tracks):
-            # Куб вокруг объекта
-            box_marker = Marker()
-            box_marker.header.stamp = stamp
-            box_marker.header.frame_id = self.frame_id
-            box_marker.ns = "flyguard_boxes"
-            box_marker.id = i
-            box_marker.type = Marker.CUBE
-            box_marker.action = Marker.ADD
-
-            x = float(getattr(trk, 'x', 0.0))
-            y = float(getattr(trk, 'y', 0.0))
-            z = float(getattr(trk, 'z', 0.0))
-
-            box_marker.pose.position.x = x
-            box_marker.pose.position.y = y
-            box_marker.pose.position.z = z
-
-            dx = float(getattr(trk, 'dx', 0.5))
-            dy = float(getattr(trk, 'dy', 0.5))
-            dz = float(getattr(trk, 'dz', 1.5))
-            box_marker.scale = Vector3(x=dx, y=dy, z=dz)
-
-            if threat_level == "EMERGENCY":
-                box_marker.color.r, box_marker.color.g, box_marker.color.b, box_marker.color.a = 1.0, 0.0, 0.0, 0.6
-            elif threat_level == "WARNING":
-                box_marker.color.r, box_marker.color.g, box_marker.color.b, box_marker.color.a = 1.0, 0.9, 0.0, 0.5
-            else:
-                box_marker.color.r, box_marker.color.g, box_marker.color.b, box_marker.color.a = 0.0, 1.0, 0.0, 0.4
-
-            marker_array.markers.append(box_marker)
-
-            # Текст с дистанцией над объектом
-            text_marker = Marker()
-            text_marker.header.stamp = stamp
-            text_marker.header.frame_id = self.frame_id
-            text_marker.ns = "flyguard_labels"
-            text_marker.id = 1000 + i
-            text_marker.type = Marker.TEXT_VIEW_FACING
-            text_marker.action = Marker.ADD
-
-            text_marker.pose.position.x = x
-            text_marker.pose.position.y = y
-            text_marker.pose.position.z = z + dz / 2.0 + 0.3
-            text_marker.scale.z = 0.4
-
-            dist = getattr(trk, 'distance', np.hypot(x, y))
-            text_marker.text = f"D: {dist:.1f}m"
-
-            text_marker.color.r, text_marker.color.g, text_marker.color.b, text_marker.color.a = 1.0, 1.0, 1.0, 1.0
-            marker_array.markers.append(text_marker)
-
-        self.pub_markers.publish(marker_array)
 
 
-def main(args=None):
+def main(args=None) -> None:
     rclpy.init(args=args)
     node = FlyGuardNode()
     try:
         rclpy.spin(node)
-    except KeyboardInterrupt:
+    except (KeyboardInterrupt, ExternalShutdownException):
         pass
     finally:
         node.destroy_node()
+        if rclpy.ok():
             rclpy.shutdown()
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     main()