forked from zbw/yiliao2026
标定+USB摄像头启动
This commit is contained in:
250
scripts/export_bag_image_odom.py
Executable file
250
scripts/export_bag_image_odom.py
Executable file
@@ -0,0 +1,250 @@
|
||||
#!/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()
|
||||
Reference in New Issue
Block a user