generate_ros2_scan_map.py

#!/usr/bin/env python3

"""Generate a 2D occupancy map from a time window in a ROS 2 bag.

The script reads LaserScan, Odometry, and TF data directly from a rosbag2
directory using ``rosbags``. Laser beams are projected into a fixed frame using
per-beam timestamps so scans collected while the robot is moving still land in a
consistent map. The output is a Nav2/map_server compatible ``.pgm`` plus
``.yaml`` sidecar.
"""

from __future__ import annotations

import argparse
from bisect import bisect_right
from collections import defaultdict, deque
from dataclasses import dataclass
import math
from pathlib import Path
from typing import Iterable

from rosbags.highlevel import AnyReader


NSEC_PER_SEC = 1_000_000_000
FREE_VALUE = 254
OCCUPIED_VALUE = 0
UNKNOWN_VALUE = 205


@dataclass(frozen=True)
class Pose2D:
    x: float
    y: float
    yaw: float


@dataclass(frozen=True)
class PoseSample:
    stamp_ns: int
    pose: Pose2D


@dataclass(frozen=True)
class ScanWindow:
    recorded_stamp_ns: int
    scan_msg: object


@dataclass(frozen=True)
class Ray:
    origin_x: float
    origin_y: float
    hit_x: float
    hit_y: float


class PoseBuffer:
    """Timestamp-sorted 2D poses with linear interpolation."""

    def __init__(self, samples: Iterable[PoseSample]):
        ordered = sorted(samples, key=lambda sample: sample.stamp_ns)
        if not ordered:
            raise ValueError('PoseBuffer requires at least one sample.')
        self.samples = ordered
        self.stamps = [sample.stamp_ns for sample in ordered]

    def interpolate(self, stamp_ns: int) -> Pose2D:
        if len(self.samples) == 1:
            return self.samples[0].pose

        if stamp_ns <= self.stamps[0]:
            return self.samples[0].pose
        if stamp_ns >= self.stamps[-1]:
            return self.samples[-1].pose

        upper_index = bisect_right(self.stamps, stamp_ns)
        first = self.samples[upper_index - 1]
        second = self.samples[upper_index]

        dt = second.stamp_ns - first.stamp_ns
        if dt <= 0:
            return second.pose

        ratio = float(stamp_ns - first.stamp_ns) / float(dt)
        return Pose2D(
            x=first.pose.x + ratio * (second.pose.x - first.pose.x),
            y=first.pose.y + ratio * (second.pose.y - first.pose.y),
            yaw=normalize_angle(first.pose.yaw + ratio * normalize_angle(second.pose.yaw - first.pose.yaw)),
        )


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description='Generate a Nav2-compatible occupancy map from a ROS 2 bag time window.',
    )
    parser.add_argument(
        'bag_path',
        nargs='?',
        default='2020-02-16-11-52-35_ros2',
        help='Path to the rosbag2 directory.',
    )
    parser.add_argument(
        '--start-offset',
        type=float,
        default=0.0,
        help='Window start offset in seconds from bag start time.',
    )
    parser.add_argument(
        '--end-offset',
        type=float,
        default=9.0,
        help='Window end offset in seconds from bag start time.',
    )
    parser.add_argument(
        '--scan-topic',
        default='/scan',
        help='LaserScan topic to map.',
    )
    parser.add_argument(
        '--odom-topic',
        default='/odom',
        help='Odometry topic used for the odom-to-base pose.',
    )
    parser.add_argument(
        '--tf-topic',
        default='/tf',
        help='TF topic used for the scan-frame extrinsics.',
    )
    parser.add_argument(
        '--fixed-frame',
        default='odom',
        help='Fixed frame for the generated map.',
    )
    parser.add_argument(
        '--resolution',
        type=float,
        default=0.05,
        help='Map resolution in meters per cell.',
    )
    parser.add_argument(
        '--padding',
        type=float,
        default=0.5,
        help='Extra map border in meters around all observed points.',
    )
    parser.add_argument(
        '--hit-weight',
        type=float,
        default=3.0,
        help='Relative weight for occupied endpoints versus free ray traversals.',
    )
    parser.add_argument(
        '--output-prefix',
        type=Path,
        help='Output path without extension. Defaults next to the bag directory.',
    )
    return parser.parse_args()


def stamp_to_ns(stamp) -> int:
    return int(stamp.sec) * NSEC_PER_SEC + int(stamp.nanosec)


def normalize_angle(theta: float) -> float:
    two_pi = 2.0 * math.pi
    normalized = theta % two_pi
    normalized = (normalized + two_pi) % two_pi
    if normalized > math.pi:
        normalized -= two_pi
    return normalized


def yaw_from_quaternion(quaternion) -> float:
    siny_cosp = 2.0 * (quaternion.w * quaternion.z + quaternion.x * quaternion.y)
    cosy_cosp = 1.0 - 2.0 * (quaternion.y * quaternion.y + quaternion.z * quaternion.z)
    return normalize_angle(math.atan2(siny_cosp, cosy_cosp))


def compose_pose(first: Pose2D, second: Pose2D) -> Pose2D:
    cos_yaw = math.cos(first.yaw)
    sin_yaw = math.sin(first.yaw)
    return Pose2D(
        x=first.x + cos_yaw * second.x - sin_yaw * second.y,
        y=first.y + sin_yaw * second.x + cos_yaw * second.y,
        yaw=normalize_angle(first.yaw + second.yaw),
    )


def invert_pose(pose: Pose2D) -> Pose2D:
    cos_yaw = math.cos(pose.yaw)
    sin_yaw = math.sin(pose.yaw)
    return Pose2D(
        x=-(cos_yaw * pose.x + sin_yaw * pose.y),
        y=-(-sin_yaw * pose.x + cos_yaw * pose.y),
        yaw=normalize_angle(-pose.yaw),
    )


def recorded_start_ns(reader: AnyReader) -> int:
    for _connection, timestamp_ns, _rawdata in reader.messages():
        return timestamp_ns
    raise ValueError('Bag is empty.')


def find_path(edge_buffers: dict[tuple[str, str], PoseBuffer], source_frame: str, target_frame: str) -> list[tuple[str, str, bool]]:
    if source_frame == target_frame:
        return []

    adjacency: dict[str, list[tuple[str, tuple[str, str], bool]]] = defaultdict(list)
    for parent, child in edge_buffers:
        adjacency[parent].append((child, (parent, child), True))
        adjacency[child].append((parent, (parent, child), False))

    queue = deque([(source_frame, [])])
    visited = {source_frame}
    while queue:
        frame, path = queue.popleft()
        for neighbor, edge_key, forward in adjacency.get(frame, []):
            if neighbor in visited:
                continue
            next_path = path + [(edge_key[0], edge_key[1], forward)]
            if neighbor == target_frame:
                return next_path
            visited.add(neighbor)
            queue.append((neighbor, next_path))

    available_frames = ', '.join(sorted(adjacency))
    raise ValueError(
        f'No transform path from {source_frame!r} to {target_frame!r}. Available frames: {available_frames}'
    )


def resolve_pose_along_path(
    edge_buffers: dict[tuple[str, str], PoseBuffer],
    path: list[tuple[str, str, bool]],
    stamp_ns: int,
) -> Pose2D:
    pose = Pose2D(0.0, 0.0, 0.0)
    for parent, child, forward in path:
        edge_pose = edge_buffers[(parent, child)].interpolate(stamp_ns)
        pose = compose_pose(pose, edge_pose if forward else invert_pose(edge_pose))
    return pose


def message_time_ns(message, fallback_recorded_ns: int) -> int:
    header_stamp = getattr(getattr(message, 'header', None), 'stamp', None)
    if header_stamp is None:
        return fallback_recorded_ns
    stamp_ns = stamp_to_ns(header_stamp)
    return stamp_ns if stamp_ns != 0 else fallback_recorded_ns


def read_inputs(
    bag_path: Path,
    scan_topic: str,
    odom_topic: str,
    tf_topic: str,
    start_offset_sec: float,
    end_offset_sec: float,
) -> tuple[int, list[ScanWindow], dict[tuple[str, str], PoseBuffer]]:
    if end_offset_sec <= start_offset_sec:
        raise ValueError('--end-offset must be greater than --start-offset.')

    edge_samples: dict[tuple[str, str], list[PoseSample]] = defaultdict(list)
    selected_scans: list[ScanWindow] = []

    with AnyReader([bag_path]) as reader:
        connections_by_topic = defaultdict(list)
        for connection in reader.connections:
            connections_by_topic[connection.topic].append(connection)

        if scan_topic not in connections_by_topic:
            available_topics = ', '.join(sorted(connections_by_topic))
            raise ValueError(f'Scan topic {scan_topic!r} not found. Available topics: {available_topics}')
        if odom_topic not in connections_by_topic:
            available_topics = ', '.join(sorted(connections_by_topic))
            raise ValueError(f'Odom topic {odom_topic!r} not found. Available topics: {available_topics}')

        bag_start_ns = recorded_start_ns(reader)
        window_start_ns = bag_start_ns + int(start_offset_sec * NSEC_PER_SEC)
        window_end_ns = bag_start_ns + int(end_offset_sec * NSEC_PER_SEC)

        wanted_topics = {scan_topic, odom_topic, tf_topic}
        wanted_connections = [
            connection
            for connection in reader.connections
            if connection.topic in wanted_topics
        ]

        for connection, recorded_ns, rawdata in reader.messages(connections=wanted_connections):
            msg = reader.deserialize(rawdata, connection.msgtype)

            if connection.topic == scan_topic:
                if window_start_ns <= recorded_ns <= window_end_ns:
                    selected_scans.append(ScanWindow(recorded_stamp_ns=recorded_ns, scan_msg=msg))
                continue

            if connection.topic == odom_topic:
                pose = msg.pose.pose
                edge_key = (msg.header.frame_id, msg.child_frame_id)
                edge_samples[edge_key].append(
                    PoseSample(
                        stamp_ns=message_time_ns(msg, recorded_ns),
                        pose=Pose2D(
                            x=float(pose.position.x),
                            y=float(pose.position.y),
                            yaw=yaw_from_quaternion(pose.orientation),
                        ),
                    )
                )
                continue

            if connection.topic == tf_topic:
                for transform in msg.transforms:
                    edge_key = (transform.header.frame_id, transform.child_frame_id)
                    edge_samples[edge_key].append(
                        PoseSample(
                            stamp_ns=message_time_ns(transform, recorded_ns),
                            pose=Pose2D(
                                x=float(transform.transform.translation.x),
                                y=float(transform.transform.translation.y),
                                yaw=yaw_from_quaternion(transform.transform.rotation),
                            ),
                        )
                    )

    if not selected_scans:
        raise ValueError('No scan messages fell within the requested bag-time window.')

    edge_buffers = {edge_key: PoseBuffer(samples) for edge_key, samples in edge_samples.items()}
    return bag_start_ns, selected_scans, edge_buffers


def build_rays(
    scans: list[ScanWindow],
    edge_buffers: dict[tuple[str, str], PoseBuffer],
    fixed_frame: str,
) -> tuple[list[Ray], str]:
    scan_frame = scans[0].scan_msg.header.frame_id
    if not scan_frame:
        raise ValueError('LaserScan header.frame_id is empty.')

    path = find_path(edge_buffers, fixed_frame, scan_frame)
    rays: list[Ray] = []

    for scan_window in scans:
        scan = scan_window.scan_msg
        if scan.header.frame_id != scan_frame:
            raise ValueError(
                f'All scans in the window must share a frame_id. Saw {scan_frame!r} and {scan.header.frame_id!r}.'
            )

        scan_start_ns = stamp_to_ns(scan.header.stamp)
        time_increment_ns = int(scan.time_increment * NSEC_PER_SEC)

        for index, laser_range in enumerate(scan.ranges):
            if not math.isfinite(laser_range):
                continue
            if laser_range < scan.range_min or laser_range > scan.range_max:
                continue

            beam_time_ns = scan_start_ns + index * time_increment_ns
            laser_pose = resolve_pose_along_path(edge_buffers, path, beam_time_ns)
            beam_angle = laser_pose.yaw + scan.angle_min + index * scan.angle_increment
            hit_x = laser_pose.x + laser_range * math.cos(beam_angle)
            hit_y = laser_pose.y + laser_range * math.sin(beam_angle)
            rays.append(
                Ray(
                    origin_x=laser_pose.x,
                    origin_y=laser_pose.y,
                    hit_x=hit_x,
                    hit_y=hit_y,
                )
            )

    if not rays:
        raise ValueError('No valid scan rays were found in the requested time window.')

    return rays, scan_frame


def grid_dimensions(rays: list[Ray], resolution: float, padding: float) -> tuple[float, float, int, int]:
    xs = [coordinate for ray in rays for coordinate in (ray.origin_x, ray.hit_x)]
    ys = [coordinate for ray in rays for coordinate in (ray.origin_y, ray.hit_y)]

    min_x = min(xs) - padding
    min_y = min(ys) - padding
    max_x = max(xs) + padding
    max_y = max(ys) + padding

    width = int(math.ceil((max_x - min_x) / resolution)) + 1
    height = int(math.ceil((max_y - min_y) / resolution)) + 1
    return min_x, min_y, width, height


def world_to_cell(x: float, y: float, min_x: float, min_y: float, resolution: float) -> tuple[int, int]:
    return int(math.floor((x - min_x) / resolution)), int(math.floor((y - min_y) / resolution))


def bresenham(x0: int, y0: int, x1: int, y1: int):
    dx = abs(x1 - x0)
    sx = 1 if x0 < x1 else -1
    dy = -abs(y1 - y0)
    sy = 1 if y0 < y1 else -1
    error = dx + dy

    while True:
        yield x0, y0
        if x0 == x1 and y0 == y1:
            return
        twice_error = 2 * error
        if twice_error >= dy:
            error += dy
            x0 += sx
        if twice_error <= dx:
            error += dx
            y0 += sy


def rasterize_map(
    rays: list[Ray],
    resolution: float,
    padding: float,
    hit_weight: float,
) -> tuple[bytearray, float, float, int, int]:
    min_x, min_y, width, height = grid_dimensions(rays, resolution, padding)
    free_counts = [0] * (width * height)
    hit_counts = [0] * (width * height)

    for ray in rays:
        origin = world_to_cell(ray.origin_x, ray.origin_y, min_x, min_y, resolution)
        hit = world_to_cell(ray.hit_x, ray.hit_y, min_x, min_y, resolution)

        cells = list(bresenham(origin[0], origin[1], hit[0], hit[1]))
        for cell_x, cell_y in cells[:-1]:
            if 0 <= cell_x < width and 0 <= cell_y < height:
                free_counts[cell_y * width + cell_x] += 1

        hit_x, hit_y = cells[-1]
        if 0 <= hit_x < width and 0 <= hit_y < height:
            hit_counts[hit_y * width + hit_x] += 1

    image = bytearray(width * height)
    for index in range(width * height):
        hits = hit_counts[index]
        frees = free_counts[index]
        if hits > 0 and hit_weight * hits >= frees:
            image[index] = OCCUPIED_VALUE
        elif frees > 0:
            image[index] = FREE_VALUE
        else:
            image[index] = UNKNOWN_VALUE

    return image, min_x, min_y, width, height


def write_pgm(path: Path, image: bytearray, width: int, height: int) -> None:
    with path.open('wb') as handle:
        handle.write(f'P5\n{width} {height}\n255\n'.encode('ascii'))
        for row in range(height - 1, -1, -1):
            start = row * width
            handle.write(image[start:start + width])


def write_yaml(path: Path, image_name: str, resolution: float, origin_x: float, origin_y: float) -> None:
    yaml_text = (
        f'image: {image_name}\n'
        f'mode: trinary\n'
        f'resolution: {resolution:.6f}\n'
        f'origin: [{origin_x:.6f}, {origin_y:.6f}, 0.000000]\n'
        f'negate: 0\n'
        f'occupied_thresh: 0.65\n'
        f'free_thresh: 0.196\n'
    )
    path.write_text(yaml_text, encoding='utf-8')


def default_output_prefix(bag_path: Path, start_offset_sec: float, end_offset_sec: float) -> Path:
    stem = bag_path.name.rstrip('/')
    window_label = f'{start_offset_sec:g}s_to_{end_offset_sec:g}s'
    return bag_path.parent / f'{stem}_{window_label}_odom_map'


def main() -> None:
    args = parse_args()

    bag_path = Path(args.bag_path)
    if not bag_path.exists():
        raise FileNotFoundError(f'Bag path {bag_path} does not exist.')
    if args.resolution <= 0.0:
        raise ValueError('--resolution must be positive.')
    if args.padding < 0.0:
        raise ValueError('--padding must be non-negative.')
    if args.hit_weight <= 0.0:
        raise ValueError('--hit-weight must be positive.')

    bag_start_ns, scans, edge_buffers = read_inputs(
        bag_path=bag_path,
        scan_topic=args.scan_topic,
        odom_topic=args.odom_topic,
        tf_topic=args.tf_topic,
        start_offset_sec=args.start_offset,
        end_offset_sec=args.end_offset,
    )
    rays, scan_frame = build_rays(scans, edge_buffers, args.fixed_frame)
    image, origin_x, origin_y, width, height = rasterize_map(
        rays=rays,
        resolution=args.resolution,
        padding=args.padding,
        hit_weight=args.hit_weight,
    )

    output_prefix = args.output_prefix or default_output_prefix(bag_path, args.start_offset, args.end_offset)
    output_prefix.parent.mkdir(parents=True, exist_ok=True)
    pgm_path = output_prefix.with_suffix('.pgm')
    yaml_path = output_prefix.with_suffix('.yaml')

    write_pgm(pgm_path, image, width, height)
    write_yaml(yaml_path, pgm_path.name, args.resolution, origin_x, origin_y)

    occupied_cells = sum(1 for value in image if value == OCCUPIED_VALUE)
    free_cells = sum(1 for value in image if value == FREE_VALUE)
    unknown_cells = len(image) - occupied_cells - free_cells
    window_start_sec = (scans[0].recorded_stamp_ns - bag_start_ns) / NSEC_PER_SEC
    window_end_sec = (scans[-1].recorded_stamp_ns - bag_start_ns) / NSEC_PER_SEC

    print(f'Bag: {bag_path}')
    print(f'Fixed frame: {args.fixed_frame}')
    print(f'Scan frame: {scan_frame}')
    print(f'Selected scans: {len(scans)}')
    print(f'Projected rays: {len(rays)}')
    print(f'Bag-time window: {window_start_sec:.3f}s to {window_end_sec:.3f}s')
    print(f'Map size: {width} x {height} cells at {args.resolution:.3f} m/cell')
    print(f'Occupied/free/unknown cells: {occupied_cells}/{free_cells}/{unknown_cells}')
    print(f'Wrote {pgm_path}')
    print(f'Wrote {yaml_path}')


if __name__ == '__main__':
    main()