20172884498

Committed 12 Dec 2025 04:18PM UTC coverage: 74.672% (+1.2%) from 73.457%

Build # 20172884498

Build Type

push

github

Committed by

web-flow

Commit Message

(v0.1.2) Path Metrics & Transformations, Python >=3.8 Support

Run Details

269 of 364 new or added lines in 8 files covered. (73.9%)

6 existing lines in 4 files now uncovered.

1082 of 1449 relevant lines covered (74.67%)

0.75 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

92.13

/src/robotdataprocess/data_types/OdometryData.py

from __future__ import annotations

from ..conversion_utils import col_to_dec_arr
import csv
from .Data import CoordinateFrame, Data
import decimal
from decimal import Decimal
from evo.core import geometry
import matplotlib.pyplot as plt
import numpy as np
from numpy.typing import NDArray
import os
import pandas as pd
from .PathData import PathData
from pathlib import Path
from ..rosbag.Ros2BagWrapper import Ros2BagWrapper
from rosbags.rosbag2 import Reader as Reader2
from rosbags.typesys import Stores, get_typestore
from rosbags.typesys.store import Typestore
from scipy.spatial.transform import Rotation as R
from typeguard import typechecked
from typing import Union, List, Tuple
import tqdm

class OdometryData(PathData):

    # Define odometry-specific data attributes
    child_frame_id: str
    poses: list # Saved nav_msgs/msg/Pose

    @typechecked
    def __init__(self, frame_id: str, child_frame_id: str, timestamps: Union[np.ndarray, list], 
                 positions: Union[np.ndarray, list], orientations: Union[np.ndarray, list], frame: CoordinateFrame):
        
        # Copy initial values into attributes
        super().__init__(frame_id, timestamps, positions, orientations, frame)
        self.child_frame_id: str = child_frame_id
        self.poses: list = []

        # Check to ensure that all arrays have same length
        if len(self.timestamps) != len(self.positions) or len(self.positions) != len(self.orientations):
            raise ValueError("Lengths of timestamp, position, and orientation arrays are not equal!")

    # =========================================================================
    # ============================ Class Methods ============================== 
    # =========================================================================  

    @classmethod
    @typechecked
    def from_ros2_bag(cls, bag_path: Union[Path, str], odom_topic: str, frame: CoordinateFrame):
        """
        Creates a class structure from a ROS2 bag file with an Odometry topic.

        Args:
            bag_path (Path | str): Path to the ROS2 bag file.
            odom_topic (str): Topic of the Odometry messages.
        Returns:
            OdometryData: Instance of this class.
        """

        # Get topic message count and typestore
        bag_wrapper = Ros2BagWrapper(bag_path, None)
        typestore: Typestore = bag_wrapper.get_typestore()
        num_msgs: int = bag_wrapper.get_topic_count(odom_topic)
        
        # Pre-allocate arrays (memory-mapped or otherwise)
        timestamps_np = np.zeros(num_msgs, dtype=Decimal)
        positions_np = np.zeros((num_msgs, 3), dtype=Decimal)
        orientations_np = np.zeros((num_msgs, 4), dtype=Decimal)

        # Setup tqdm bar & counter
        pbar = tqdm.tqdm(total=num_msgs, desc="Extracting Odometry...", unit=" msgs")
        i = 0

        # Extract the odometry information
        frame_id, child_frame_id = None, None
        with Reader2(str(bag_path)) as reader:

            # Extract frames from first message
            connections = [x for x in reader.connections if x.topic == odom_topic]
            for conn, timestamp, rawdata in reader.messages(connections=connections):  
                msg = typestore.deserialize_cdr(rawdata, conn.msgtype)
                frame_id = msg.header.frame_id
                child_frame_id = msg.child_frame_id
                break

            # Extract individual message data
            connections = [x for x in reader.connections if x.topic == odom_topic]
            for conn, timestamp, rawdata in reader.messages(connections=connections):
                msg = typestore.deserialize_cdr(rawdata, conn.msgtype)
                
                timestamps_np[i] = bag_wrapper.extract_timestamp(msg)
                pos = msg.pose.pose.position
                positions_np[i] = np.array([Decimal(pos.x), Decimal(pos.y), Decimal(pos.z)])
                ori = msg.pose.pose.orientation
                orientations_np[i] = np.array([Decimal(ori.x), Decimal(ori.y), Decimal(ori.z), Decimal(ori.w)])

                # NOTE: Doesn't support Twist information currently

                # Increment the count
                i += 1
                pbar.update(1)

        # Create an OdometryData class
        return cls(frame_id, child_frame_id, timestamps_np, positions_np, orientations_np, frame)
    
    @classmethod
    @typechecked
    def from_csv(cls, csv_path: Union[Path, str], frame_id: str, child_frame_id: str, frame: CoordinateFrame,          
                 header_included: bool, column_to_data: Union[List[int], None] = None, 
                 separator: Union[str, None] = None, filter: Union[Tuple[str, str], None] = None,
                 ts_in_ns: bool = False):
        """
        Creates a class structure from a csv file.

        Args:
            csv_path (Path | str): Path to the CSV file.
            frame_id (str): The frame that this odometry is relative to.
            child_frame_id (str): The frame whose pose is represented by this odometry.
            frame (CoordinateFrame): The coordinate system convention of this data.
            header_included (bool): If this csv file has a header, so we can remove it.
            column_to_data (list[int]): Tells the algorithms which columns in the csv contain which
                of the following data: ['timestamp', 'x', 'y', 'z', 'qw', 'qx', 'qy', 'qz']. Thus, 
                index 0 of column_to_data should be the column that timestamp data is found in the 
                csv file. Set to None to use [0,1,2,3,4,5,6,7].
            separator (str | None): The separator used in the csv file. If None, will use a comma by default.
            filter: A tuple of (column_name, value), where only rows with column_name == value will be kept. If
                csv file has no headers, then `column_name` should be the index of the column as a string.
            ts_in_ns (bool): If True, assumes timestamps are in nanoseconds and converts to seconds. Otherwise,
                assumes timestamps are already in seconds.

        Returns:
            OdometryData: Instance of this class.
        """

        # If column_to_data is None, assume default
        if column_to_data is None:
            column_to_data = [0,1,2,3,4,5,6,7]
        else:
            # Check column_to_data values are valid
            assert np.all(np.array(column_to_data) >= 0)
            assert len(column_to_data) == 8

        # Read the csv file
        header = 0 if header_included else None
        df1 = pd.read_csv(str(csv_path), header=header, index_col=False, sep=separator, engine='python')

        # Rename columns to standard names
        rename_dict = {}
        desired_data = ['timestamp', 'x', 'y', 'z', 'qw', 'qx', 'qy', 'qz']
        for j, ind in enumerate(column_to_data):
            rename_dict[df1.columns[ind]] = desired_data[j]
        df1 = df1.rename(columns=rename_dict)

        # Using the filter if provided, remove unwanted rows
        if filter is not None:
            df1 = df1[df1[filter[0]] == filter[1]]

        # Convert columns to NDArray[Decimal]
        timestamps_np = np.array([Decimal(str(ts)) for ts in df1['timestamp']], dtype=object)
        positions_np = np.array([[Decimal(str(x)), Decimal(str(y)), Decimal(str(z))] 
                                 for x, y, z in zip(df1['x'], df1['y'], df1['z'])], dtype=object)
        orientations_np = np.array([[Decimal(str(qx)), Decimal(str(qy)), Decimal(str(qz)), Decimal(str(qw))]
                                    for qx, qy, qz, qw in zip(df1['qx'], df1['qy'], df1['qz'], df1['qw'])], dtype=object)
        
        # If timestamps are in ns, convert to s
        if ts_in_ns:
            timestamps_np = timestamps_np / Decimal('1e9')

        # Create an OdometryData class
        return cls(frame_id, child_frame_id, timestamps_np, positions_np, orientations_np, frame)
    
    @classmethod
    @typechecked
    def from_txt_file(cls, file_path: Union[Path, str], frame_id: str, child_frame_id: str, frame: CoordinateFrame):
        """
        Creates a class structure from a text file, where the order of values
        in the files follows ['timestamp', 'x', 'y', 'z', 'qw', 'qx', 'qy', 'qz'].

        Args:
            file_path (Path | str): Path to the file containing the odometry data.
            frame_id (str): The frame where this odometry is relative to.
            child_frame_id (str): The frame whose pose is represented by this odometry.
            frame (CoordinateFrame): The coordinate system convention of this data.
        Returns:
            OdometryData: Instance of this class.
        """
        
        # Count the number of lines in the file
        line_count = 0
        with open(str(file_path), 'r') as file:
            for _ in file: 
                line_count += 1

        # Setup arrays to hold data
        timestamps_np = np.zeros((line_count), dtype=object)
        positions_np = np.zeros((line_count, 3), dtype=object)
        orientations_np = np.zeros((line_count, 4), dtype=object)

        # Open the txt file and read in the data
        with open(str(file_path), 'r') as file:
            for i, line in enumerate(file):
                line_split = line.split(' ')
                timestamps_np[i] = line_split[0]
                positions_np[i] = line_split[1:4]
                orientations_np[i] =  line_split[5:8] + [line_split[4]]
        
        # Create an OdometryData class
        return cls(frame_id, child_frame_id, timestamps_np, positions_np, orientations_np, frame)
    
    # =========================================================================
    # ========================= Manipulation Methods ========================== 
    # =========================================================================  
    
    def add_folded_guassian_noise_to_position(self, xy_noise_std_per_frame: float,
            z_noise_std_per_frame: float):
        """
        This method simulates odometry drift by adding folded gaussian noise
        to the odometry positions on a per frame basis. It also accumulates
        it over time. NOTE: It completely ignores the timestamps, and the "folded
        guassian noise" distribution stds might not align with the stds of the 
        guassian used internally, so this is not a robust function at all.

        Args:
            xy_noise_std_per_frame (float): Standard deviation of the gaussian 
                distribution for xy, whose output is then run through abs().
            z_noise_std_per_frame (float): Same as above, but for z.
        """

        # Track cumulative noise for each field
        cumulative_noise_pos = {'x': 0.0, 'y': 0.0, 'z': 0.0}

        # For each position
        for i in range(len(self.timestamps)):

            # Sample noise and accumulate
            noise_pos = {'x': np.random.normal(0, xy_noise_std_per_frame),
                        'y': np.random.normal(0, xy_noise_std_per_frame),
                        'z': np.random.normal(0, z_noise_std_per_frame)}
            for key in cumulative_noise_pos:
                cumulative_noise_pos[key] += abs(noise_pos[key])

            # Update positions
            self.positions[i][0] += Decimal(cumulative_noise_pos['x'])
            self.positions[i][1] += Decimal(cumulative_noise_pos['y'])
            self.positions[i][2] += Decimal(cumulative_noise_pos['z'])

    @typechecked
    def shift_position(self, x_shift: float, y_shift: float, z_shift: float):
        """
        Shifts the positions of the odometry.

        Args:
            x_shift (float): Shift in x-axis.
            y_shift (float): Shift in y_axis.
            z_shift (float): Shift in z_axis.
        """
        self.positions[:,0] += Decimal(x_shift)
        self.positions[:,1] += Decimal(y_shift)
        self.positions[:,2] += Decimal(z_shift)

    def shift_to_start_at_identity(self):
        """
        Alter the positions and orientations based so that the first pose 
        starts at Identity.
        """

        # Get pose of first robot position w.r.t world
        R_o = R.from_quat(self.orientations[0]).as_matrix()
        T_o = np.expand_dims(self.positions[0], axis=1)
        
        # Calculate the inverse (pose of world w.r.t first robot location)
        R_inv = R_o.T

        # Rotate positions and orientations
        self.positions = (R_inv @ (self.positions.T - T_o).astype(float)).T
        for i in range(self.len()):
            self.orientations[i] = R.from_matrix((R_inv @ R.from_quat(self.orientations[i]).as_matrix())).as_quat()

        # Convert back to decimal array
        self.positions = col_to_dec_arr(self.positions)
        self.orientations = col_to_dec_arr(self.orientations)

    def crop_data(self, start: Decimal, end: Decimal):
        """ Will crop the data so only values within [start, end] inclusive are kept. """

        # Create boolean mask of data to keep
        mask = (self.timestamps >= start) & (self.timestamps <= end)

        # Apply mask
        self.timestamps = self.timestamps[mask]
        self.positions = self.positions[mask]
        self.orientations = self.orientations[mask]

        # Empty poses as they might need to be recalculated
        self.poses = []

    # =========================================================================
    # ============================ Export Methods ============================= 
    # =========================================================================  

    @typechecked
    def to_csv(self, csv_path: Union[Path, str]):
        """
        Writes the odometry data to a .csv file. Note that data will be
        saved in the following order: timestamp, pos.x, pos.y, pos.z,
        ori.w, ori.x, ori.y, ori.z. Timestamp is in seconds.

        Args:
            csv_path (Path | str): Path to the output csv file.
            odom_topic (str): Topic of the Odometry messages.
        Returns:
            OdometryData: Instance of this class.
        """

        # setup tqdm 
        pbar = tqdm.tqdm(total=None, desc="Saving to csv... ", unit=" frames")

        # Check that file path doesn't already exist
        file_path = Path(csv_path)
        if os.path.exists(file_path):
            raise ValueError(f"Output file already exists: {file_path}")
        
        # Open csv file
        with open(file_path, 'w', newline='') as csvfile:
            writer = csv.writer(csvfile, delimiter=',')

            # Write the first row
            writer.writerow(['timestamp', 'x', 'y', 'z', 'qw', 'qx', 'qy', 'qz'])
                
            # Write message data to the csv file
            for i in range(len(self.timestamps)):
                writer.writerow([str(self.timestamps[i]), 
                    str(self.positions[i][0]), str(self.positions[i][1]), str(self.positions[i][2]),
                    str(self.orientations[i][3]), str(self.orientations[i][0]), str(self.orientations[i][1]), 
                    str(self.orientations[i][2])])
                pbar.update(1)

    # =========================================================================
    # =========================== Frame Conversions =========================== 
    # ========================================================================= 
    def to_FLU_frame(self):
        # If we are already in the FLU frame, return
        if self.frame == CoordinateFrame.FLU:
            print("Data already in FLU coordinate frame, returning...")
            return

        # If in NED, run the conversion
        elif self.frame == CoordinateFrame.NED:
            # Define the rotation matrix
            R_NED = np.array([[1,  0,  0],
                              [0, -1,  0],
                              [0,  0, -1]])

            # Do a change of basis to update the frame
            self._convert_frame(R_NED)

            # Update frame
            self.frame = CoordinateFrame.FLU

        # Otherwise, throw an error
        else:
            raise RuntimeError(f"OdometryData class is in an unexpected frame: {self.frame}!")
    
    @typechecked
    def _convert_frame(self, R_frame: np.ndarray):
        """ Uses a change of basis to update the positions and orientations. """
        R_frame_Q = R.from_matrix(R_frame)
        self.positions = (R_frame @ self.positions.T).T
        self._ori_change_of_basis(R_frame_Q)

    @typechecked
    def _ori_apply_rotation(self, R_i: R):
        """ Applies a rotation (not a change of basis) to orientations, thus stays in the same frame. """
        for i in range(self.len()):
            self.orientations[i] = (R_i * R.from_quat(self.orientations[i])).as_quat()

    @typechecked
    def _ori_change_of_basis(self, R_i: R):
        """ Applies a change of basis to orientations """
        for i in range(self.len()):
            self.orientations[i] = (R_i * R.from_quat(self.orientations[i]) * R_i.inv()).as_quat()

    # =========================================================================
    # =========================== Conversion to ROS =========================== 
    # ========================================================================= 

    @typechecked
    @staticmethod
    def get_ros_msg_type(msg_type: str = "Odometry") -> str:
        """ Return the __msgtype__ for an Imu msg. """
        typestore = get_typestore(Stores.ROS2_HUMBLE)
        if msg_type == "Odometry":
            return typestore.types['nav_msgs/msg/Odometry'].__msgtype__
        elif msg_type == "Path":
            return typestore.types['nav_msgs/msg/Path'].__msgtype__
        else:
            raise ValueError(f"Unsupported msg_type for OdometryData: {msg_type}")
    
    @typechecked
    def extract_seconds_and_nanoseconds(self, i: int):
        seconds = int(self.timestamps[i])
        nanoseconds = (self.timestamps[i] - self.timestamps[i].to_integral_value(rounding=decimal.ROUND_DOWN)) \
                        * Decimal("1e9").to_integral_value(decimal.ROUND_HALF_EVEN)
        return seconds, nanoseconds
    
    def calculate_stamped_poses(self):
        """ Pre-calculates all stamped poses if they aren't calculated yet."""

        if len(self.poses) != self.len():
            # Get ROS2 message types
            typestore = get_typestore(Stores.ROS2_HUMBLE)
            Header = typestore.types['std_msgs/msg/Header']
            Time = typestore.types['builtin_interfaces/msg/Time']
            PoseStamped = typestore.types['geometry_msgs/msg/PoseStamped']
            Pose = typestore.types['geometry_msgs/msg/Pose']
            Point = typestore.types['geometry_msgs/msg/Point']
            Quaternion = typestore.types['geometry_msgs/msg/Quaternion']

            # Pre-calculate all the poses
            for i in range(self.len()):
                seconds, nanoseconds = self.extract_seconds_and_nanoseconds(i)
                self.poses.append(PoseStamped(Header(stamp=Time(sec=int(seconds), 
                                                                nanosec=int(nanoseconds)),
                                                    frame_id=self.frame_id),
                                            pose=Pose(position=Point(x=self.positions[i][0],
                                                                    y=self.positions[i][1],
                                                                    z=self.positions[i][2]),
                            orientation=Quaternion(x=self.orientations[i][0],
                                                    y=self.orientations[i][1],
                                                    z=self.orientations[i][2],
                                                    w=self.orientations[i][3]))))

    @typechecked
    def get_ros_msg(self, i: int, msg_type: str = "Odometry"):
        """
        Gets an Image ROS2 Humble message corresponding to the odometry in index i.
        
        Args:
            i (int): The index of the odometry data to convert.
        Raises:
            ValueError: If i is outside the data bounds.
        """

        # Calculate the Stamped Poses
        self.calculate_stamped_poses()

        # Make sure our data is in the FLU frame, otherwise throw an error
        if self.frame != CoordinateFrame.FLU:
            raise RuntimeError("Convert this Odometry Data to a FLU frame before writing to a ROS2 bag!")

        # Check to make sure index is within data bounds
        if i < 0 or i >= self.len():
            raise IndexError(f"Index {i} is out of bounds!")

        # Get ROS2 message classes
        typestore = get_typestore(Stores.ROS2_HUMBLE)
        Odometry = typestore.types['nav_msgs/msg/Odometry']
        Header = typestore.types['std_msgs/msg/Header']
        Time = typestore.types['builtin_interfaces/msg/Time']
        PoseWithCovariance = typestore.types['geometry_msgs/msg/PoseWithCovariance']
        TwistWithCovariance = typestore.types['geometry_msgs/msg/TwistWithCovariance']
        Twist = typestore.types['geometry_msgs/msg/Twist']
        Vector3 = typestore.types['geometry_msgs/msg/Vector3']
        Path = typestore.types['nav_msgs/msg/Path']
        Pose = typestore.types['geometry_msgs/msg/Pose']
        Point = typestore.types['geometry_msgs/msg/Point']
        Quaternion = typestore.types['geometry_msgs/msg/Quaternion']

        # Write the data into the new msg
        if msg_type == "Odometry":
            seconds, nanoseconds = self.extract_seconds_and_nanoseconds(i)
            return Odometry(Header(stamp=Time(sec=int(seconds), 
                                              nanosec=int(nanoseconds)), 
                            frame_id=self.frame_id),
                            child_frame_id=self.child_frame_id,
                            pose=PoseWithCovariance(pose=Pose(position=Point(x=self.positions[i][0],
                                                                 y=self.positions[i][1],
                                                                 z=self.positions[i][2]),
                                                    orientation=Quaternion(x=self.orientations[i][0],
                                                                            y=self.orientations[i][1],
                                                                            z=self.orientations[i][2],
                                                                            w=self.orientations[i][3])),
                                                    covariance=np.zeros(36)),
                            twist=TwistWithCovariance(twist=Twist(linear=Vector3(x=0, # Currently doesn't support Twist
                                                                                y=0,
                                                                                z=0,),
                                                                angular=Vector3(x=0,
                                                                                y=0,
                                                                                z=0,)),
                                                    covariance=np.zeros(36)))
        elif msg_type == "Path":
            seconds, nanoseconds = self.extract_seconds_and_nanoseconds(i)
            return Path(Header(stamp=Time(sec=int(seconds), 
                                          nanosec=int(nanoseconds)),
                               frame_id=self.frame_id),
                               poses=self.poses[0:i+1:10])
        else:
            raise ValueError(f"Unsupported msg_type for OdometryData: {msg_type}")

lunarlab-gatech / robotdataprocess / 20172884498

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous