diff --git a/flyguard/flyguard_ros2_node.py b/flyguard/flyguard_ros2_node.py index 6f68801..9b3b6f5 100644 --- a/flyguard/flyguard_ros2_node.py +++ b/flyguard/flyguard_ros2_node.py @@ -1,216 +1,235 @@ -""" -ROS 2 узел для биоинспирированного пайплайна FlyGuard (ЛЦТ-2026 Кейс 05). -Связывает входные данные 3D-лидара с ядром FlyGuard и публикует детекции/маркеры. +"""ROS 2 adapter for the FlyGuard lidar-obstacle detector. + +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 -try: - from flyguard.pipeline import FlyGuard, Params - from flyguard.mushroom_body import MushroomBody -except ImportError as e: - FlyGuard = None - Params = None - MushroomBody = None +from flyguard.cdr import PointCloud2 as CorePointCloud2 +from flyguard.export import export_frame +from flyguard.mbon_readout import MbonReadout +from flyguard.mushroom_body import MushroomBody +from flyguard.pipeline import FlyGuard, Params class FlyGuardNode(Node): - def __init__(self): - super().__init__('flyguard_node') + """Consume lidar clouds and publish final FlyGuard decisions.""" - # --- Объявление параметров --- - self.declare_parameter('lidar_topic', '/pandar_points') - self.declare_parameter('frame_id', 'pandar') - self.declare_parameter('memory_path', '') - self.declare_parameter('fov_deg', 30.0) + def __init__(self) -> None: + super().__init__("flyguard_node") + self.declare_parameter("lidar_topic", "/pandar_points") + # Empty means preserve the PointCloud2 header frame for downstream use. + 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.frame_id = self.get_parameter('frame_id').get_parameter_value().string_value - memory_path = self.get_parameter('memory_path').get_parameter_value().string_value - fov_deg = self.get_parameter('fov_deg').get_parameter_value().double_value + self.lidar_topic = self.get_parameter("lidar_topic").value + self.frame_id_override = self.get_parameter("frame_id").value + memory_path = self.get_parameter("memory_path").value + mbon_path = self.get_parameter("mbon_path").value + fov_deg = self.get_parameter("fov_deg").value - # --- Инициализация ядра FlyGuard --- - if FlyGuard is not None: - self.get_logger().info("Инициализация ядра FlyGuard...") - memory = MushroomBody.load(memory_path) if memory_path else None - params = Params(fov_deg=fov_deg) - self.fg = FlyGuard(params, memory=memory) - else: - self.get_logger().error("Модуль flyguard.pipeline не найден!") - self.fg = None + memory = self._load_model(memory_path, "memory", MushroomBody.load) + readout = self._load_model(mbon_path, "MBON readout", MbonReadout.load) + self.fg = FlyGuard(Params(fov_deg=float(fov_deg)), memory=memory, + readout=readout) - # --- Подписчики и Публикаторы --- + # 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, - self.lidar_topic, - self.pointcloud_callback, - 10 + RosPointCloud2, self.lidar_topic, self.pointcloud_callback, + qos_profile_sensor_data, + ) + 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 "", + memory_path or "disabled", mbon_path or "disabled") ) - 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) + def _load_model(self, value: str, label: str, loader): + """Load an optional model and fail early for a supplied bad path.""" + 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): - """Обработка каждого кадра лидара (10 Гц)""" - start_time = self.get_clock().now() - - # 1. Читаем точки из PointCloud2 в structured numpy array - # (сохраняет стандартную структуру PointCloud) - pc_data = pc2.read_points(msg, skip_nans=True) - if pc_data.size == 0: + def pointcloud_callback(self, msg: RosPointCloud2) -> None: + self.received_frames += 1 + started = self.get_clock().now() + try: + cloud = self._to_core_cloud(msg) + if cloud.n_points == 0: + self.get_logger().warning("Ignoring an empty PointCloud2 frame") + 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 - if self.fg is None: + # The first calib_frames scans intentionally produce no result. + if result is None: return - res = self.fg.process(pc_data) - if res is None: - return + self.processed_frames += 1 + decision = result.decision + 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 - tracks = getattr(self.fg.cx, 'tracks', []) + self.pub_threat.publish(String(data=threat_level)) + self.publish_detections(exported.boxes, msg.header.stamp, frame_id) + self.publish_rviz_markers(exported.boxes, msg.header.stamp, frame_id) - # 3. Определяем уровень угрозы - if d.detected: - threat_level = "EMERGENCY" if getattr(d, 'emergency', False) else "WARNING" - else: - threat_level = "CLEAR" + elapsed_ms = (self.get_clock().now() - started).nanoseconds / 1e6 + self.get_logger().debug( + "frame=%d processed=%d %.1f ms: %s, %d object(s)" + % (self.received_frames, self.processed_frames, elapsed_ms, + threat_level, len(exported.boxes)) + ) - # 4. Публикация статуса - threat_msg = String() - threat_msg.data = threat_level - self.pub_threat.publish(threat_msg) + 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) - # 5. Публикация маркеров и рамок - stamp = msg.header.stamp - self.publish_detections(tracks, stamp) - self.publish_rviz_markers(tracks, threat_level, d, stamp) + 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) - elapsed_ms = (self.get_clock().now() - start_time).nanoseconds / 1e6 - self.get_logger().debug(f"Кадр за {elapsed_ms:.1f} мс | Статус: {threat_level} | Треков: {len(tracks)}") - - def publish_detections(self, tracks, stamp): - """Формирование vision_msgs/Detection3DArray""" - det_array = Detection3DArray() - 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 + 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: - box_marker.color.r, box_marker.color.g, box_marker.color.b, box_marker.color.a = 0.0, 1.0, 0.0, 0.4 + cube.color.r, cube.color.g, cube.color.b, cube.color.a = 1.0, 0.85, 0.0, 0.65 + message.markers.append(cube) - 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) + label = Marker() + label.header.stamp = stamp + label.header.frame_id = frame_id + label.ns = "flyguard_labels" + label.id = 100000 + int(box.track_id) + label.type = Marker.TEXT_VIEW_FACING + label.action = Marker.ADD + label.pose.position.x, label.pose.position.y = box.x, box.y + label.pose.position.z = box.z + box.dz * 0.5 + 0.35 + label.pose.orientation.w = 1.0 + label.scale.z = 0.40 + label.color.r = label.color.g = label.color.b = label.color.a = 1.0 + label.text = "ID:%d | %.1f m | conf:%.2f" % ( + box.track_id, box.distance_along_track, box.confidence) + message.markers.append(label) + self.pub_markers.publish(message) -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() - rclpy.shutdown() + if rclpy.ok(): + rclpy.shutdown() -if __name__ == '__main__': - main() \ No newline at end of file +if __name__ == "__main__": + main()