Add linear velocity calculation and visualization to IMU data processing
- Implemented Euler to rotation matrix conversion for IMU data. - Added function to compute linear velocity from accelerometer data with gravity compensation. - Enhanced CSV and PNG saving functionality to automatically create directories. - Updated statistics printing to include linear acceleration and velocity metrics. - Expanded plotting functionality to visualize linear acceleration and velocity alongside existing data. - Added new PNG files for IMU quality visualization.
This commit is contained in:
Brunsmeier
208
visualise.py
208
visualise.py
@ -184,21 +184,146 @@ def write_csv(samples, path):
|
||||
writer = csv.DictWriter(f, fieldnames=columns)
|
||||
writer.writeheader()
|
||||
writer.writerows(samples)
|
||||
|
||||
def euler_to_rot_matrix(roll_deg, pitch_deg, yaw_deg):
|
||||
"""
|
||||
Convert IMU Euler angles to body->world rotation matrix.
|
||||
Assumption:
|
||||
angle_x = roll
|
||||
angle_y = pitch
|
||||
angle_z = yaw
|
||||
Rotation order:
|
||||
R = Rz(yaw) @ Ry(pitch) @ Rx(roll)
|
||||
|
||||
注意：这个顺序要和 IWT603 官方定义确认。
|
||||
如果发现补偿后静止时 linear_acc 不接近 0.需要调整欧拉角顺序或正负号。
|
||||
"""
|
||||
roll = np.deg2rad(roll_deg)
|
||||
pitch = np.deg2rad(pitch_deg)
|
||||
yaw = np.deg2rad(yaw_deg)
|
||||
|
||||
cr, sr = np.cos(roll), np.sin(roll)
|
||||
cp, sp = np.cos(pitch), np.sin(pitch)
|
||||
cy, sy = np.cos(yaw), np.sin(yaw)
|
||||
|
||||
Rx = np.array([
|
||||
[1, 0, 0],
|
||||
[0, cr, -sr],
|
||||
[0, sr, cr],
|
||||
])
|
||||
|
||||
Ry = np.array([
|
||||
[cp, 0, sp],
|
||||
[0, 1, 0],
|
||||
[-sp, 0, cp],
|
||||
])
|
||||
|
||||
Rz = np.array([
|
||||
[cy, -sy, 0],
|
||||
[sy, cy, 0],
|
||||
[0, 0, 1],
|
||||
])
|
||||
|
||||
return Rz @ Ry @ Rx
|
||||
|
||||
|
||||
def add_linear_velocity(samples, gravity_sign=1.0, acc_deadband=0.15, vel_decay=0.995, bias_seconds=1.0):
|
||||
if not samples:
|
||||
return samples
|
||||
|
||||
# 先计算每一帧的重力补偿线加速度
|
||||
raw_linear_acc_list = []
|
||||
|
||||
for sample in samples:
|
||||
acc_body_g = np.array([
|
||||
sample["acc_x_g"],
|
||||
sample["acc_y_g"],
|
||||
sample["acc_z_g"],
|
||||
], dtype=float)
|
||||
|
||||
R = euler_to_rot_matrix(
|
||||
sample["angle_x_deg"],
|
||||
sample["angle_y_deg"],
|
||||
sample["angle_z_deg"],
|
||||
)
|
||||
|
||||
acc_world_g = R @ acc_body_g
|
||||
linear_acc_world_g = acc_world_g - np.array([0.0, 0.0, gravity_sign])
|
||||
linear_acc_world = linear_acc_world_g * 9.81
|
||||
|
||||
raw_linear_acc_list.append(linear_acc_world)
|
||||
|
||||
raw_linear_acc_arr = np.array(raw_linear_acc_list)
|
||||
|
||||
# 用前 bias_seconds 秒估计静止零偏
|
||||
t0 = samples[0]["t_s"]
|
||||
bias_indices = [
|
||||
i for i, s in enumerate(samples)
|
||||
if s["t_s"] - t0 <= bias_seconds
|
||||
]
|
||||
|
||||
if bias_indices:
|
||||
acc_bias = np.mean(raw_linear_acc_arr[bias_indices], axis=0)
|
||||
else:
|
||||
acc_bias = np.zeros(3)
|
||||
|
||||
print("Estimated linear acc bias m/s^2:", acc_bias)
|
||||
|
||||
v = np.zeros(3, dtype=float)
|
||||
last_t = samples[0]["t_s"]
|
||||
|
||||
for i, sample in enumerate(samples):
|
||||
t = sample["t_s"]
|
||||
dt = t - last_t
|
||||
last_t = t
|
||||
|
||||
if dt <= 0 or dt > 0.2:
|
||||
dt = 0.0
|
||||
|
||||
linear_acc_world = raw_linear_acc_arr[i] - acc_bias
|
||||
|
||||
if np.linalg.norm(linear_acc_world) < acc_deadband:
|
||||
linear_acc_world[:] = 0.0
|
||||
|
||||
v = v * vel_decay + linear_acc_world * dt
|
||||
|
||||
sample["lin_acc_x_ms2"] = linear_acc_world[0]
|
||||
sample["lin_acc_y_ms2"] = linear_acc_world[1]
|
||||
sample["lin_acc_z_ms2"] = linear_acc_world[2]
|
||||
sample["vel_x_ms"] = v[0]
|
||||
sample["vel_y_ms"] = v[1]
|
||||
sample["vel_z_ms"] = v[2]
|
||||
|
||||
return samples
|
||||
def unique_path(path):
|
||||
"""
|
||||
If path already exists, append timestamp before extension.
|
||||
Example:
|
||||
imu_quality.png -> imu_quality_20260701_123045.png
|
||||
Automatically save CSV into ./csv/
|
||||
Automatically save PNG into ./png/
|
||||
Create folders if they do not exist.
|
||||
"""
|
||||
|
||||
if not path:
|
||||
return path
|
||||
|
||||
if not os.path.exists(path):
|
||||
return path
|
||||
root, ext = os.path.splitext(os.path.basename(path))
|
||||
ext = ext.lower()
|
||||
|
||||
if ext == ".csv":
|
||||
folder = "csv"
|
||||
elif ext == ".png":
|
||||
folder = "png"
|
||||
else:
|
||||
folder = "output"
|
||||
|
||||
os.makedirs(folder, exist_ok=True)
|
||||
|
||||
new_path = os.path.join(folder, root + ext)
|
||||
|
||||
if not os.path.exists(new_path):
|
||||
return new_path
|
||||
|
||||
root, ext = os.path.splitext(path)
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
return "{}_{}{}".format(root, timestamp, ext)
|
||||
return os.path.join(folder, "{}_{}{}".format(root, timestamp, ext))
|
||||
|
||||
def print_stats(samples):
|
||||
if not samples:
|
||||
@ -238,6 +363,35 @@ def print_stats(samples):
|
||||
float(np.max(freq)),
|
||||
)
|
||||
)
|
||||
if "lin_acc_x_ms2" in samples[0]:
|
||||
print("\nLINEAR ACC m/s^2")
|
||||
for field in ("lin_acc_x_ms2", "lin_acc_y_ms2", "lin_acc_z_ms2"):
|
||||
values = np.array([sample[field] for sample in samples], dtype=float)
|
||||
diffs = np.diff(values)
|
||||
diff_std = float(np.std(diffs)) if len(diffs) else 0.0
|
||||
print(
|
||||
" {:16s} mean={: .6f} std={:.6f} ptp={:.6f} diff_std={:.6f}".format(
|
||||
field,
|
||||
float(np.mean(values)),
|
||||
float(np.std(values)),
|
||||
float(np.ptp(values)),
|
||||
diff_std,
|
||||
)
|
||||
)
|
||||
|
||||
if "vel_x_ms" in samples[0]:
|
||||
print("\nVELOCITY m/s")
|
||||
for field in ("vel_x_ms", "vel_y_ms", "vel_z_ms"):
|
||||
values = np.array([sample[field] for sample in samples], dtype=float)
|
||||
print(
|
||||
" {:16s} mean={: .6f} std={:.6f} ptp={:.6f} final={: .6f}".format(
|
||||
field,
|
||||
float(np.mean(values)),
|
||||
float(np.std(values)),
|
||||
float(np.ptp(values)),
|
||||
float(values[-1]),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def plot_samples(samples, output=None, show=True):
|
||||
@ -245,24 +399,42 @@ def plot_samples(samples, output=None, show=True):
|
||||
raise ValueError("No samples to plot.")
|
||||
|
||||
t = np.array([sample["t_s"] for sample in samples], dtype=float)
|
||||
fig, axes = plt.subplots(4, 1, sharex=True, figsize=(13, 9))
|
||||
fig, axes = plt.subplots(6, 1, sharex=True, figsize=(13, 13))
|
||||
|
||||
plot_group(
|
||||
axes[0],
|
||||
t,
|
||||
samples,
|
||||
("acc_x_g", "acc_y_g", "acc_z_g"),
|
||||
"Acceleration (g)",
|
||||
"Raw Acc (g)",
|
||||
)
|
||||
|
||||
plot_group(
|
||||
axes[1],
|
||||
t,
|
||||
samples,
|
||||
("lin_acc_x_ms2", "lin_acc_y_ms2", "lin_acc_z_ms2"),
|
||||
"Linear Acc (m/s²)",
|
||||
)
|
||||
|
||||
plot_group(
|
||||
axes[2],
|
||||
t,
|
||||
samples,
|
||||
("vel_x_ms", "vel_y_ms", "vel_z_ms"),
|
||||
"Velocity (m/s)",
|
||||
)
|
||||
|
||||
plot_group(
|
||||
axes[3],
|
||||
t,
|
||||
samples,
|
||||
("gyro_x_dps", "gyro_y_dps", "gyro_z_dps"),
|
||||
"Gyro (deg/s)",
|
||||
)
|
||||
|
||||
plot_group(
|
||||
axes[2],
|
||||
axes[4],
|
||||
t,
|
||||
samples,
|
||||
("angle_x_deg", "angle_y_deg", "angle_z_deg"),
|
||||
@ -270,11 +442,11 @@ def plot_samples(samples, output=None, show=True):
|
||||
)
|
||||
|
||||
freq = np.array([sample["freq_hz"] for sample in samples], dtype=float)
|
||||
axes[3].plot(t, freq, label="freq_hz", color="tab:purple", linewidth=1.0)
|
||||
axes[3].set_ylabel("Hz")
|
||||
axes[3].set_xlabel("Time (s)")
|
||||
axes[3].grid(True, alpha=0.3)
|
||||
axes[3].legend(loc="upper right")
|
||||
axes[5].plot(t, freq, label="freq_hz", linewidth=1.0)
|
||||
axes[5].set_ylabel("Hz")
|
||||
axes[5].set_xlabel("Time (s)")
|
||||
axes[5].grid(True, alpha=0.3)
|
||||
axes[5].legend(loc="upper right")
|
||||
|
||||
fig.suptitle("IMU Data Quality Overview")
|
||||
fig.tight_layout()
|
||||
@ -328,6 +500,12 @@ def main():
|
||||
args.max_gap_ms,
|
||||
)
|
||||
|
||||
# 关键：没有样本就直接退出，避免后面 CSV/plot 误报或报错
|
||||
if not samples:
|
||||
print("No samples parsed. Check input mode, serial port, baudrate, or log file format.")
|
||||
return
|
||||
|
||||
samples = add_linear_velocity(samples, gravity_sign=1.0)
|
||||
print_stats(samples)
|
||||
|
||||
if args.csv:
|
||||
|
||||
Reference in New Issue
Block a user