#!/usr/bin/env python3 """Export matched /image and /odom_combined samples from a ROS2 bag.""" import argparse import bisect import math from dataclasses import dataclass from pathlib import Path from typing import Dict, Iterable, List, Optional, Sequence, Tuple DEFAULT_IMAGE_TOPIC = "/image" DEFAULT_ODOM_TOPIC = "/odom_combined" DEFAULT_OUTPUT_ROOT = Path("/home/sunrise/yiliao_ws/bag_outputs") DEFAULT_MAX_DELTA_S = 0.2 @dataclass(frozen=True) class OdomSample: stamp: float x: float y: float cos_yaw: float sin_yaw: float @dataclass(frozen=True) class ImageRecord: stamp: float msg: object msg_type: str def stamp_to_seconds(stamp) -> float: return float(stamp.sec) + float(stamp.nanosec) * 1e-9 def message_time_seconds(msg, fallback_nanoseconds: int) -> float: header = getattr(msg, "header", None) stamp = getattr(header, "stamp", None) if stamp is not None and (getattr(stamp, "sec", 0) != 0 or getattr(stamp, "nanosec", 0) != 0): return stamp_to_seconds(stamp) return fallback_nanoseconds * 1e-9 def yaw_sin_cos_from_quaternion(q) -> Tuple[float, float]: siny_cosp = 2.0 * (q.w * q.z + q.x * q.y) cosy_cosp = 1.0 - 2.0 * (q.y * q.y + q.z * q.z) yaw = math.atan2(siny_cosp, cosy_cosp) return math.cos(yaw), math.sin(yaw) def odom_sample_from_msg(stamp: float, msg) -> OdomSample: pose = msg.pose.pose cos_yaw, sin_yaw = yaw_sin_cos_from_quaternion(pose.orientation) return OdomSample( stamp=stamp, x=float(pose.position.x), y=float(pose.position.y), cos_yaw=cos_yaw, sin_yaw=sin_yaw, ) def find_nearest_odom( odom_samples: Sequence[OdomSample], image_stamp: float, max_delta_s: float, ) -> Optional[OdomSample]: if not odom_samples: return None stamps = [sample.stamp for sample in odom_samples] index = bisect.bisect_left(stamps, image_stamp) candidates = [] if index < len(odom_samples): candidates.append(odom_samples[index]) if index > 0: candidates.append(odom_samples[index - 1]) nearest = min(candidates, key=lambda sample: abs(sample.stamp - image_stamp)) if abs(nearest.stamp - image_stamp) > max_delta_s: return None return nearest def read_bag_records( bag_path: Path, image_topic: str, odom_topic: str, ) -> Tuple[List[ImageRecord], List[OdomSample]]: from rclpy.serialization import deserialize_message from rosbag2_py import ConverterOptions, SequentialReader, StorageOptions from rosidl_runtime_py.utilities import get_message reader = SequentialReader() reader.open( StorageOptions(uri=str(bag_path), storage_id="sqlite3"), ConverterOptions(input_serialization_format="cdr", output_serialization_format="cdr"), ) topic_types: Dict[str, str] = { topic_metadata.name: topic_metadata.type for topic_metadata in reader.get_all_topics_and_types() } required = [topic for topic in (image_topic, odom_topic) if topic not in topic_types] if required: available = ", ".join(sorted(topic_types)) raise ValueError(f"topic(s) missing from bag: {', '.join(required)}. Available topics: {available}") msg_classes = { image_topic: get_message(topic_types[image_topic]), odom_topic: get_message(topic_types[odom_topic]), } image_records: List[ImageRecord] = [] odom_samples: List[OdomSample] = [] while reader.has_next(): topic, data, bag_time_ns = reader.read_next() if topic not in msg_classes: continue msg = deserialize_message(data, msg_classes[topic]) stamp = message_time_seconds(msg, bag_time_ns) if topic == image_topic: image_records.append(ImageRecord(stamp=stamp, msg=msg, msg_type=topic_types[image_topic])) elif topic == odom_topic: odom_samples.append(odom_sample_from_msg(stamp, msg)) odom_samples.sort(key=lambda sample: sample.stamp) image_records.sort(key=lambda record: record.stamp) return image_records, odom_samples def bag_output_name(bag_path: Path) -> str: return bag_path.resolve().name def ensure_clean_output_dirs(image_dir: Path, odom_dir: Path, overwrite: bool) -> None: image_dir.mkdir(parents=True, exist_ok=True) odom_dir.mkdir(parents=True, exist_ok=True) existing = list(image_dir.glob("*.png")) + list(odom_dir.glob("*.txt")) if existing and not overwrite: raise FileExistsError( f"output already contains exported files: {image_dir} or {odom_dir}. " "Use --overwrite to replace numbered outputs." ) if overwrite: for path in existing: path.unlink() def write_image_png(msg, msg_type: str, path: Path) -> None: import cv2 import numpy as np if msg_type == "sensor_msgs/msg/CompressedImage": encoded = np.frombuffer(msg.data, dtype=np.uint8) image = cv2.imdecode(encoded, cv2.IMREAD_UNCHANGED) if image is None: raise ValueError("failed to decode compressed image") elif msg_type == "sensor_msgs/msg/Image": from cv_bridge import CvBridge image = CvBridge().imgmsg_to_cv2(msg, desired_encoding="passthrough") else: raise ValueError(f"unsupported image message type: {msg_type}") if not cv2.imwrite(str(path), image): raise IOError(f"failed to write image: {path}") def write_odom_txt(sample: OdomSample, path: Path) -> None: path.write_text( f"{sample.x:.9f} {sample.y:.9f} {sample.cos_yaw:.9f} {sample.sin_yaw:.9f}\n", encoding="utf-8", ) def export_matched_records( image_records: Iterable[ImageRecord], odom_samples: Sequence[OdomSample], image_dir: Path, odom_dir: Path, max_delta_s: float, ) -> Tuple[int, int]: exported = 0 dropped = 0 for record in image_records: odom = find_nearest_odom(odom_samples, record.stamp, max_delta_s) if odom is None: dropped += 1 continue exported += 1 stem = f"{exported:04d}" write_image_png(record.msg, record.msg_type, image_dir / f"{stem}.png") write_odom_txt(odom, odom_dir / f"{stem}.txt") return exported, dropped def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="Export /image frames and nearest /odom_combined poses from a ROS2 bag." ) parser.add_argument("bag_path", type=Path, help="Path to the ROS2 bag directory.") parser.add_argument("--image-topic", default=DEFAULT_IMAGE_TOPIC, help=f"Image topic, default: {DEFAULT_IMAGE_TOPIC}") parser.add_argument("--odom-topic", default=DEFAULT_ODOM_TOPIC, help=f"Odom topic, default: {DEFAULT_ODOM_TOPIC}") parser.add_argument( "--output-root", type=Path, default=DEFAULT_OUTPUT_ROOT, help=f"Output root, default: {DEFAULT_OUTPUT_ROOT}", ) parser.add_argument( "--max-delta", type=float, default=DEFAULT_MAX_DELTA_S, help=f"Maximum image/odom timestamp difference in seconds, default: {DEFAULT_MAX_DELTA_S}", ) parser.add_argument("--overwrite", action="store_true", help="Delete existing numbered png/txt files first.") return parser.parse_args() def main() -> None: args = parse_args() bag_path = args.bag_path.expanduser().resolve() if not bag_path.exists(): raise FileNotFoundError(f"bag path does not exist: {bag_path}") output_name = bag_output_name(bag_path) image_dir = args.output_root / "images" / output_name odom_dir = args.output_root / "odom" / output_name ensure_clean_output_dirs(image_dir, odom_dir, args.overwrite) image_records, odom_samples = read_bag_records(bag_path, args.image_topic, args.odom_topic) exported, dropped = export_matched_records( image_records=image_records, odom_samples=odom_samples, image_dir=image_dir, odom_dir=odom_dir, max_delta_s=args.max_delta, ) print(f"bag: {bag_path}") print(f"images read: {len(image_records)}") print(f"odom read: {len(odom_samples)}") print(f"exported pairs: {exported}") print(f"dropped images: {dropped}") print(f"image output: {image_dir}") print(f"odom output: {odom_dir}") if __name__ == "__main__": main()