Files
stabilization/visualization.py

98 lines
2.7 KiB
Python

import serial
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
from matplotlib.animation import FuncAnimation
import re
# Serial settings (change COM port and baud rate as needed)
SERIAL_PORT = 'COM3' # e.g., 'COM3' on Windows or '/dev/ttyUSB0' on Linux/Mac
BAUD_RATE = 9600
# Open serial connection
ser = serial.Serial(SERIAL_PORT, BAUD_RATE, timeout=1)
# Parse "roll: xx, pitch: xx, yaw: xx"
def parse_data(line):
match = re.search(r'roll:\s*(-?\d+\.?\d*),\s*pitch:\s*(-?\d+\.?\d*),\s*yaw:\s*(-?\d+\.?\d*)', line)
if match:
roll = float(match.group(1))
pitch = float(match.group(2))
yaw = float(match.group(3))
return roll, pitch, yaw
return None
def euler_to_rotation_matrix(roll, pitch, yaw):
roll = np.radians(roll)
pitch = np.radians(pitch)
yaw = np.radians(yaw)
Rx = np.array([[1, 0, 0],
[0, np.cos(roll), -np.sin(roll)],
[0, np.sin(roll), np.cos(roll)]])
Ry = np.array([[np.cos(pitch), 0, np.sin(pitch)],
[0, 1, 0],
[-np.sin(pitch), 0, np.cos(pitch)]])
Rz = np.array([[np.cos(yaw), -np.sin(yaw), 0],
[np.sin(yaw), np.cos(yaw), 0],
[0, 0, 1]])
return Rz @ Ry @ Rx
def create_box():
l, w, h = 1, 0.5, 0.2
return np.array([
[-l/2, -w/20, -h/2],
[ l/2, -w/2, -h/2],
[ l/2, w/2, -h/2],
[-l/2, w/20, -h/2],
[-l/2, -w/20, h/2],
[ l/2, -w/2, h/2],
[ l/2, w/2, h/2],
[-l/2, w/20, h/2],
])
def get_faces(vertices):
return [
[vertices[j] for j in [0,1,2,3]],
[vertices[j] for j in [4,5,6,7]],
[vertices[j] for j in [0,1,5,4]],
[vertices[j] for j in [2,3,7,6]],
[vertices[j] for j in [1,2,6,5]],
[vertices[j] for j in [3,0,4,7]],
]
# Create plot
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
box = create_box()
collection = Poly3DCollection([], facecolors='skyblue', edgecolors='black', alpha=0.8)
ax.add_collection3d(collection)
ax.set_xlim([-1, 1])
ax.set_ylim([-1, 1])
ax.set_zlim([-1, 1])
ax.set_xlabel("X")
ax.set_ylabel("Y")
ax.set_zlabel("Z")
ax.view_init(elev=30, azim=45)
# Animation update function
def update(frame):
line = ser.readline().decode('utf-8').strip()
result = parse_data(line)
if result:
roll, pitch, yaw = result
R = euler_to_rotation_matrix(roll, pitch, yaw)
rotated_box = box @ R.T
faces = get_faces(rotated_box)
collection.set_verts(faces)
ax.set_title(f"Roll={roll:.1f}°, Pitch={pitch:.1f}°, Yaw={yaw:.1f}°")
ani = FuncAnimation(fig, update, interval=50)
plt.tight_layout()
plt.show()