Brainrot_Muxa/flyguard/export.py
Zhirik1337 062bdfd68c аугментации бенчмарка, зонный пол в колее, фильтр платформы, экспорт RViz
- tools/make_benchmark.py: аугментации спавна d_start ∈ [40, 200] м, боковой дрейф
  v_lat, шум продольной координаты δs (устранение инверсии s_std), сценарии стоянки;
- flyguard/lobula.py: зонный пол h_lo_core = 0.16 м в межрельсовой колее,
  динамическое расширение габарита в кривых W_eff(d) по Corridor.sigma(d),
  отсечение плоскости настила платформы;
- flyguard/central_complex.py: поддержка лежащих препятствий в колее без
  штрафа за вытянутость формы;
- flyguard/synth.py: добавлен класс «человек_лежа» (1.8×0.5×0.3 м);
- flyguard/descending.py: дальний мягкий канал предупреждения на дистанциях >90 м;
- flyguard/export.py: экспорт детекций в 3D BBox, уровни угрозы, маркеры RViz MarkerArray;
- tests/test_pipeline.py, tests/run_tests.py: 38 юнит-тестов и автономный раннер.
2026-09-22 21:25:39 +03:00

258 lines
10 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: вверх)
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
)