diff --git a/dataset_converter/pohang.py b/dataset_converter/pohang.py index 6c14ec8e8bd584be2fe325d407edbbaf752d27f1..6b7e22714cfc2e047d6ae3c444ee0ebda45cfd36 100644 --- a/dataset_converter/pohang.py +++ b/dataset_converter/pohang.py @@ -5,10 +5,14 @@ import cv_bridge import dataset_converter.util as util from pathlib import Path from sensor_msgs.msg import Imu +from geometry_msgs.msg import PoseStamped # used by process func # map "folder_name(file_name)" : ("LINK_NAME", "TOPIC_NAME") -NAVIGATION_PROCESS = {"ahrs.txt": ("gx5_link", "/gx5/imu/data")} +NAVIGATION_PROCESS = { + "ahrs.txt": ("gx5_link", "/gx5/imu/data"), + "baseline.txt": ("gx5_link", "/gx5/baseline"), +} STEREO_PROCESS = { "left_images": ("stereo_left_link", "/stereo_cam/left_img/compressed"), "right_images": ("stereo_right_link", "/stereo_cam/right_img/compressed"), @@ -76,9 +80,7 @@ def process_images( images = util.get_sorted_file_list(data_path) # read all lines and divide into [[timestamp, seq], ... ] - with open(data_path.parent / "timestamp.txt", "r") as timedata: - timestamps = timedata.readlines() - timestamps = [t.strip().split("\t") for t in timestamps] + timestamps = txt_preprocess(data_path.parent / "timestamp.txt") if not len(images) == len(timestamps): print( @@ -122,10 +124,8 @@ def process_images( def process_navigation(data_folder: Path, rosbag_path: Path): data_path = data_folder.joinpath("navigation") - ahrs = data_path.joinpath("ahrs.txt") - with open(ahrs, "r") as f: - data = f.readlines() - data = [l.strip().split("\t") for l in data] + # AHRS Part + data = txt_preprocess(data_path.joinpath("ahrs.txt")) bagmode = check_bag_mode(rosbag_path) with rosbag.Bag(rosbag_path, bagmode) as bag: @@ -152,6 +152,30 @@ def process_navigation(data_folder: Path, rosbag_path: Path): print("AHRS Done!") + # Baseline Part + # frame_id unknown + data = txt_preprocess(data_path.joinpath("baseline.txt")) + bagmode = check_bag_mode(rosbag_path) + with rosbag.Bag(rosbag_path, bagmode) as bag: + for idx, l in enumerate(data): + l = list(map(float, l)) + t = rospy.Time.from_sec(l[0]) + + pose = PoseStamped() + pose.header.seq = idx + pose.header.stamp = t + pose.header.frame_id = NAVIGATION_PROCESS["baseline.txt"][0] + pose.pose.orientation.x = l[1] + pose.pose.orientation.y = l[2] + pose.pose.orientation.z = l[3] + pose.pose.orientation.w = l[4] + pose.pose.position.x = l[5] + pose.pose.position.y = l[6] + pose.pose.position.z = l[7] + + bag.write(NAVIGATION_PROCESS["baseline.txt"][1], pose, t) + + print("Baseline Done!") def process_stereo(data_folder: Path, rosbag_path: Path): data_path = data_folder.joinpath("stereo") @@ -182,7 +206,15 @@ def process_omni(data_folder: Path, rosbag_path: Path): print("Omni Done!") -def check_bag_mode(rosbag_path): +def txt_preprocess(path: Path): + with open(path, "r") as f: + data = f.readlines() + data = [l.strip().split("\t") for l in data] + + return data + + +def check_bag_mode(rosbag_path: Path): if rosbag_path.exists(): bagmode = "a" else: