forked from Dan4ick/Lidar_Muxa
b8e95eccbe
Ретина, ламина, медулла, лобула, грибовидное тело, веерное тело, центральный комплекс, нисходящие нейроны. Обучение памяти тоннеля и считывания MBON, оценка leave-one-bag-out, полигон дальности, 24 теста. Реальный объект на 55 м — 98.9 % кадров, ложных 7.5 трека на км, кадр обрабатывается за 33 мс на CPU.
130 lines
4.4 KiB
Python
130 lines
4.4 KiB
Python
"""Разбор sensor_msgs/msg/PointCloud2 из CDR без зависимости от ROS.
|
||
|
||
Нужен для двух сценариев:
|
||
* офлайн-эксперименты на машине без ROS (Windows);
|
||
* прямое чтение rosbag внутри контейнера, минуя `ros2 bag play`.
|
||
|
||
Внутри ROS-ноды сообщение приходит уже разобранным, и этот модуль не используется.
|
||
"""
|
||
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)
|