add baseline

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: