120 lines
4.1 KiB
Python
120 lines
4.1 KiB
Python
|
"""
|
|||
|
|
# 收集实飞数据，记录位置、姿态、图像，用于离线fine-tuning (保存至save_dir)
|
|||
|
|
# 注意： 由于里程计漂移，可能utils/pointcloud_clip需要对地图进行微调，需对无人机位置和yaw, pitch, roll做相同的变换
|
|||
|
|
# 注意保证地图和里程计处于同一坐标系，同时录包+保存地图
|
|||
|
|
"""
|
|||
|
|
import cv2
|
|||
|
|
import numpy as np
|
|||
|
|
import time, os, sys
|
|||
|
|
from cv_bridge import CvBridge, CvBridgeError
|
|||
|
|
import rospy
|
|||
|
|
from sensor_msgs.msg import Image
|
|||
|
|
from nav_msgs.msg import Odometry
|
|||
|
|
from scipy.spatial.transform import Rotation
|
|||
|
|
|
|||
|
|
depth_img = np.zeros([270, 480])
|
|||
|
|
pos = np.array([0, 0, 0])
|
|||
|
|
quat = np.array([1, 0, 0, 0])
|
|||
|
|
positions = []
|
|||
|
|
quaternions = []
|
|||
|
|
frame_id = 0
|
|||
|
|
new_depth = False
|
|||
|
|
new_odom = False
|
|||
|
|
first_frame = True
|
|||
|
|
last_time = time.time()
|
|||
|
|
save_dir = os.environ["FLIGHTMARE_PATH"] + "/run/depth_realworld"
|
|||
|
|
label_path = save_dir + "/label.npz"
|
|||
|
|
if not os.path.exists(save_dir):
|
|||
|
|
os.mkdir(save_dir)
|
|||
|
|
# Due to odometry drift, the map is adjusted, and the drone's position is also adjusted accordingly.
|
|||
|
|
R_no = Rotation.from_euler('ZYX', [15, 3, 0.0], degrees=True) # yaw, pitch, roll
|
|||
|
|
translation_no = np.array([0, 0, 2])
|
|||
|
|
|
|||
|
|
|
|||
|
|
def callback_odometry(data):
|
|||
|
|
# NWU
|
|||
|
|
global pos, quat, new_odom, R_no, translation_no
|
|||
|
|
p_ob = np.array([[data.pose.pose.position.x],
|
|||
|
|
[data.pose.pose.position.y],
|
|||
|
|
[data.pose.pose.position.z]])
|
|||
|
|
q_ob = np.array([data.pose.pose.orientation.x,
|
|||
|
|
data.pose.pose.orientation.y,
|
|||
|
|
data.pose.pose.orientation.z,
|
|||
|
|
data.pose.pose.orientation.w])
|
|||
|
|
R_ob = Rotation.from_quat(q_ob) # old->body (xyzw)
|
|||
|
|
quat_xyzw = (R_no * R_ob).as_quat() # new->body (xyzw)
|
|||
|
|
quat = np.array([quat_xyzw[3], quat_xyzw[0], quat_xyzw[1], quat_xyzw[2]])
|
|||
|
|
pos = np.squeeze(np.dot(R_no.as_matrix(), p_ob)) + translation_no
|
|||
|
|
new_odom = True
|
|||
|
|
|
|||
|
|
|
|||
|
|
def callback_depth(data):
|
|||
|
|
global depth_img, new_depth
|
|||
|
|
max_dis = 20.0
|
|||
|
|
min_dis = 0.03
|
|||
|
|
height = 270
|
|||
|
|
width = 480
|
|||
|
|
scale = 0.001
|
|||
|
|
bridge = CvBridge()
|
|||
|
|
try:
|
|||
|
|
depth_ = bridge.imgmsg_to_cv2(data, "32FC1")
|
|||
|
|
except:
|
|||
|
|
print("CV_bridge ERROR: Your ros and python path has something wrong!")
|
|||
|
|
|
|||
|
|
if depth_.shape[0] != height or depth_.shape[1] != width:
|
|||
|
|
depth_ = cv2.resize(depth_, (width, height), interpolation=cv2.INTER_NEAREST)
|
|||
|
|
depth_ = np.minimum(depth_ * scale, max_dis) / max_dis
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
nan_mask = np.isnan(depth_) | (depth_ < min_dis)
|
|||
|
|
depth_ = cv2.inpaint(np.uint8(depth_ * 255), np.uint8(nan_mask), 3, cv2.INPAINT_NS)
|
|||
|
|
depth_ = depth_.astype(np.float32) / 255.0
|
|||
|
|
except:
|
|||
|
|
print("Interpolation failed")
|
|||
|
|
|
|||
|
|
# Not necessary, but encountered some inexplicable errors previously, so temporarily kept.
|
|||
|
|
if np.sum(np.isnan(depth_)) > 0:
|
|||
|
|
depth_[np.isnan(depth_)] = 0
|
|||
|
|
print("WARN: Have NAN values in depth image")
|
|||
|
|
|
|||
|
|
depth_img = depth_.copy()
|
|||
|
|
new_depth = True
|
|||
|
|
|
|||
|
|
|
|||
|
|
def save_data(_timer):
|
|||
|
|
global pos, quat, new_odom, depth_img, new_depth, last_time, first_frame
|
|||
|
|
global save_dir, label_path, frame_id, positions, quaternions
|
|||
|
|
if not (new_odom and new_depth):
|
|||
|
|
if not first_frame and time.time() - last_time > 1:
|
|||
|
|
np.savez(
|
|||
|
|
label_path,
|
|||
|
|
positions=np.asarray(positions),
|
|||
|
|
quaternions=np.asarray(quaternions),
|
|||
|
|
)
|
|||
|
|
print("Record Done!")
|
|||
|
|
sys.exit()
|
|||
|
|
return
|
|||
|
|
new_odom, new_depth = False, False
|
|||
|
|
|
|||
|
|
image_path = save_dir + "/img_" + str(frame_id) + ".tif"
|
|||
|
|
cv2.imwrite(image_path, depth_img)
|
|||
|
|
positions.append(pos)
|
|||
|
|
quaternions.append(quat)
|
|||
|
|
|
|||
|
|
last_time = time.time()
|
|||
|
|
first_frame = False
|
|||
|
|
frame_id = frame_id + 1
|
|||
|
|
|
|||
|
|
|
|||
|
|
def main():
|
|||
|
|
rospy.init_node('data_collect', anonymous=False)
|
|||
|
|
odom_ref_sub = rospy.Subscriber("/odometry/imu", Odometry, callback_odometry, queue_size=1)
|
|||
|
|
depth_sub = rospy.Subscriber("/camera/depth/image_rect_raw", Image, callback_depth, queue_size=1)
|
|||
|
|
timer = rospy.Timer(rospy.Duration(0.033), save_data)
|
|||
|
|
print("Data Collection Node Ready!")
|
|||
|
|
rospy.spin()
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
main()
|