-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathplot.py
126 lines (110 loc) · 4.99 KB
/
plot.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
import numpy as np
from scipy.integrate import cumtrapz
import matplotlib.pyplot as plt
from mpl_toolkits import mplot3d
from math import degrees, atan2, asin, sin, cos
import line_notifier as notifier
from sensor_fusion import R_x, R_y, R_z
def plot_trajectory(df, dt):
# Double integrate accelerations to find positions
x = cumtrapz(cumtrapz(df['EARTH LINEAR ACCELERATION X'], dx=dt), dx=dt)
y = cumtrapz(cumtrapz(df['EARTH LINEAR ACCELERATION Y'], dx=dt), dx=dt)
z = cumtrapz(cumtrapz(df['EARTH LINEAR ACCELERATION Z'], dx=dt), dx=dt)
# Plot 3D Trajectory
fig, ax = plt.subplots()
fig.suptitle('3D Trajectory of IMU', fontsize=20)
ax = plt.axes(projection='3d')
ax.plot3D(x, y, z, c='red', lw=5, label='IMU trajectory')
ax.set_xlabel('X position (m)')
ax.set_ylabel('Y position (m)')
ax.set_zlabel('Z position (m)')
plt.savefig('trajectory.png')
notifier.send_pic("IMU trajectory: ", "trajectory.png")
# plt.show()
def plot_trajectory_attenuated(df, dt):
# Double integrate accelerations to calculate coordinate positions
x = cumtrapz(cumtrapz(df['x_ifft'], dx=dt), dx=dt)
y = cumtrapz(cumtrapz(df['y_ifft'], dx=dt), dx=dt)
z = cumtrapz(cumtrapz(df['z_ifft'], dx=dt), dx=dt)
# Add XYZ axis arrows to indicate phone pose
# Earth 3 axis unit vectors
earth_x = np.array([[1,0,0],] * len(x)).T
earth_y = np.array([[0,1,0],] * len(x)).T
earth_z = np.array([[0,0,1],] * len(x)).T
# Initialize body Vectors
body_x = np.empty(earth_x.shape)
body_y = np.empty(earth_y.shape)
body_z = np.empty(earth_z.shape)
# Perform inverse frame transformations (body frame <-- earth frame)
# body_vectors = (RxRyRz)(earth_vectors)
roll = df['ROLL']
pitch = df['PITCH']
yaw = df['YAW']
for i in range(x.shape[0]):
# use negative angles to reverse rotation
body_x[:,i] = R_x(-roll[i]) @ R_y(-pitch[i]) @ R_z(-yaw[i]) @earth_x[:,i]
body_y[:,i] = R_x(-roll[i]) @ R_y(-pitch[i]) @ R_z(-yaw[i]) @ earth_y[:,i]
body_z[:,i] = R_x(-roll[i]) @ R_y(-pitch[i]) @ R_z(-yaw[i]) @ earth_z[:,i]
# Set length of quiver arrows
distance = np.sqrt(x[-1]**2 + y[-1]**2 + z[-1]**2)
length = 0.05 * distance
# Plot x vectors
# downsampling to every 5th arrow ([::5])
fig6, ax4 = plt.subplots()
fig6.suptitle('IMU trajectory and pose',fontsize=20)
ax4 = plt.axes(projection='3d')
ax4.plot3D(x, y, z, c='red', lw=5, label='Attenuated IMU trajectory')
# plot z vectors
ax4.quiver(x[::5],y[::5],z[::5],
body_z[0][::5],body_z[1][::5],body_z[2][::5],
color='g',label = 'z axis',length = length)
ax4.set_xlabel('X position (m)')
ax4.set_ylabel('Y position (m)')
ax4.set_zlabel('Z position (m)')
plt.savefig('attenuated_trajectory_z_axiz.png')
notifier.send_pic("Attenuated IMU trajectory with Z axis: ", "attenuated_trajectory_z_axiz.png")
plt.show()
def fft(df, dt):
# Discrete Fourier Transform sample frequencies
freq = np.fft.rfftfreq(df['EARTH LINEAR ACCELERATION X'].size, d=dt)
# Compute the Fast Fourier Transform (FFT) of acceleration signals
fft_x = np.fft.rfft(df['EARTH LINEAR ACCELERATION X'])
fft_y = np.fft.rfft(df['EARTH LINEAR ACCELERATION Y'])
fft_z = np.fft.rfft(df['EARTH LINEAR ACCELERATION Z'])
# Plot Frequency spectrum
fig, [ax1, ax2, ax3] = plt.subplots(3, 1, sharex=True, sharey=True)
for f, fx, fy, fz in zip(freq, abs(fft_x), abs(fft_y), abs(fft_z)):
print(f'freq = {f:.2f} ({fx:.3f}, {fy:.3f}, {fz:.3f})')
fig.suptitle('Frequency Spectrum',fontsize=20)
ax1.plot(freq, abs(fft_x), c='r', label='x frequency')
ax1.legend()
ax2.plot(freq, abs(fft_y), c='b', label='y frequency')
ax2.legend()
ax3.plot(freq, abs(fft_z), c='g', label='z frequency')
ax3.legend()
ax3.set_xlabel('Freqeuncy (Hz)')
plt.savefig('frequency_spectrum.png')
notifier.send_pic("Frequency Spectrum: ", "frequency_spectrum.png")
# Attenuate noise in X, Y, Z below 0.1Hz and above 6Hz by 10 dB
atten_x_fft = np.where((freq > 6) | (freq < 0.1), fft_x * 0.1, fft_x)
atten_y_fft = np.where((freq > 6) | (freq < 0.1), fft_y * 0.1, fft_y)
atten_z_fft = np.where((freq > 6) | (freq < 0.1), fft_z * 0.1, fft_z)
# Compute inverse of discrete Fourier Transform
df['x_ifft'] = np.fft.irfft(atten_x_fft, n=df.shape[0])
df['y_ifft'] = np.fft.irfft(atten_y_fft, n=df.shape[0])
df['z_ifft'] = np.fft.irfft(atten_z_fft, n=df.shape[0])
plt.subplot(1, 3, 1)
plt.plot(df['EARTH LINEAR ACCELERATION X'], label='Before')
plt.plot(df['x_ifft'], label='After')
plt.legend()
plt.subplot(1, 3, 2)
plt.plot(df['EARTH LINEAR ACCELERATION Y'], label='Before')
plt.plot(df['y_ifft'], label='After')
plt.legend()
plt.subplot(1, 3, 3)
plt.plot(df['EARTH LINEAR ACCELERATION Z'], label='Before')
plt.plot(df['z_ifft'], label='After')
plt.legend()
plt.savefig('ifft.png')
notifier.send_pic("IFFT: ", "ifft.png")
return df