-
Notifications
You must be signed in to change notification settings - Fork 5
/
vis.py
112 lines (91 loc) · 4.55 KB
/
vis.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
import torch
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from utils import *
import io
from PIL import Image
from liegroups.torch import SE3
def canvas_to_array(fig):
#Convert matplotlib figure to a C X H X W numpy array for TensorboardX
canvas = fig.canvas
canvas.draw()
width, height = fig.get_size_inches() * fig.get_dpi()
image_np = np.fromstring(canvas.tostring_rgb(), dtype='uint8').reshape(height.astype(np.uint32), width.astype(np.uint32), 3)
#PIL expects CXHXW
return np.rollaxis(image_np, 2)
def _plot_sigma(x, y, y_mean, y_sigma, y_sigma_2, label, ax):
ax.fill_between(x, y_mean-3*y_sigma, y_mean+3*y_sigma, alpha=0.5, label='$\pm 3\sigma$ Total')
ax.fill_between(x, y_mean - 3 * y_sigma_2, y_mean + 3 * y_sigma_2, alpha=0.5, color='red', label='$\pm 3\sigma$ Direct')
ax.scatter(x, y, s=0.5, c='black')
ax.set_ylabel(label)
return
def _plot_sigma_with_gt(x, y_est, y_gt, y_sigma, y_sigma_2, label, ax, y_lim=None):
ax.fill_between(x, y_est-3*y_sigma, y_est+3*y_sigma, alpha=0.5, label='$\pm 3\sigma$ Total')
ax.fill_between(x, y_est - 3 * y_sigma_2, y_est + 3 * y_sigma_2, alpha=0.5, color='red', label='$\pm 3\sigma$ Direct')
ax.scatter(x, y_est, s=0.5, c='green')
ax.scatter(x, y_gt, s=0.5, c='black')
ax.set_ylabel(label)
if y_lim is not None:
ax.set_ylim(y_lim)
return
def plot_errors_with_sigmas(q_gt, q_est, R_est, R_direct_est, filename='sigma_plot.pdf'):
fig, ax = plt.subplots(3, 1, sharex='col', sharey='row')
x_labels = np.arange(0, q_gt.shape[0])
phi_errs = quat_log_diff(q_est, q_gt).numpy()
R_est = R_est.numpy()
R_direct_est = R_direct_est.numpy()
_plot_sigma(x_labels, phi_errs[:, 0], 0., np.sqrt(R_est[:,0,0].flatten()), np.sqrt(R_direct_est[:,0,0].flatten()), '$\Theta_1$ err', ax[0])
_plot_sigma(x_labels, phi_errs[:, 1], 0., np.sqrt(R_est[:,1,1].flatten()), np.sqrt(R_direct_est[:,1,1].flatten()), '$\Theta_2$ err', ax[1])
_plot_sigma(x_labels, phi_errs[:, 2], 0., np.sqrt(R_est[:,2,2].flatten()), np.sqrt(R_direct_est[:,2,2].flatten()), '$\Theta_3$ err', ax[2])
ax[2].legend()
#image_array = canvas_to_array(fig)
fig.savefig(filename, bbox_inches='tight')
plt.close(fig)
def plot_abs_with_sigmas(q_gt, q_est, R_est, R_direct_est, filename='sigma_plot.pdf'):
fig, ax = plt.subplots(3, 1, sharex='col', sharey='row')
x_labels = np.arange(0, q_gt.shape[0])
phi_est = quat_log(q_est).numpy()
phi_gt = quat_log(q_gt).numpy()
R_est = R_est.numpy()
R_direct_est = R_direct_est.numpy()
_plot_sigma_with_gt(x_labels, phi_est[:, 0], phi_gt[:, 0], np.sqrt(R_est[:,0,0].flatten()), np.sqrt(R_direct_est[:,0,0].flatten()), '$\Theta_1$', ax[0])
_plot_sigma_with_gt(x_labels, phi_est[:, 1], phi_gt[:, 1], np.sqrt(R_est[:,1,1].flatten()), np.sqrt(R_direct_est[:,1,1].flatten()), '$\Theta_2$', ax[1])
_plot_sigma_with_gt(x_labels, phi_est[:, 2], phi_gt[:, 2], np.sqrt(R_est[:,2,2].flatten()), np.sqrt(R_direct_est[:,2,2].flatten()), '$\Theta_3$', ax[2])
ax[2].legend()
#image_array = canvas_to_array(fig)
fig.savefig(filename, bbox_inches='tight')
plt.close(fig)
def plot_nees(q_gt, q_est, R_est, filename='nees_plot.pdf'):
fig, ax = plt.subplots(1, 1, sharex='col', sharey='row')
xlabels = np.arange(0, q_gt.shape[0])
residuals = quat_log_diff(q_est, q_gt).unsqueeze(2)
nees = (1./3.) * residuals.transpose(1, 2).bmm(R_est.inverse()).bmm(residuals)
nees.squeeze_()
ax.plot(xlabels, nees.numpy(), label='nees')
ax.legend()
#image_array = canvas_to_array(fig)
fig.savefig(filename, bbox_inches='tight')
plt.close(fig)
def plot_3D(tau_gt, tau_odom, tau_est, l_true=None):
fig = plt.figure()
ax = Axes3D(fig)
# ax.set_zlim([0, 2.5])
# ax.set_ylim([0, 5])
# ax.set_xlim([1, 3.5])
if l_true is not None:
ax.scatter(l_true[0,:], l_true[1,:], l_true[2,:], label='Landmarks')
r_odom = SE3.exp(torch.from_numpy(tau_odom)).inv().trans.numpy()
r_est = SE3.exp(torch.from_numpy(tau_est)).inv().trans.numpy()
r_gt = SE3.exp(torch.from_numpy(tau_gt)).inv().trans.numpy()
ax.plot3D(r_odom[:,0], r_odom[:, 1],r_odom[:, 2], '-r', label='IMU Integration')
ax.plot3D(r_est[:,0], r_est[:, 1], r_est[:, 2], '-g', label='Corrected')
ax.plot3D(r_gt[:,0], r_gt[:, 1], r_gt[:, 2], '--k', label='Ground Truth')
ax.legend()
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_zlabel('z')
ax.view_init(elev=90., azim=0.)
image_array = canvas_to_array(fig)
plt.close(fig)
return image_array