303 lines
13 KiB
Python
303 lines
13 KiB
Python
"""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
|