b8e95eccbe
Ретина, ламина, медулла, лобула, грибовидное тело, веерное тело, центральный комплекс, нисходящие нейроны. Обучение памяти тоннеля и считывания MBON, оценка leave-one-bag-out, полигон дальности, 24 теста. Реальный объект на 55 м — 98.9 % кадров, ложных 7.5 трека на км, кадр обрабатывается за 33 мс на CPU.
133 lines
5.4 KiB
Python
133 lines
5.4 KiB
Python
"""Чтение rosbag2 (storage sqlite3) без установленного ROS.
|
||
|
||
Поддерживает многошардовые бэги (`new_data` — 221 файл `*.db3`), произвольные
|
||
имена топиков и порядок шардов по числовому суффиксу. Метаданные `metadata.yaml`
|
||
не требуются: список топиков берётся из самой БД, что снимает зависимость от PyYAML.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import re
|
||
import sqlite3
|
||
from contextlib import closing
|
||
from dataclasses import dataclass
|
||
from pathlib import Path
|
||
from typing import Iterator
|
||
|
||
from .cdr import PointCloud2, parse_pointcloud2
|
||
|
||
_SHARD_RE = re.compile(r"_(\d+)\.db3$")
|
||
_POINTCLOUD_TYPE = "sensor_msgs/msg/PointCloud2"
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class BagTopic:
|
||
name: str
|
||
type: str
|
||
count: int
|
||
|
||
|
||
class Bag:
|
||
"""Последовательное чтение облаков точек из rosbag2."""
|
||
|
||
def __init__(self, path: str | Path, topic: str | None = None):
|
||
path = Path(path)
|
||
if path.is_dir():
|
||
shards = sorted(path.glob("*.db3"), key=self._shard_key)
|
||
elif path.suffix == ".db3":
|
||
shards = [path]
|
||
else:
|
||
raise FileNotFoundError(f"не бэг и не .db3: {path}")
|
||
if not shards:
|
||
raise FileNotFoundError(f"в {path} нет файлов *.db3")
|
||
|
||
self.path = path
|
||
self.shards = shards
|
||
self.topics = self._scan_topics()
|
||
self.topic = topic or self._pick_topic()
|
||
|
||
@staticmethod
|
||
def _shard_key(p: Path) -> tuple[int, str]:
|
||
m = _SHARD_RE.search(p.name)
|
||
return (int(m.group(1)) if m else 0, p.name)
|
||
|
||
def _scan_topics(self) -> dict[str, BagTopic]:
|
||
found: dict[str, BagTopic] = {}
|
||
for shard in self.shards:
|
||
with closing(self._connect(shard)) as con:
|
||
rows = con.execute(
|
||
"SELECT t.name, t.type, count(m.id) FROM topics t "
|
||
"LEFT JOIN messages m ON m.topic_id = t.id GROUP BY t.id"
|
||
).fetchall()
|
||
for name, type_, count in rows:
|
||
prev = found.get(name)
|
||
found[name] = BagTopic(name, type_, (prev.count if prev else 0) + count)
|
||
return found
|
||
|
||
def _pick_topic(self) -> str:
|
||
clouds = [t for t in self.topics.values() if t.type == _POINTCLOUD_TYPE]
|
||
if not clouds:
|
||
raise ValueError(f"в {self.path} нет топиков {_POINTCLOUD_TYPE}: "
|
||
f"{sorted(self.topics)}")
|
||
# при нескольких облачных топиках берём самый наполненный
|
||
return max(clouds, key=lambda t: t.count).name
|
||
|
||
@staticmethod
|
||
def _connect(shard: Path) -> sqlite3.Connection:
|
||
return sqlite3.connect(f"file:{shard.as_posix()}?mode=ro&immutable=1", uri=True)
|
||
|
||
def __len__(self) -> int:
|
||
return self.topics[self.topic].count
|
||
|
||
def frames(self, start: int = 0, stop: int | None = None,
|
||
stride: int = 1) -> Iterator[tuple[int, PointCloud2]]:
|
||
"""Выдать (timestamp_ns, облако) для сообщений выбранного топика.
|
||
|
||
Индексация сквозная по всему бэгу; шарды читаются по порядку, внутри
|
||
шарда — по возрастанию времени.
|
||
"""
|
||
idx = 0
|
||
for shard in self.shards:
|
||
# closing(), а не сам connection: у sqlite3 `with` управляет
|
||
# транзакцией и файл остаётся открытым — на Windows его потом
|
||
# невозможно удалить
|
||
with closing(self._connect(shard)) as con:
|
||
row = con.execute("SELECT id FROM topics WHERE name = ?",
|
||
(self.topic,)).fetchone()
|
||
if row is None:
|
||
continue
|
||
topic_id = row[0]
|
||
n_here = con.execute(
|
||
"SELECT count(*) FROM messages WHERE topic_id = ?", (topic_id,)
|
||
).fetchone()[0]
|
||
if stop is not None and idx >= stop:
|
||
return
|
||
if idx + n_here <= start:
|
||
idx += n_here
|
||
continue
|
||
cur = con.execute(
|
||
"SELECT timestamp, data FROM messages WHERE topic_id = ? "
|
||
"ORDER BY timestamp", (topic_id,))
|
||
for ts, blob in cur:
|
||
if stop is not None and idx >= stop:
|
||
return
|
||
if idx >= start and (idx - start) % stride == 0:
|
||
yield ts, parse_pointcloud2(blob)
|
||
idx += 1
|
||
|
||
def describe(self) -> str:
|
||
lines = [f"бэг: {self.path}",
|
||
f"шардов: {len(self.shards)}",
|
||
f"топик: {self.topic}"]
|
||
for t in sorted(self.topics.values(), key=lambda t: -t.count):
|
||
mark = "*" if t.name == self.topic else " "
|
||
lines.append(f" {mark} {t.name} [{t.type}] {t.count} сообщений")
|
||
return "\n".join(lines)
|
||
|
||
|
||
def find_bags(root: str | Path) -> list[Path]:
|
||
"""Найти все каталоги-бэги под указанным корнем."""
|
||
root = Path(root)
|
||
if not root.exists():
|
||
return []
|
||
out = {p.parent for p in root.rglob("*.db3")}
|
||
return sorted(out)
|