"""ROS2 и визуализационный экспорт решений FlyGuard. Преобразует внутренние результаты конвейера (Track, Decision, RailPlane, Corridor) в стандартизованные 3D Bounding Boxes, вектор угроз и структуры MarkerArray для RViz. Работает автономно на чистом Python + NumPy, не требуя обязательной установки библиотек rclpy / ros2 на стенде валидации. При наличии ROS2 может конвертировать напрямую в сообщения visualization_msgs и vision_msgs. """ from __future__ import annotations from dataclasses import asdict, dataclass, field from enum import IntEnum import math import numpy as np from .descending import Decision, DetectedObject from .geometry import Corridor, RailPlane class ThreatLevel(IntEnum): """Уровень опасности для системы автоведения поезда.""" CLEAR = 0 # Путь свободен WARNING = 1 # Заблаговременное предупреждение (DNp02/DNp11 soft-warning) EMERGENCY = 2 # Экстренное торможение (DNp01 / Giant Fiber) @dataclass class BoundingBox3D: """3D ориентированный параллелепипед в координатах сенсора лидара.""" # Центр бокса в системе сенсора (x: вправо, y: вперёд (-d), z: вверх) x: float y: float z: float # Размеры бокса (м) dx: float # ширина поперёк пути dy: float # протяжённость вдоль пути dz: float # высота # Ориентация (рыскание относительно оси лидара, рад) yaw: float # Метрики движения и трекинга distance_along_track: float lateral_offset: float height_above_rail: float confidence: float novelty: float ttc: float track_id: int threat_level: ThreatLevel def to_dict(self) -> dict: d = asdict(self) d["threat_level"] = self.threat_level.name return d @dataclass class ExportResult: """Полный экспортный пакет за один кадр.""" stamp: float threat_level: ThreatLevel nearest_distance: float ttc: float stopping_distance: float speed_mps: float speed_kmh: float boxes: list[BoundingBox3D] = field(default_factory=list) corridor_points_xyz: list[tuple[float, float, float]] = field(default_factory=list) def to_dict(self) -> dict: return { "stamp": self.stamp, "threat_level": self.threat_level.name, "threat_code": int(self.threat_level), "nearest_distance": self.nearest_distance, "ttc": self.ttc, "stopping_distance": self.stopping_distance, "speed_kmh": self.speed_kmh, "n_objects": len(self.boxes), "boxes": [b.to_dict() for b in self.boxes], "corridor_points": self.corridor_points_xyz, } def to_rviz_markers(self, frame_id: str = "hesai_pandar") -> list[dict]: """Генерация словарей, готовых для преобразования в visualization_msgs/Marker.""" markers = [] now_sec = int(self.stamp) now_nanosec = int((self.stamp - now_sec) * 1e9) # 1. Линия коридора пути (LINE_STRIP, type 4) if self.corridor_points_xyz: markers.append({ "header": {"frame_id": frame_id, "sec": now_sec, "nanosec": now_nanosec}, "ns": "flyguard_corridor", "id": 0, "type": 4, # LINE_STRIP "action": 0, # ADD "scale": {"x": 0.12}, "color": {"r": 0.2, "g": 0.8, "b": 1.0, "a": 0.8}, "points": [{"x": p[0], "y": p[1], "z": p[2]} for p in self.corridor_points_xyz] }) # 2. Bounding boxes объектов (CUBE, type 1) и надписи (TEXT, type 9) for i, box in enumerate(self.boxes): if box.threat_level == ThreatLevel.EMERGENCY: color = {"r": 1.0, "g": 0.1, "b": 0.1, "a": 0.75} # Красный elif box.threat_level == ThreatLevel.WARNING: color = {"r": 1.0, "g": 0.85, "b": 0.0, "a": 0.65} # Жёлтый else: color = {"r": 0.2, "g": 0.8, "b": 0.2, "a": 0.50} # Зелёный # Кватернион поворота вокруг оси Z (yaw) cy = math.cos(box.yaw * 0.5) sy = math.sin(box.yaw * 0.5) # CUBE маркер markers.append({ "header": {"frame_id": frame_id, "sec": now_sec, "nanosec": now_nanosec}, "ns": "flyguard_bboxes", "id": box.track_id * 2, "type": 1, # CUBE "action": 0, "pose": { "position": {"x": box.x, "y": box.y, "z": box.z}, "orientation": {"x": 0.0, "y": 0.0, "z": sy, "w": cy} }, "scale": {"x": max(box.dx, 0.2), "y": max(box.dy, 0.2), "z": max(box.dz, 0.2)}, "color": color }) # TEXT_VIEW_FACING над объектом ttc_str = f"{box.ttc:.1f}s" if math.isfinite(box.ttc) else "inf" label = f"ID:{box.track_id} | {box.distance_along_track:.1f}m | TTC:{ttc_str}" markers.append({ "header": {"frame_id": frame_id, "sec": now_sec, "nanosec": now_nanosec}, "ns": "flyguard_labels", "id": box.track_id * 2 + 1, "type": 9, # TEXT_VIEW_FACING "action": 0, "pose": { "position": {"x": box.x, "y": box.y, "z": box.z + box.dz * 0.5 + 0.35}, "orientation": {"x": 0.0, "y": 0.0, "z": 0.0, "w": 1.0} }, "scale": {"z": 0.40}, # Высота шрифта "color": {"r": 1.0, "g": 1.0, "b": 1.0, "a": 0.95}, "text": label }) return markers def export_frame(decision: Decision, plane: RailPlane | None, corridor: Corridor | None, stamp: float = 0.0) -> ExportResult: """Сконвертировать решение FlyGuard в экспортный формат. Parameters ---------- decision : Decision Итоговый вердикт системы за кадр. plane : RailPlane, optional Плоскость головок рельсов (z = a·d + b·u + c). corridor : Corridor, optional Осевая линия тоннеля (парабола u(d)). stamp : float Временная метка кадра. """ if decision.emergency: threat = ThreatLevel.EMERGENCY elif decision.detected: threat = ThreatLevel.WARNING else: threat = ThreatLevel.CLEAR # Параметры плоскости пути: z = a*d + b*u + c a = plane.a if plane is not None else 0.0 b = plane.b if plane is not None else 0.0 c = plane.c if plane is not None else -1.80 # высота лидара над рельсами ~1.8м boxes: list[BoundingBox3D] = [] for obj in decision.objects: d = obj.distance u = obj.lateral h = obj.height # Координаты в системе сенсора (x: вправо, y: вперёд (-d), z: вверх) x_sensor = float(u) y_sensor = float(-d) z_sensor = float(h + a * d + b * u + c) # Касательная к оси коридора: dyaw / dd if corridor is not None and corridor.n_slices > 0: c0, c1, c2 = corridor.coef dm = max(corridor.d_max_seen, 1.0) d_in = min(d, dm) slope = c1 + 2.0 * c2 * d_in yaw = float(-math.atan(slope)) else: yaw = 0.0 # Уровень опасности для конкретного объекта if decision.emergency and d <= max(decision.stopping_distance, 25.0): obj_threat = ThreatLevel.EMERGENCY elif decision.detected: obj_threat = ThreatLevel.WARNING else: obj_threat = ThreatLevel.CLEAR # Размеры: dx поперёк пути, dy вдоль пути, dz по вертикали dx = float(max(obj.width, 0.35)) dy = float(max(obj.size_v if hasattr(obj, "depth") else 0.50, 0.40)) dz = float(max(obj.size_v, 0.40)) boxes.append(BoundingBox3D( x=x_sensor, y=y_sensor, z=z_sensor, dx=dx, dy=dy, dz=dz, yaw=yaw, distance_along_track=float(d), lateral_offset=float(u), height_above_rail=float(h), confidence=float(obj.confidence), novelty=float(obj.novelty), ttc=float(obj.ttc), track_id=int(obj.track_id), threat_level=obj_threat, )) # Траектория коридора вперед (на 10..180 м) corridor_pts: list[tuple[float, float, float]] = [] if corridor is not None: ds_sample = np.linspace(10.0, min(max(corridor.d_max_seen, 50.0), 180.0), 25) us_sample = corridor.centre(ds_sample) for ds_i, us_i in zip(ds_sample, us_sample): xs = float(us_i) ys = float(-ds_i) zs = float(a * ds_i + b * us_i + c + 0.1) # чуть над рельсом corridor_pts.append((xs, ys, zs)) speed = decision.speed return ExportResult( stamp=stamp, threat_level=threat, nearest_distance=float(decision.distance), ttc=float(decision.ttc), stopping_distance=float(decision.stopping_distance), speed_mps=float(speed), speed_kmh=float(speed * 3.6), boxes=boxes, corridor_points_xyz=corridor_pts )