Files
imu_decode/visualize_3d.py
2026-05-26 04:19:57 +08:00

129 lines
4.2 KiB
Python

"""
IMU 3D 轨迹实时可视化
matplotlib 3D 窗口, 30Hz 刷新, 显示:
- 蓝色轨迹线 (最近 800 点, Tracker 端已限长)
- 当前点红点 + 朝向指示线 (RGB = body XYZ)
- 原点世界坐标系指示
- 等比例坐标轴
"""
import numpy as np
import matplotlib.pyplot as plt
class TrajectoryViewer:
"""3D 轨迹实时显示窗口"""
def __init__(self, title="IMU 3D Odometry"):
self.fig = plt.figure(figsize=(8, 7))
self.fig.canvas.manager.set_window_title(title)
self.ax = self.fig.add_subplot(111, projection='3d')
# 轨迹线 (蓝色)
self.traj_line, = self.ax.plot([], [], [], 'b-', linewidth=1.0, label='Trajectory')
# 当前点 (红色)
self.current_point, = self.ax.plot([], [], [], 'ro', markersize=6, label='Current')
# 朝向指示线 (body XYZ, 用 plot 可原地更新)
self.body_x_line, = self.ax.plot([], [], [], 'r-', linewidth=2.0, alpha=0.9)
self.body_y_line, = self.ax.plot([], [], [], 'g-', linewidth=2.0, alpha=0.9)
self.body_z_line, = self.ax.plot([], [], [], 'b-', linewidth=2.0, alpha=0.9)
# 世界坐标系指示 (固定在原点)
axis_len = 0.3
self.ax.quiver(0, 0, 0, axis_len, 0, 0, color='r', arrow_length_ratio=0.15)
self.ax.quiver(0, 0, 0, 0, axis_len, 0, color='g', arrow_length_ratio=0.15)
self.ax.quiver(0, 0, 0, 0, 0, axis_len, color='b', arrow_length_ratio=0.15)
self._setup_axes()
self._frame = 0
plt.show(block=False)
plt.pause(0.1)
def _setup_axes(self):
self.ax.set_xlabel('X (front)')
self.ax.set_ylabel('Y (left)')
self.ax.set_zlabel('Z (up)')
self.ax.set_title("IMU 3D Trajectory (Z-up)")
self.ax.set_xlim([-1, 1])
self.ax.set_ylim([-1, 1])
self.ax.set_zlim([-1, 1])
try:
self.ax.set_box_aspect([1, 1, 1])
except NotImplementedError:
pass
self.ax.grid(True)
def update(self, history_array, rotation_matrix=None):
"""更新显示
Args:
history_array: Nx3 位置历史 (已由 Tracker 限长为 800)
rotation_matrix: 3x3 body→world 旋转矩阵 (None 则不显示朝向)
"""
n = len(history_array)
if n < 1:
return
x, y, z = history_array[:, 0], history_array[:, 1], history_array[:, 2]
self.traj_line.set_data(x, y)
self.traj_line.set_3d_properties(z)
# 当前点
last = history_array[-1]
self.current_point.set_data([last[0]], [last[1]])
self.current_point.set_3d_properties([last[2]])
# 朝向 (原地更新, 不重建对象)
if rotation_matrix is not None:
self._update_orientation(last, rotation_matrix)
else:
self.body_x_line.set_data([], [])
self.body_x_line.set_3d_properties([])
self.body_y_line.set_data([], [])
self.body_y_line.set_3d_properties([])
self.body_z_line.set_data([], [])
self.body_z_line.set_3d_properties([])
# 坐标轴范围 (每 5 帧更新一次)
self._frame += 1
if self._frame % 8 == 0:
self._auto_scale(x, y, z)
def _update_orientation(self, position, R):
"""更新 body 系朝向指示线 (不重建对象)
Args:
position: (3,) 当前位置
R: 3x3 body→world 旋转矩阵
"""
px, py, pz = position
length = 0.15
for line, col in [(self.body_x_line, 0),
(self.body_y_line, 1),
(self.body_z_line, 2)]:
d = R[:, col] * length
line.set_data([px, px + d[0]], [py, py + d[1]])
line.set_3d_properties([pz, pz + d[2]])
def _auto_scale(self, x, y, z):
all_coords = np.concatenate([x, y, z])
margin = max(np.ptp(all_coords) * 0.2, 0.5)
mid = (all_coords.min() + all_coords.max()) / 2
half = (all_coords.max() - all_coords.min()) / 2 + margin
self.ax.set_xlim([mid - half, mid + half])
self.ax.set_ylim([mid - half, mid + half])
self.ax.set_zlim([mid - half, mid + half])
def close(self):
plt.close(self.fig)