feat: add compare_trajectories.py and extract_pose_from_rosbag.py (…

…#216) * Added compare_trajectories.py and extract_pose_from_rosbag.py Signed-off-by: Shintaro Sakoda <[email protected]> * style(pre-commit): autofix * Rewrite as a literal Signed-off-by: Shintaro Sakoda <[email protected]> * Added "cspell: ignore rotvec" Signed-off-by: Shintaro Sakoda <[email protected]> --------- Signed-off-by: Shintaro Sakoda <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
autowarefoundation · Feb 27, 2025 · 295fc01 · 295fc01
1 parent 52e1711
commit 295fc01
Show file tree

Hide file tree

Showing 7 changed files with 639 additions and 0 deletions.
diff --git a/localization/autoware_localization_evaluation_scripts/CMakeLists.txt b/localization/autoware_localization_evaluation_scripts/CMakeLists.txt
@@ -5,6 +5,8 @@ find_package(autoware_cmake REQUIRED)
 autoware_package()
 
 install(PROGRAMS
+    scripts/compare_trajectories.py
+    scripts/extract_pose_from_rosbag.py
     scripts/plot_diagnostics.py
     DESTINATION lib/${PROJECT_NAME}
 )

diff --git a/localization/autoware_localization_evaluation_scripts/README.md b/localization/autoware_localization_evaluation_scripts/README.md
@@ -59,3 +59,61 @@ $HOME/driving_log_replayer_v2/out/latest/diagnostics_result
 [Example : ndt_scan_matcher__scan_matching_status.png]
 
 ![ndt_scan_matcher__scan_matching_status.png](./media/ndt_scan_matcher__scan_matching_status.png)
+
+## extract_pose_from_rosbag.py
+
+```bash
+ros2 run autoware_localization_evaluation_scripts extract_pose_from_rosbag.py \
+   <rosbag_path> \
+   --save_dir=/your/path (default:rosbag_path/../)
+```
+
+[Example]
+
+```bash
+$ ros2 run autoware_localization_evaluation_scripts extract_pose_from_rosbag.py \
+    $HOME/driving_log_replayer_v2/out/latest/result_bag \
+    --target_topics "/localization/kinematic_state" \
+                    "/localization/pose_estimator/pose_with_covariance"
+```
+
+This command outputs tsv files for each pose.
+
+The file names are the topic names with “/” replaced with “\_\_”.
+
+```bash
+$ tree $HOME/driving_log_replayer_v2/out/latest/pose_tsv
+$HOME/driving_log_replayer_v2/out/latest/pose_tsv
+├── localization__kinematic_state.tsv
+└── localization__pose_estimator__pose_with_covariance.tsv
+
+0 directories, 2 files
+```
+
+## compare_trajectories.py
+
+```bash
+ros2 run autoware_localization_evaluation_scripts compare_trajectories.py \
+   <tsv_file_subject> <tsv_file_reference>
+```
+
+[Example]
+
+```bash
+$ ros2 run autoware_localization_evaluation_scripts compare_trajectories.py \
+    $HOME/driving_log_replayer_v2/out/latest/pose_tsv/localization__kinematic_state.tsv \
+    $HOME/driving_log_replayer_v2/out/latest/pose_tsv/localization__pose_estimator__pose_with_covariance.tsv
+```
+
+This command outputs tsv files for each pose.
+
+```bash
+$ tree $HOME/driving_log_replayer_v2/out/latest/pose_tsv/localization__kinematic_state_result
+$HOME/driving_log_replayer_v2/out/latest/pose_tsv/localization__kinematic_state_result
+├── compare_trajectories.png
+├── relative_pose.png
+├── relative_pose.tsv
+└── relative_pose_summary.tsv
+
+0 directories, 4 files
+```
diff --git a/localization/autoware_localization_evaluation_scripts/scripts/compare_trajectories.py b/localization/autoware_localization_evaluation_scripts/scripts/compare_trajectories.py
@@ -0,0 +1,138 @@
+#!/usr/bin/env python3
+"""A script to compare two trajectories."""
+
+import argparse
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import pandas as pd
+from utils.calc_relative_pose import calc_relative_pose
+from utils.interpolate_pose import interpolate_pose
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("subject_tsv", type=Path)
+    parser.add_argument("reference_tsv", type=Path)
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    subject_tsv = args.subject_tsv
+    reference_tsv = args.reference_tsv
+
+    result_name = subject_tsv.stem
+    save_dir = subject_tsv.parent / f"{result_name}_result"
+    save_dir.mkdir(parents=True, exist_ok=True)
+
+    df_sub = pd.read_csv(subject_tsv, sep="\t")
+    df_ref = pd.read_csv(reference_tsv, sep="\t")
+
+    # plot
+    plt.plot(df_sub["position.x"], df_sub["position.y"], label="subject")
+    plt.plot(df_ref["position.x"], df_ref["position.y"], label="reference")
+    plt.legend()
+    plt.axis("equal")
+    plt.xlabel("x [m]")
+    plt.ylabel("y [m]")
+    plt.savefig(
+        f"{save_dir}/compare_trajectories.png",
+        bbox_inches="tight",
+        pad_inches=0.05,
+        dpi=300,
+    )
+    plt.close()
+
+    # sort by timestamp
+    df_sub = df_sub.sort_values(by="timestamp")
+    df_ref = df_ref.sort_values(by="timestamp")
+
+    # interpolate
+    timestamp = df_sub["timestamp"]
+    ok_mask = (timestamp > df_ref["timestamp"].min()) * (timestamp < df_ref["timestamp"].max())
+    df_sub = df_sub[ok_mask]
+    timestamp = timestamp[ok_mask]
+    df_ref = interpolate_pose(df_ref, timestamp)
+
+    # reset index
+    df_sub = df_sub.reset_index(drop=True)
+    df_ref = df_ref.reset_index(drop=True)
+
+    assert len(df_sub) == len(df_ref), f"len(df_pr)={len(df_sub)}, len(df_gt)={len(df_ref)}"
+
+    # calc mean error
+    diff_x = df_sub["position.x"].to_numpy() - df_ref["position.x"].to_numpy()
+    diff_y = df_sub["position.y"].to_numpy() - df_ref["position.y"].to_numpy()
+    diff_z = df_sub["position.z"].to_numpy() - df_ref["position.z"].to_numpy()
+    diff_meter = (diff_x**2 + diff_y**2 + diff_z**2) ** 0.5
+
+    # calc relative pose
+    df_relative = calc_relative_pose(df_sub, df_ref)
+    df_relative.to_csv(f"{save_dir}/relative_pose.tsv", sep="\t", index=False)
+
+    x_diff_mean = df_relative["position.x"].abs().mean()
+    y_diff_mean = df_relative["position.y"].abs().mean()
+    z_diff_mean = df_relative["position.z"].abs().mean()
+    angle_x_diff_mean = df_relative["angle.x"].abs().mean()
+    angle_y_diff_mean = df_relative["angle.y"].abs().mean()
+    angle_z_diff_mean = df_relative["angle.z"].abs().mean()
+    error_norm = df_relative["position.norm"]
+    df_summary = pd.DataFrame(
+        {
+            "x_diff_mean": [x_diff_mean],
+            "y_diff_mean": [y_diff_mean],
+            "z_diff_mean": [z_diff_mean],
+            "error_mean": [error_norm.mean()],
+            "roll_diff_mean": [angle_x_diff_mean],
+            "pitch_diff_mean": [angle_y_diff_mean],
+            "yaw_diff_mean": [angle_z_diff_mean],
+        },
+    )
+    df_summary.to_csv(
+        f"{save_dir}/relative_pose_summary.tsv",
+        sep="\t",
+        index=False,
+        float_format="%.4f",
+    )
+    print(f"mean error: {error_norm.mean():.3f} m")
+
+    plot_target_list = ["position", "angle"]
+    GUIDELINE_POSITION = 0.5  # [m]
+    GUIDELINE_ANGLE = 0.5  # [degree]
+
+    for i, plot_target in enumerate(plot_target_list):
+        plt.subplot(2, 1, i + 1)
+        plt.plot(df_relative[f"{plot_target}.x"], label="x")
+        plt.plot(df_relative[f"{plot_target}.y"], label="y")
+        plt.plot(df_relative[f"{plot_target}.z"], label="z")
+        guide = GUIDELINE_POSITION if plot_target == "position" else GUIDELINE_ANGLE
+        unit = "degree" if plot_target == "angle" else "m"
+        plt.plot(
+            [0, len(df_relative)],
+            [guide, guide],
+            linestyle="dashed",
+            color="red",
+            label=f"guideline = {guide} [{unit}]",
+        )
+        plt.plot(
+            [0, len(df_relative)],
+            [-guide, -guide],
+            linestyle="dashed",
+            color="red",
+        )
+        bottom, top = plt.ylim()
+        plt.ylim(bottom=min(bottom, -guide * 2), top=max(top, guide * 2))
+        plt.legend(loc="upper left", bbox_to_anchor=(1, 1))
+        plt.xlabel("frame number")
+        plt.ylabel(f"relative {plot_target} [{unit}]")
+        plt.grid()
+    plt.tight_layout()
+    plt.savefig(
+        f"{save_dir}/relative_pose.png",
+        bbox_inches="tight",
+        pad_inches=0.05,
+        dpi=300,
+    )
+    print(f"saved to {save_dir}/relative_pose.png")
+    plt.close()
diff --git a/localization/autoware_localization_evaluation_scripts/scripts/extract_pose_from_rosbag.py b/localization/autoware_localization_evaluation_scripts/scripts/extract_pose_from_rosbag.py
@@ -0,0 +1,43 @@
+#!/usr/bin/env python3
+"""A script to extract pose from rosbag and save as tsv."""
+
+import argparse
+from pathlib import Path
+
+from utils.parse_functions import parse_rosbag
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("rosbag_path", type=Path)
+    parser.add_argument("--save_dir", type=Path, default=None)
+    parser.add_argument("--target_topics", type=str, required=True, nargs="+")
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    rosbag_path = args.rosbag_path
+    target_topics = args.target_topics
+    save_dir = args.save_dir
+
+    if save_dir is None:
+        if rosbag_path.is_dir():  # if specified directory containing db3 files
+            save_dir = rosbag_path.parent / "pose_tsv"
+        else:  # if specified db3 file directly
+            save_dir = rosbag_path.parent.parent / "pose_tsv"
+    save_dir.mkdir(parents=True, exist_ok=True)
+
+    df_dict = parse_rosbag(str(rosbag_path), target_topics)
+
+    for target_topic in target_topics:
+        save_name = "__".join(target_topic.split("/")[1:])
+        df = df_dict[target_topic]
+        df.to_csv(
+            f"{save_dir}/{save_name}.tsv",
+            index=False,
+            sep="\t",
+            float_format="%.9f",
+        )
+
+    print(f"Saved pose tsv files to {save_dir}")
diff --git a/localization/autoware_localization_evaluation_scripts/scripts/utils/calc_relative_pose.py b/localization/autoware_localization_evaluation_scripts/scripts/utils/calc_relative_pose.py
@@ -0,0 +1,64 @@
+"""A script to compare two pose_lists."""
+
+import numpy as np
+import pandas as pd
+from scipy.spatial.transform import Rotation
+
+# cspell: ignore rotvec
+
+
+def calc_relative_pose(df_prd: pd.DataFrame, df_ref: pd.DataFrame) -> pd.DataFrame:
+    """Calculate the relative position and orientation of df_prd with respect to df_ref."""
+    position_keys = ["position.x", "position.y", "position.z"]
+    orientation_keys = [
+        "orientation.x",
+        "orientation.y",
+        "orientation.z",
+        "orientation.w",
+    ]
+    assert len(df_prd) == len(df_ref)
+
+    df_relative = df_prd.copy()
+
+    # Calculate the relative orientation
+    rotation_prd = Rotation.from_quat(df_prd[orientation_keys].values)
+    rotation_ref = Rotation.from_quat(df_ref[orientation_keys].values)
+    df_relative[orientation_keys] = (rotation_prd * rotation_ref.inv()).as_quat()
+
+    # Calculate the relative position
+    df_relative[position_keys] = df_prd[position_keys].to_numpy() - df_ref[position_keys].to_numpy()
+    # Rotate the relative position of each frame by rotation_true.inv()
+    # This makes the relative position based on the pose of df_ref
+    df_relative[position_keys] = rotation_ref.inv().apply(
+        df_relative[position_keys].values,
+    )
+
+    # Add norm
+    df_relative["position.norm"] = np.linalg.norm(
+        df_relative[position_keys].values,
+        axis=1,
+    )
+
+    # Add rpy angle
+    r = Rotation.from_quat(
+        df_relative[["orientation.x", "orientation.y", "orientation.z", "orientation.w"]],
+    )
+    euler = r.as_euler("xyz", degrees=True)
+    df_relative["angle.x"] = euler[:, 0]
+    df_relative["angle.y"] = euler[:, 1]
+    df_relative["angle.z"] = euler[:, 2]
+    df_relative["angle.norm"] = np.linalg.norm(r.as_rotvec(), axis=1)
+
+    # Arrange the order of columns
+    return df_relative[
+        [
+            "timestamp",
+            *position_keys,
+            "position.norm",
+            *orientation_keys,
+            "angle.x",
+            "angle.y",
+            "angle.z",
+            "angle.norm",
+        ]
+    ]
diff --git a/localization/autoware_localization_evaluation_scripts/scripts/utils/interpolate_pose.py b/localization/autoware_localization_evaluation_scripts/scripts/utils/interpolate_pose.py
@@ -0,0 +1,73 @@
+"""A script to interpolate poses to match the timestamp in target_timestamp."""
+
+import pandas as pd
+from scipy.spatial.transform import Rotation
+from scipy.spatial.transform import Slerp
+
+
+def interpolate_pose(df_pose: pd.DataFrame, target_timestamp: pd.Series) -> pd.DataFrame:
+    """Interpolate each pose in df_pose to match the timestamp in target_timestamp."""
+    # check monotonicity
+    assert df_pose["timestamp"].is_monotonic_increasing
+    assert target_timestamp.is_monotonic_increasing
+
+    # check length
+    assert len(df_pose) > 0, f"{len(df_pose)=}"
+    assert len(target_timestamp) > 0, f"{len(target_timestamp)=}"
+
+    # check range
+    assert df_pose.iloc[0]["timestamp"] <= target_timestamp.iloc[0]
+    assert target_timestamp.iloc[-1] <= df_pose.iloc[-1]["timestamp"]
+
+    position_keys = ["position.x", "position.y", "position.z"]
+    orientation_keys = [
+        "orientation.x",
+        "orientation.y",
+        "orientation.z",
+        "orientation.w",
+    ]
+
+    df_pose = df_pose.reset_index(drop=True)
+    target_timestamp = target_timestamp.reset_index(drop=True)
+
+    target_index = 0
+    df_index = 0
+    data_dict = {
+        "timestamp": [],
+        **{key: [] for key in position_keys},
+        **{key: [] for key in orientation_keys},
+    }
+    while df_index < len(df_pose) - 1 and target_index < len(target_timestamp):
+        curr_time = df_pose.iloc[df_index]["timestamp"]
+        next_time = df_pose.iloc[df_index + 1]["timestamp"]
+        target_time = target_timestamp[target_index]
+
+        # Find a df_index that includes target_time
+        if not (curr_time <= target_time <= next_time):
+            df_index += 1
+            continue
+
+        curr_weight = (next_time - target_time) / (next_time - curr_time)
+        next_weight = 1.0 - curr_weight
+
+        curr_position = df_pose.iloc[df_index][position_keys]
+        next_position = df_pose.iloc[df_index + 1][position_keys]
+        target_position = curr_position * curr_weight + next_position * next_weight
+
+        curr_orientation = df_pose.iloc[df_index][orientation_keys]
+        next_orientation = df_pose.iloc[df_index + 1][orientation_keys]
+        curr_r = Rotation.from_quat(curr_orientation)
+        next_r = Rotation.from_quat(next_orientation)
+        slerp = Slerp([curr_time, next_time], Rotation.concatenate([curr_r, next_r]))
+        target_orientation = slerp([target_time]).as_quat()[0]
+
+        data_dict["timestamp"].append(target_time)
+        data_dict[position_keys[0]].append(target_position[0])
+        data_dict[position_keys[1]].append(target_position[1])
+        data_dict[position_keys[2]].append(target_position[2])
+        data_dict[orientation_keys[0]].append(target_orientation[0])
+        data_dict[orientation_keys[1]].append(target_orientation[1])
+        data_dict[orientation_keys[2]].append(target_orientation[2])
+        data_dict[orientation_keys[3]].append(target_orientation[3])
+        target_index += 1
+    return pd.DataFrame(data_dict)