-
Notifications
You must be signed in to change notification settings - Fork 11
/
Copy pathreset_imu_rotation.py
executable file
·151 lines (114 loc) · 7.43 KB
/
reset_imu_rotation.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
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
#!/usr/bin/env python
# coding=UTF-8
import rospy
import rosbag
import argparse
import tf2_msgs.msg
import geometry_msgs.msg
import sensor_msgs.msg
from tf import transformations
def str2bool(v):
return v.lower() in ("yes", "true", "t", "1")
def reset_quaternion(quaternion, reset_config):
print ' Reset configuration: %s' % (reset_config)
reset_yaw = 0
reset_pitch = 0
reset_roll = 0
yaw, pitch, roll = transformations.euler_from_quaternion(quaternion, axes='szyx')
if 'y' in reset_config:
print ' Resetting yaw'
reset_yaw = yaw
if 'p' in reset_config:
print ' Resetting pitch'
reset_pitch = pitch
if 'r' in reset_config:
print ' Resetting roll'
reset_roll = roll
return transformations.quaternion_from_euler(reset_yaw, reset_pitch, reset_roll, axes='szyx')
def print_reset_quaternion(quaternion, message_header):
reset_yaw, reset_pitch, reset_roll = transformations.euler_from_quaternion(quaternion, axes='szyx')
print '%s | quaternion [qx: %f, qy: %f, qz: %f, qw: %f] and fixed euler angles [yaw: %f, pitch: %f, roll: %f]' % (message_header, quaternion[0], quaternion[1], quaternion[2], quaternion[3], reset_yaw, reset_pitch, reset_roll)
def print_reset_coversion(original_quaternion, new_quaternion):
original_yaw, original_pitch, original_roll = transformations.euler_from_quaternion(original_quaternion, axes='szyx')
new_yaw, new_pitch, new_roll = transformations.euler_from_quaternion(new_quaternion, axes='szyx')
print ' Reset orientation [qx: %f, qy: %f, qz: %f, qw: %f] -> [qx: %f, qy: %f, qz: %f, qw: %f] | [y: %f, p: %f, r: %f] -> [y: %f, p: %f, r: %f]' % (original_quaternion[0], original_quaternion[1], original_quaternion[2], original_quaternion[3], new_quaternion[0], new_quaternion[1], new_quaternion[2], new_quaternion[3], original_yaw, original_pitch, original_roll, new_yaw, new_pitch, new_roll)
def compute_quaternion_euler_offset(original_quaternion, new_quaternion):
offset_quaternion = transformations.quaternion_multiply(transformations.quaternion_inverse(new_quaternion), original_quaternion)
yaw, pitch, roll = transformations.euler_from_quaternion(offset_quaternion, axes='szyx')
return (abs(yaw) + abs(pitch) + abs(roll))
def reset_rotation(inbag_filename, outbag_filename, imu_topic, pose_topic, reset_config, euler_angle_offset_to_discard_initial_inconsistent_msgs, reset_rotation_msgs_initial_window, print_verbose_info):
print ' Resetting IMU orientation in topic %s within bag %s' % (imu_topic, inbag_filename)
print ' Resetting PoseStamped orientation in topic %s within bag %s' % (pose_topic, inbag_filename)
inbag = rosbag.Bag(inbag_filename,'r')
outbag = rosbag.Bag(outbag_filename, 'w', rosbag.bag.Compression.BZ2)
reset_quaternion_imu_valid = False
reset_quaternion_pose_valid = False
number_imu_msgs_reset = 0
number_pose_msgs_reset = 0
for topic, msg, t in inbag.read_messages():
if topic == imu_topic:
current_quaternion = [msg.orientation.x, msg.orientation.y, msg.orientation.z, msg.orientation.w]
if not reset_quaternion_imu_valid:
quaternion_imu = current_quaternion
reset_quaternion_imu = transformations.quaternion_inverse(current_quaternion)
if reset_config:
reset_quaternion_imu = reset_quaternion(reset_quaternion_imu, reset_config)
reset_quaternion_imu_valid = True
print_reset_quaternion(reset_quaternion_imu, ' Resetting IMU messages\t\t')
if number_imu_msgs_reset < reset_rotation_msgs_initial_window:
euler_offset = compute_quaternion_euler_offset(quaternion_imu, current_quaternion)
if euler_offset > euler_angle_offset_to_discard_initial_inconsistent_msgs:
print ' Discarded inconsistent reset imu msg (nº %i with offset %f)' % (number_imu_msgs_reset, euler_offset)
reset_quaternion_imu_valid = False
new_quaternion = transformations.quaternion_multiply(reset_quaternion_imu, current_quaternion)
msg.orientation.x = new_quaternion[0]
msg.orientation.y = new_quaternion[1]
msg.orientation.z = new_quaternion[2]
msg.orientation.w = new_quaternion[3]
if print_verbose_info:
print_reset_coversion(current_quaternion, new_quaternion)
number_imu_msgs_reset += 1
if topic == pose_topic:
current_quaternion = [msg.pose.orientation.x, msg.pose.orientation.y, msg.pose.orientation.z, msg.pose.orientation.w]
if not reset_quaternion_pose_valid:
quaternion_pose = current_quaternion
reset_quaternion_pose = transformations.quaternion_inverse(current_quaternion)
if reset_config:
reset_quaternion_pose = reset_quaternion(reset_quaternion_pose, reset_config)
reset_quaternion_pose_valid = True
print_reset_quaternion(reset_quaternion_pose, ' Resetting PoseStamped messages\t')
if number_pose_msgs_reset < reset_rotation_msgs_initial_window:
euler_offset = compute_quaternion_euler_offset(quaternion_pose, current_quaternion)
if euler_offset > euler_angle_offset_to_discard_initial_inconsistent_msgs:
print ' Discarded inconsistent reset pose msg (nº %i with offset %f)' % (number_pose_msgs_reset, euler_offset)
reset_quaternion_pose_valid = False
new_quaternion = transformations.quaternion_multiply(reset_quaternion_pose, current_quaternion)
msg.pose.orientation.x = new_quaternion[0]
msg.pose.orientation.y = new_quaternion[1]
msg.pose.orientation.z = new_quaternion[2]
msg.pose.orientation.w = new_quaternion[3]
if print_verbose_info:
print_reset_coversion(current_quaternion, new_quaternion)
number_pose_msgs_reset += 1
outbag.write(topic, msg, t)
print 'Closing output bagfile %s' % (outbag_filename)
inbag.close()
outbag.close()
if __name__ == "__main__":
parser = argparse.ArgumentParser(description='Script to reset IMU and PoseStamped orientation')
parser.register('type', 'bool', str2bool)
parser.add_argument('-i', metavar='INPUT_BAGFILE', required=True, help='Input bagfile')
parser.add_argument('-o', metavar='OUTPUT_BAGFILE', required=True, help='Output bagfile')
parser.add_argument('-t', metavar='TOPIC_IMU', required=True, help='Topic to reset IMU rotation')
parser.add_argument('-p', metavar='TOPIC_POSE', required=True, help='Topic to reset PoseStamped rotation')
parser.add_argument('-r', metavar='RESET_CONFIG', required=False, default='', help='Reset config (string with rpy to reset each / combined fixed euler angles)')
parser.add_argument('-s', type=float, required=False, default=0.17, help='Euler angle offset to discard initial inconsistent msgs')
parser.add_argument('-w', type=int, required=False, default=3, help='Reset rotation msgs initial window')
parser.add_argument('-v', metavar='PRINT_VERBOSE_INFO', type='bool', required=False, default=False, help='Print verbose convertion info')
args = parser.parse_args()
try:
reset_rotation(args.i, args.o, args.t, args.p, args.r, args.s, args.w, args.v)
exit(0)
except Exception, e:
import traceback
traceback.print_exc()