Brainrot_Muxa/flyguard/cdr.py

"""Разбор sensor_msgs/msg/PointCloud2 из CDR без зависимости от ROS.

Нужен для двух сценариев:
  * офлайн-эксперименты на машине без ROS (Windows);
  * прямое чтение rosbag внутри контейнера, минуя `ros2 bag play`.

Внутри ROS-ноды сообщение приходит уже разобранным; его переводит в тот же
`PointCloud2` функция `from_ros_message`.
"""
from __future__ import annotations

import struct
from dataclasses import dataclass

import numpy as np

# sensor_msgs/msg/PointField: код типа -> (numpy dtype, размер в байтах)
_PF_DTYPES = {
    1: ("i1", 1), 2: ("u1", 1), 3: ("i2", 2), 4: ("u2", 2),
    5: ("i4", 4), 6: ("u4", 4), 7: ("f4", 4), 8: ("f8", 8),
}


@dataclass(frozen=True)
class PointCloud2:
    """Минимальное представление облака точек."""

    stamp: float
    frame_id: str
    height: int
    width: int
    point_step: int
    is_dense: bool
    points: np.ndarray  # структурированный массив длиной height*width

    @property
    def n_points(self) -> int:
        return int(self.points.shape[0])


class _CdrReader:
    """Чтение little-endian CDR с выравниванием примитивов относительно тела сообщения."""

    __slots__ = ("buf", "origin", "pos")

    def __init__(self, buf: bytes | memoryview):
        self.buf = buf
        self.origin = 4  # заголовок инкапсуляции
        self.pos = 4

    def _align(self, size: int) -> None:
        self.pos += (-(self.pos - self.origin)) % size

    def u8(self) -> int:
        v = self.buf[self.pos]
        self.pos += 1
        return v

    def u32(self) -> int:
        self._align(4)
        v = struct.unpack_from("<I", self.buf, self.pos)[0]
        self.pos += 4
        return v

    def i32(self) -> int:
        self._align(4)
        v = struct.unpack_from("<i", self.buf, self.pos)[0]
        self.pos += 4
        return v

    def string(self) -> str:
        n = self.u32()
        s = bytes(self.buf[self.pos:self.pos + max(n - 1, 0)]).decode("utf-8", "replace")
        self.pos += n
        return s


def point_dtype(fields: list[tuple[str, int, int, int]], point_step: int) -> np.dtype:
    """Собрать numpy-dtype по описанию полей, явно добивая пропуски паддингом.

    Поля лидара невыровнены (`timestamp` float64 по смещению 18), поэтому
    структурированный dtype строится вручную, а не через `np.dtype(align=True)`.
    """
    spec: list[tuple[str, str]] = []
    used = 0
    for name, offset, datatype, count in fields:
        kind, size = _PF_DTYPES[datatype]
        if offset > used:
            spec.append((f"_pad{used}", f"V{offset - used}"))
        elif offset < used:
            raise ValueError(f"перекрывающиеся поля в PointCloud2: {name}")
        spec.append((name, kind if count == 1 else f"{count}{kind}"))
        used = offset + size * count
    if point_step > used:
        spec.append((f"_pad{used}", f"V{point_step - used}"))
    dt = np.dtype(spec)
    if dt.itemsize != point_step:
        raise ValueError(f"dtype {dt.itemsize} байт != point_step {point_step}")
    return dt


def parse_pointcloud2(blob: bytes | memoryview) -> PointCloud2:
    """Разобрать CDR-сериализованное sensor_msgs/msg/PointCloud2."""
    r = _CdrReader(blob)
    sec = r.i32()
    nsec = r.u32()
    frame_id = r.string()
    height = r.u32()
    width = r.u32()

    fields = []
    for _ in range(r.u32()):
        name = r.string()
        offset = r.u32()
        datatype = r.u8()
        count = r.u32()
        fields.append((name, offset, datatype, count))

    r.u8()  # is_bigendian: в данных всегда 0, little-endian
    point_step = r.u32()
    r.u32()  # row_step
    n_bytes = r.u32()
    data = memoryview(blob)[r.pos:r.pos + n_bytes]
    r.pos += n_bytes
    is_dense = bool(r.u8())

    dt = point_dtype(fields, point_step)
    points = np.frombuffer(data, dtype=dt, count=height * width)
    return PointCloud2(stamp=sec + nsec * 1e-9, frame_id=frame_id, height=height,
                       width=width, point_step=point_step, is_dense=is_dense,
                       points=points)


def from_ros_message(msg) -> PointCloud2:
    """Перевести разобранное rclpy-сообщение sensor_msgs/msg/PointCloud2.

    Все height·width точек сохраняются, включая NaN: сетчатка раскладывает облако
    в решётку азимут × кольцо по порядку точек, и выброшенная точка сдвинула бы
    всю решётку (`sensor_msgs_py.read_points(skip_nans=True)` делает именно это).
    """
    if msg.is_bigendian:
        raise ValueError("big-endian PointCloud2 не поддерживается")
    fields = [(f.name, f.offset, f.datatype, f.count) for f in msg.fields]
    dt = point_dtype(fields, msg.point_step)
    points = np.frombuffer(msg.data, dtype=dt, count=msg.height * msg.width)
    stamp = msg.header.stamp.sec + msg.header.stamp.nanosec * 1e-9
    return PointCloud2(stamp=stamp, frame_id=msg.header.frame_id, height=msg.height,
                       width=msg.width, point_step=msg.point_step,
                       is_dense=bool(msg.is_dense), points=points)