correct_odometry_bag_ros2.py

#!/usr/bin/env python3

"""Write a copy of a ROS 2 bag with direction-specific odometry yaw correction.

This is an offline calibration utility, not a live odometry filter.  It reads
the complete ``/odom`` stream first, unwraps its successive planar-yaw changes,
and applies an independently calibrated scale to counterclockwise and clockwise
changes.  The resulting corrected yaw is continuous, so it remains meaningful
even when the platform makes several complete turns.

The correction is intentionally limited to orientation.  Every topic,
connection, recorded timestamp, and non-orientation field is copied from the
input bag unchanged.  Both ``/odom`` and the matching ``odom -> base_link``
transform in ``/tf`` receive the same yaw correction; updating only one of
them would leave RViz and other TF consumers with contradictory poses.

The default direction-specific scales come from the raw-lidar ICP measurements
reported in ``odometry_lidar_analysis_report.md``.  They compensate for the
measured accumulated heading error in this recorded dataset and are not a
general replacement for estimating an odometry system's calibration.
"""

from __future__ import annotations

import argparse
from bisect import bisect_right
from dataclasses import dataclass
from math import atan2, cos, pi, sin
from pathlib import Path

from rosbags.highlevel import AnyReader
from rosbags.rosbag2 import Writer


NSEC_PER_SEC = 1_000_000_000


@dataclass(frozen=True)
class YawSample:
    """One odometry sample in both raw and continuously corrected yaw space."""

    stamp_ns: int
    raw_yaw: float
    corrected_yaw: float


class YawCorrection:
    """Look up the continuously varying yaw correction at an arbitrary time.

    TF and odometry streams are not guaranteed to publish at the same times.
    Storing ``corrected_yaw - raw_yaw`` at odometry samples and interpolating
    that *difference* lets the correction be applied consistently to TF
    messages without assuming their timestamps align with ``/odom``.
    """

    def __init__(self, samples: list[YawSample]):
        if len(samples) < 2:
            raise ValueError('Need at least two odometry messages to construct a yaw correction.')
        self.samples = samples
        self.stamps = [sample.stamp_ns for sample in samples]

    def corrected_yaw(self, stamp_ns: int, raw_yaw: float) -> float:
        # Outside odometry coverage, retain the endpoint correction rather than
        # extrapolating a turn rate that the recorded data does not establish.
        if stamp_ns <= self.stamps[0]:
            correction = self.samples[0].corrected_yaw - self.samples[0].raw_yaw
        elif stamp_ns >= self.stamps[-1]:
            correction = self.samples[-1].corrected_yaw - self.samples[-1].raw_yaw
        else:
            # ``bisect_right`` brackets the requested timestamp by adjacent
            # odometry samples.  Interpolating only the correction preserves
            # the message's own raw yaw while smoothly varying its adjustment.
            upper = bisect_right(self.stamps, stamp_ns)
            first = self.samples[upper - 1]
            second = self.samples[upper]
            ratio = (stamp_ns - first.stamp_ns) / (second.stamp_ns - first.stamp_ns)
            first_correction = first.corrected_yaw - first.raw_yaw
            second_correction = second.corrected_yaw - second.raw_yaw
            correction = first_correction + ratio * (second_correction - first_correction)
        return raw_yaw + correction


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description='Create a ROS 2 bag with direction-specific yaw correction in /odom and /tf.',
    )
    parser.add_argument('input_bag', type=Path, help='Input rosbag2 directory.')
    parser.add_argument('output_bag', type=Path, help='New output rosbag2 directory; it must not exist.')
    parser.add_argument('--odom-topic', default='/odom', help='Odometry topic to correct.')
    parser.add_argument('--tf-topic', default='/tf', help='TF topic containing the odom-to-base transform.')
    parser.add_argument('--parent-frame', default='odom', help='Parent frame of the corrected TF transform.')
    parser.add_argument('--child-frame', default='base_link', help='Child frame of the corrected TF transform.')
    parser.add_argument(
        '--ccw-scale', type=float, default=0.99909,
        help='Scale for positive (counterclockwise) incremental yaw; default derives from scan matching.',
    )
    parser.add_argument(
        '--cw-scale', type=float, default=0.99691,
        help='Scale for negative (clockwise) incremental yaw magnitude; default derives from scan matching.',
    )
    return parser.parse_args()


def normalize_angle(angle: float) -> float:
    """Wrap an angle to [-pi, pi), the shortest signed angular displacement."""

    return (angle + pi) % (2.0 * pi) - pi


def yaw_from_quaternion(quaternion) -> float:
    """Extract planar Z-axis yaw from a ROS quaternion.

    The source poses are treated as planar.  Roll and pitch are therefore not
    preserved when the corrected orientation is written back below.
    """

    return atan2(
        2.0 * (quaternion.w * quaternion.z + quaternion.x * quaternion.y),
        1.0 - 2.0 * (quaternion.y * quaternion.y + quaternion.z * quaternion.z),
    )


def set_planar_yaw(quaternion, yaw: float) -> None:
    """Replace a quaternion with the zero-roll, zero-pitch orientation at yaw."""

    quaternion.x = 0.0
    quaternion.y = 0.0
    quaternion.z = sin(yaw / 2.0)
    quaternion.w = cos(yaw / 2.0)


def message_stamp_ns(message, fallback_ns: int) -> int:
    """Prefer a message header timestamp, falling back to the bag record time.

    A TF transform has its own header, while ordinary topic messages generally
    use the containing message header.  A zero stamp is treated as unavailable
    because it cannot locate the message on the correction timeline.
    """

    stamp = getattr(getattr(message, 'header', None), 'stamp', None)
    if stamp is None:
        return fallback_ns
    stamp_ns = int(stamp.sec) * NSEC_PER_SEC + int(stamp.nanosec)
    return stamp_ns if stamp_ns else fallback_ns


def build_yaw_correction(input_bag: Path, odom_topic: str, ccw_scale: float, cw_scale: float) -> YawCorrection:
    """Build an unwrapped, direction-scaled yaw timeline from ``/odom``.

    Quaternion yaw is inherently wrapped at +/- pi.  Normalizing each adjacent
    difference before accumulation converts it into the local signed turn,
    avoiding a false near-2*pi jump whenever the heading crosses that boundary.
    """

    raw_samples: list[tuple[int, float]] = []
    with AnyReader([input_bag]) as reader:
        connections = [connection for connection in reader.connections if connection.topic == odom_topic]
        if not connections:
            available = ', '.join(sorted({connection.topic for connection in reader.connections}))
            raise ValueError(f'Odometry topic {odom_topic!r} not found. Available topics: {available}')
        for connection, recorded_ns, rawdata in reader.messages(connections=connections):
            message = reader.deserialize(rawdata, connection.msgtype)
            raw_samples.append((message_stamp_ns(message, recorded_ns), yaw_from_quaternion(message.pose.pose.orientation)))

    # Anchor the corrected track to the first raw pose.  Calibration changes
    # accumulated rotation, not the chosen absolute heading at bag start.
    corrected_samples = [YawSample(raw_samples[0][0], raw_samples[0][1], raw_samples[0][1])]
    for (stamp_ns, raw_yaw), (_, previous_raw_yaw) in zip(raw_samples[1:], raw_samples):
        previous = corrected_samples[-1]
        delta = normalize_angle(raw_yaw - previous_raw_yaw)
        # ``delta`` is signed: positive is CCW and negative is CW in ROS's
        # standard planar convention.  The scale is applied before accumulation
        # so a long sequence of turns remains continuous beyond +/- pi.
        scale = ccw_scale if delta >= 0.0 else cw_scale
        corrected_samples.append(YawSample(stamp_ns, raw_yaw, previous.corrected_yaw + delta * scale))
    return YawCorrection(corrected_samples)


def write_corrected_bag(args: argparse.Namespace, correction: YawCorrection) -> tuple[int, int]:
    """Copy the input bag while replacing only the two coordinated yaw fields."""

    odom_updates = 0
    tf_updates = 0
    with AnyReader([args.input_bag]) as reader, Writer(args.output_bag, version=9) as writer:
        output_connections = {
            # Recreate every source connection before copying messages.  The
            # source connection id is used as the key because it distinguishes
            # otherwise identical topic/type connections in a rosbag2 file.
            connection.id: writer.add_connection(
                connection.topic,
                connection.msgtype,
                msgdef=connection.msgdef.data,
                rihs01=connection.digest,
                serialization_format=connection.ext.serialization_format,
                offered_qos_profiles=(),
            )
            for connection in reader.connections
        }
        for connection, recorded_ns, rawdata in reader.messages():
            if connection.topic == args.odom_topic:
                # Deserialize and reserialize only messages whose payload must
                # change.  All other raw CDR payloads pass through byte-for-byte.
                message = reader.deserialize(rawdata, connection.msgtype)
                quaternion = message.pose.pose.orientation
                set_planar_yaw(quaternion, correction.corrected_yaw(message_stamp_ns(message, recorded_ns), yaw_from_quaternion(quaternion)))
                rawdata = reader.typestore.serialize_cdr(message, connection.msgtype)
                odom_updates += 1
            elif connection.topic == args.tf_topic:
                message = reader.deserialize(rawdata, connection.msgtype)
                changed = False
                for transform_stamped in message.transforms:
                    # A TFMessage may bundle unrelated transforms, so alter
                    # only the explicitly selected parent/child frame pair.
                    if (
                        transform_stamped.header.frame_id == args.parent_frame
                        and transform_stamped.child_frame_id == args.child_frame
                    ):
                        quaternion = transform_stamped.transform.rotation
                        set_planar_yaw(
                            quaternion,
                            correction.corrected_yaw(message_stamp_ns(transform_stamped, recorded_ns), yaw_from_quaternion(quaternion)),
                        )
                        changed = True
                        tf_updates += 1
                if changed:
                    rawdata = reader.typestore.serialize_cdr(message, connection.msgtype)
            writer.write(output_connections[connection.id], recorded_ns, rawdata)
    return odom_updates, tf_updates


def main() -> None:
    args = parse_args()
    if not args.input_bag.is_dir():
        raise SystemExit(f'Input bag directory does not exist: {args.input_bag}')
    if args.output_bag.exists():
        raise SystemExit(f'Output bag path already exists: {args.output_bag}')
    if args.ccw_scale <= 0.0 or args.cw_scale <= 0.0:
        raise SystemExit('Yaw scales must be positive.')

    correction = build_yaw_correction(args.input_bag, args.odom_topic, args.ccw_scale, args.cw_scale)
    odom_updates, tf_updates = write_corrected_bag(args, correction)
    print(f'Wrote corrected bag: {args.output_bag}')
    print(f'Applied scales: CCW={args.ccw_scale:.8f}, CW={args.cw_scale:.8f}')
    print(f'Updated orientations: /odom={odom_updates}, {args.parent_frame} -> {args.child_frame} TF={tf_updates}')


if __name__ == '__main__':
    main()