266 lines
11 KiB
Python
266 lines
11 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.
|
||
if corridor is not None and corridor.n_slices > 0:
|
||
x_sensor = 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 = 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
|
||
)
|