diff --git a/CMakeLists.txt b/CMakeLists.txt index d65df4b5d56bd0c665c8cf4aa52c286dde62f750..d949ee3a584eeef39e1b1f861fcd41d1148a686c 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -167,6 +167,7 @@ include_directories( ## in contrast to setup.py, you can choose the destination catkin_install_python(PROGRAMS scripts/pohang2rosbag.py + scripts/pohang_add_missing_imu.py DESTINATION ${CATKIN_PACKAGE_BIN_DESTINATION} ) diff --git a/dataset_converter/pohang.py b/dataset_converter/pohang.py index c8448a595305441eeb690d3610e5fa3448b62c16..c5cbadb0bfc5a234ae965e22063cdf00041e58a5 100644 --- a/dataset_converter/pohang.py +++ b/dataset_converter/pohang.py @@ -35,6 +35,17 @@ LIDAR_PROCESS = { "/lidar_starboard/os_cloud_node/points", ), } +LIDAR_IMU_PROCESS = { + "lidar_front/imu.txt": ( + "lidar_front/imu", + "/lidar_front/imu/data", + ), + "lidar_port/imu.txt": ("lidar_port/imu", "/lidar_port/imu/data"), + "lidar_starboard/imu.txt": ( + "lidar_starboard/imu", + "/lidar_starboard/imu/data", + ), +} STEREO_PROCESS = { "left_images": ("stereo_left_link", "/stereo_cam/left_img/compressed"), "right_images": ("stereo_right_link", "/stereo_cam/right_img/compressed"), @@ -243,6 +254,35 @@ def process_imu( imu.header.seq = idx imu.header.stamp = t imu.header.frame_id = frame_id + imu.orientation.x = 0 + imu.orientation.y = 0 + imu.orientation.z = 0 + imu.orientation.w = 1 + imu.angular_velocity.x = l[1] + imu.angular_velocity.y = l[2] + imu.angular_velocity.z = l[3] + imu.linear_acceleration.x = l[4] + imu.linear_acceleration.y = l[5] + imu.linear_acceleration.z = l[5] + + bag.write(topic, imu, t) + + +def process_navigation(data_folder: Path, rosbag_path: Path): + data_path = data_folder.joinpath("navigation") + + # AHRS Part + data = txt_preprocess(data_path.joinpath("ahrs.txt")) + + with rosbag.Bag(rosbag_path, check_bag_mode(rosbag_path)) as bag: + for idx, l in enumerate(data): + l = list(map(float, l)) + t = rospy.Time.from_sec(l[0]) + + imu = Imu() + imu.header.seq = idx + imu.header.stamp = t + imu.header.frame_id = NAVIGATION_PROCESS["ahrs.txt"][0] imu.orientation.x = l[1] imu.orientation.y = l[2] imu.orientation.z = l[3] @@ -254,19 +294,7 @@ def process_imu( imu.linear_acceleration.y = l[9] imu.linear_acceleration.z = l[10] - bag.write(topic, imu, t) - - -def process_navigation(data_folder: Path, rosbag_path: Path): - data_path = data_folder.joinpath("navigation") - - # AHRS Part - process_imu( - data_path.joinpath("ahrs.txt"), - rosbag_path, - NAVIGATION_PROCESS["ahrs.txt"][1], - NAVIGATION_PROCESS["ahrs.txt"][0], - ) + bag.write(NAVIGATION_PROCESS["ahrs.txt"][1], imu, t) # Baseline Part # frame_id unknown @@ -338,6 +366,9 @@ def process_lidar(data_folder: Path, rosbag_path: Path): for k, v in LIDAR_PROCESS.items(): process_pointcloud(data_path / k, rosbag_path, v[1], v[0]) + for k, v in LIDAR_IMU_PROCESS.items(): + process_imu(data_path / k, rosbag_path, v[1], v[0]) + def process_infrared(data_folder: Path, rosbag_path: Path): data_path = data_folder.joinpath("infrared") diff --git a/scripts/pohang_add_missing_imu.py b/scripts/pohang_add_missing_imu.py new file mode 100644 index 0000000000000000000000000000000000000000..8340dc6b739d0ebfac2292a42dce2a54f4a2a4ce --- /dev/null +++ b/scripts/pohang_add_missing_imu.py @@ -0,0 +1,26 @@ +#!/usr/bin/env python3 + +# add missing lidar imu data +import dataset_converter.pohang as pohang +from pathlib import Path +from concurrent.futures import ProcessPoolExecutor + +# pohang folder -> rosbag path +DATA_MAP = { + "/workspace/data/pohang00_pier/lidar": "/workspace/data/pohang00_pier.bag", + "/workspace/data/pohang00_canal/lidar": "/workspace/data/pohang00_canal.bag", +} + + +def add_missing_imu(pohang_data_folder: Path, bag_path: Path): + for k, v in pohang.LIDAR_IMU_PROCESS.items(): + pohang.process_imu(pohang_data_folder.joinpath(k), bag_path, v[1], v[0]) + + print(f"{pohang_data_folder} Done") + + +if __name__ == "__main__": + + for data_folder, bag_path in DATA_MAP.items(): + add_missing_imu(Path(data_folder), Path(bag_path)) +