Brainrot_Muxa/flyguard/bag.py

"""Чтение rosbag2 (storage sqlite3) без установленного ROS.

Поддерживает многошардовые бэги (`new_data` — 221 файл `*.db3`), произвольные
имена топиков и порядок шардов по числовому суффиксу. Метаданные `metadata.yaml`
не требуются: список топиков берётся из самой БД, что снимает зависимость от PyYAML.
"""
from __future__ import annotations

import logging
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

logger = logging.getLogger("flyguard.bag")

_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()
        self.n_frames_failed = 0  # обновляется по ходу последнего frames()

    @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, облако) для сообщений выбранного топика.

        Индексация сквозная по всему бэгу; шарды читаются по порядку, внутри
        шарда — по возрастанию времени. Сообщение, которое не удаётся разобрать
        (битый CDR у реального сенсора), пропускается с предупреждением в лог,
        а не останавливает чтение бэга — `start`/`stop` считают по позиции
        сообщения в бэге, а не по числу успешно разобранных.
        """
        idx = 0
        self.n_frames_failed = 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:
                        try:
                            pc = parse_pointcloud2(blob)
                        except Exception as e:  # noqa: BLE001
                            self.n_frames_failed += 1
                            logger.warning(
                                "кадр %d в %s не разобран (%s: %s) — пропущен",
                                idx, shard.name, type(e).__name__, e)
                        else:
                            yield ts, pc
                    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)