"""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: вверх). # # Боковое смещение трека отсчитано от ОСИ ПУТИ, а не от оси сенсора # (TrackFrame.lateral), поэтому в кривой к нему прибавляется положение # оси на этой дальности: при радиусе 1300 м это 1.2 м на 55 м и 8.6 м # на 150 м — без поправки рамка рисовалась в стене. Направление оси # пути в кадре — (наклон, −1), и длинная ось рамки (её локальная y) # совпадает с ним при повороте на +atan(наклон). Проверка обоих — # test_export_box_follows_a_curved_track. # # Если трек знает своё смещение в системе лидара (`sensor_x`), берётся # оно: габарит — объединение прямого и изогнутого, и у предмета, # попавшего в прямой, `u` отсчитан не от кривой, а от оси лидара. sx = float(getattr(obj, "sensor_x", float("nan"))) if corridor is not None and corridor.n_slices > 0: x_sensor = (sx if math.isfinite(sx) else float(u + corridor.centre(np.array([d], np.float32))[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: x_sensor = sx if math.isfinite(sx) else float(u) yaw = 0.0 y_sensor = float(-d) z_sensor = float(h + a * d + b * x_sensor + c) # Уровень опасности для конкретного объекта 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 = 0.50 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 ) def gauge_outline(boxes, centre=None, step: float = 2.0, frame_every: float = 20.0) -> list[tuple[float, float, float]]: """Отрезки контура габарита (пары точек для LINE_LIST) в координатах облака обзора. `boxes` — секции `(полуширина, низ, верх, от, до)`; `centre(d)` — ось пути (`Corridor.centre`), None — прямой короб вдоль x = 0. Облако обзора в кривой не выпрямлено, поэтому габарит вдоль изогнутой оси рисуется изогнутым — там, где узел его и проверяет. Вперёд = −y, вправо = +x, вверх = +z над головкой рельса, как у облака обзора. Рамки поперёк — через `frame_every`. """ pts: list[tuple[float, float, float]] = [] for hw, z0, z1, d0, d1 in boxes: if d1 <= d0: continue n = 1 if centre is None else max(int(np.ceil((d1 - d0) / step)), 1) ds = np.linspace(d0, d1, n + 1) cx = np.zeros_like(ds) if centre is None else np.asarray(centre(ds), np.float64) corners = [(-hw, z0), (hw, z0), (hw, z1), (-hw, z1)] for x, z in corners: # вдоль пути for a in range(n): pts.append((float(cx[a] + x), float(-ds[a]), float(z))) pts.append((float(cx[a + 1] + x), float(-ds[a + 1]), float(z))) fd = np.arange(d0, d1 + 1e-6, frame_every) # поперёк fc = np.zeros_like(fd) if centre is None else np.asarray(centre(fd), np.float64) for dd, c in zip(fd, fc): for (xa, za), (xb, zb) in zip(corners, corners[1:] + corners[:1]): pts.append((float(c + xa), float(-dd), float(za))) pts.append((float(c + xb), float(-dd), float(zb))) return pts