1+ #!/usr/bin/env python
2+ # Author: Kota Kondo
3+ # example use: python plot_animation.py /media/kota/T7/frame/real_time_fastsam_exp/run.bag NX04 /media/kota/T7/frame/real_time_fastsam_exp
4+
5+ import os
6+ import argparse
7+ import rosbag
8+ import rospy
9+ from scipy .spatial .transform import Rotation as R
10+ import numpy as np
11+ import matplotlib .pyplot as plt
12+ from itertools import permutations , combinations
13+ from tf_bag import BagTfTransformer
14+ from scipy .interpolate import make_interp_spline
15+
16+ def smooth_data (x , y ):
17+ X_Y_Spline = make_interp_spline (x , y )
18+ X_ = np .linspace (min (x ), max (x ), 10 )
19+ Y_ = X_Y_Spline (X_ )
20+ return X_ , Y_
21+
22+ def get_transformation (bag , veh_pair ):
23+
24+ # ego name
25+ ego_name = veh_pair [0 ]
26+ # other name
27+ other_name = veh_pair [1 ]
28+
29+ # topic name
30+ tpn_frame_align = f'/{ ego_name }
31+
32+ # get euler angles from data_frame_align (transofrmation matrix from ego to other)
33+ euler_frame_align = []
34+ offsets_frame_align = []
35+ t_frame_align = []
36+ for topic , msg , t in bag .read_messages (topics = [tpn_frame_align ]):
37+ if topic == tpn_frame_align and msg .frame_src == other_name and msg .frame_dest == ego_name :
38+ t_frame_align .append (t .to_sec ())
39+ rotation_matrix = [[msg .transform [0 ], msg .transform [1 ], msg .transform [2 ]], [msg .transform [4 ], msg .transform [5 ], msg .transform [6 ]], [msg .transform [8 ], msg .transform [9 ], msg .transform [10 ]]]
40+ r = R .from_matrix (rotation_matrix )
41+ euler_frame_align .append (r .as_euler ('xyz' , degrees = True ))
42+ offsets_frame_align .append ([msg .transform [3 ], msg .transform [7 ], msg .transform [11 ]])
43+
44+ return euler_frame_align , offsets_frame_align , t_frame_align
45+
46+ def get_ground_truth_transformation (bag , frame_name , t_frame_align ):
47+
48+ # get bag's start and end time
49+ bag_start_time = bag .get_start_time ()
50+ bag_end_time = bag .get_end_time ()
51+ time_discreted = np .linspace (bag_start_time , bag_end_time , len (t_frame_align )* 5 )
52+
53+ # bag transformer
54+ bag_transformer = BagTfTransformer (bag )
55+
56+ # get euler_actual_drift and offsets_actual_drift
57+ euler_actual_drift = []
58+ offsets_actual_drift = []
59+ for t in time_discreted :
60+ translation , quaternion = bag_transformer .lookupTransform ("world" , frame_name , rospy .Time .from_sec (t ))
61+ r = R .from_quat (quaternion )
62+ euler_actual_drift .append (r .as_euler ('xyz' , degrees = True ))
63+ offsets_actual_drift .append (translation )
64+
65+ return euler_actual_drift , offsets_actual_drift , time_discreted
66+
67+ def get_drift (bag , veh_name ):
68+
69+ # topic name
70+ tpn_drift = f'/{ veh_name }
71+
72+ # get drift from data_drift
73+ drift = []
74+ t_drift = []
75+ for topic , msg , t in bag .read_messages (topics = [tpn_drift ]):
76+ if topic == tpn_drift :
77+ t_drift .append (t .to_sec ())
78+ drift .append (msg )
79+
80+ return drift , t_drift
81+
82+ def main ():
83+
84+ # data extraction from bag file
85+ # Parse command line arguments
86+ parser = argparse .ArgumentParser (description = "Extract images from a ROS bag." )
87+ parser .add_argument ("--bag_folder" , help = "Input directory." )
88+ args = parser .parse_args ()
89+ VEH_NAMES = ["SQ01s" , "SQ02s" ]
90+
91+ # get bags from input directory
92+ bags = []
93+ for file in os .listdir (args .bag_folder ):
94+ if file .endswith (".bag" ):
95+ bags .append (os .path .join (args .bag_folder , file ))
96+
97+ # sort bags by time
98+ bags .sort ()
99+
100+ # for loop
101+ for bag_text in bags :
102+
103+ # open the bag
104+ bag = rosbag .Bag (bag_text , "r" )
105+
106+ # get permutations of veh_names
107+ # veh_name_permutations = permutations(VEH_NAMES)
108+
109+ # get combinations of veh_names
110+ veh_name_combinations = combinations (VEH_NAMES ,2 )
111+
112+ euler_frame_align = []
113+ offsets_frame_align = []
114+ t_frame_align = []
115+ for veh_pair in veh_name_combinations :
116+ euler_frame_align , offsets_frame_align , t_frame_align = get_transformation (bag , veh_pair )
117+
118+ # get the actual transformatoin between the two vehicles
119+ # note: not sure why but directly comparing SQ01s/corrupted_world to SQ02s/corrupted_world by lookupTransform doesn't work
120+ # note: so just get each vehicle's transformation to the world and then subtract them
121+ # note: maybe the order of "world" and frame_name in lookupTransform matters?
122+ # agent 1
123+ # euler_gt1, offsets_gt1, t_actual = get_ground_truth_transformation(bag, VEH_NAMES[0], t_frame_align)
124+ # euler_corrupted1, offsets_corrupted1, _ = get_ground_truth_transformation(bag, VEH_NAMES[0]+"/corrupted_world", t_frame_align)
125+
126+ # euler_actual_drift1 = []
127+ # for i in range(len(euler_gt1)):
128+ # euler_actual_drift1.append(np.array(euler_corrupted1[i]) - np.array(euler_gt1[i]))
129+
130+ # offsets_actual_drift1 = []
131+ # for i in range(len(offsets_gt1)):
132+ # offsets_actual_drift1.append(np.array(offsets_corrupted1[i]) - np.array(offsets_gt1[i]))
133+
134+ # # agent 2
135+ # euler_gt2, offsets_gt2, _ = get_ground_truth_transformation(bag, VEH_NAMES[1], t_frame_align)
136+ # euler_corrupted2, offsets_corrupted2, _ = get_ground_truth_transformation(bag, VEH_NAMES[1]+"/corrupted_world", t_frame_align)
137+
138+ # euler_actual_drift2 = []
139+ # for i in range(len(euler_gt2)):
140+ # euler_actual_drift2.append(np.array(euler_corrupted2[i]) - np.array(euler_gt2[i]))
141+
142+ # offsets_actual_drift2 = []
143+ # for i in range(len(offsets_gt2)):
144+ # offsets_actual_drift2.append(np.array(offsets_corrupted2[i]) - np.array(offsets_gt2[i]))
145+
146+ # # subtract the two
147+ # euler_actual_drift = []
148+ # for i in range(len(euler_actual_drift1)):
149+ # euler_actual_drift.append(np.array(euler_actual_drift1[i]) - np.array(euler_actual_drift2[i]))
150+ # offsets_actual_drift = []
151+ # for i in range(len(offsets_actual_drift1)):
152+ # offsets_actual_drift.append(np.array(offsets_actual_drift1[i]) - np.array(offsets_actual_drift2[i]))
153+
154+ # get drift
155+ drift , t_plot = get_drift (bag , VEH_NAMES [0 ])
156+ # get euler_actual_drift and offsets_actual_drift
157+ euler_actual_drift_roll = []
158+ euler_actual_drift_pitch = []
159+ euler_actual_drift_yaw = []
160+ offsets_actual_drift_x = []
161+ offsets_actual_drift_y = []
162+ for i in range (len (drift )):
163+ euler_actual_drift_roll .append (np .rad2deg (drift [i ].drift_euler [0 ]))
164+ euler_actual_drift_pitch .append (np .rad2deg (drift [i ].drift_euler [1 ]))
165+ euler_actual_drift_yaw .append (np .rad2deg (drift [i ].drift_euler [2 ]))
166+ offsets_actual_drift_x .append (drift [i ].drift_pos [0 ])
167+ offsets_actual_drift_y .append (drift [i ].drift_pos [1 ])
168+
169+ # close the bag
170+ bag .close ()
171+
172+ # sometimes the artificially introduced drift has a little spikes in it (im guessing due to comm delay) so just smooth out the data
173+ # t_plot, euler_actual_drift_roll = smooth_data(t_actual, np.array(euler_actual_drift)[:, 0])
174+ # t_plot, euler_actual_drift_pitch = smooth_data(t_actual, np.array(euler_actual_drift)[:, 1])
175+ # t_plot, euler_actual_drift_yaw = smooth_data(t_actual, np.array(euler_actual_drift)[:, 2])
176+ # t_plot, offsets_actual_drift_x = smooth_data(t_actual, np.array(offsets_actual_drift)[:, 0])
177+ # t_plot, offsets_actual_drift_y = smooth_data(t_actual, np.array(offsets_actual_drift)[:, 1])
178+
179+ # plot
180+ fig , axs = plt .subplots (2 , 1 , figsize = (10 , 10 ))
181+ axs [0 ].plot (t_frame_align , np .array (euler_frame_align )[:, 0 ], label = 'pred roll drift' )
182+ axs [0 ].plot (t_frame_align , np .array (euler_frame_align )[:, 1 ], label = 'pred pitch drift' )
183+ axs [0 ].plot (t_frame_align , np .array (euler_frame_align )[:, 2 ], label = 'pred yaw drift' )
184+ axs [0 ].plot (t_plot , euler_actual_drift_roll , label = 'actual roll drift' )
185+ axs [0 ].plot (t_plot , euler_actual_drift_pitch , label = 'actual pitch drift' )
186+ axs [0 ].plot (t_plot , euler_actual_drift_yaw , label = 'actual yaw drift' )
187+ axs [0 ].set_xlabel ('time [s]' )
188+ axs [0 ].set_ylabel ('euler angles [deg]' )
189+ axs [0 ].legend ()
190+ axs [0 ].grid ()
191+ axs [1 ].plot (t_frame_align , np .array (offsets_frame_align )[:, 0 ], label = 'pred x drift' )
192+ axs [1 ].plot (t_frame_align , np .array (offsets_frame_align )[:, 1 ], label = 'pred y drift' )
193+ # axs[1].plot(t_frame_align, np.array(offsets_frame_align)[:, 2], label='pred z drift')
194+ axs [1 ].plot (t_plot , offsets_actual_drift_x , label = 'actual x drift' )
195+ axs [1 ].plot (t_plot , offsets_actual_drift_y , label = 'actual y drift' )
196+ # axs[1].plot(t_frame_align, np.array(offsets_actual_drift)[:, 2], label='actual z drift')
197+ axs [1 ].set_xlabel ('time [s]' )
198+ axs [1 ].set_ylabel ('offsets [m]' )
199+ axs [1 ].legend ()
200+ axs [1 ].grid ()
201+ plt .tight_layout ()
202+ plt .savefig (os .path .join (args .bag_folder , os .path .splitext (os .path .basename (bag_text ))[0 ] + '.png' ))
203+ plt .show ()
204+
205+ if __name__ == '__main__' :
206+ main ()
0 commit comments