Brainrot_Muxa/tools/make_benchmark.py

"""Размеченный полигон: сценарии сближения с синтетическим препятствием.

Для каждого бэга и каждого типа предмета строится сценарий: предмет ставится
в фиксированную точку тоннеля далеко впереди, поезд к нему подъезжает, и на
каждом кадре известна истинная дистанция. Отсюда получаются именно те цифры,
которые просит ТЗ: с какой дальности предмет уверенно виден, сколько ложных
тревог и как это зависит от размера.

Первый проход считает собственное движение по чистым данным (это и есть
разметка по дистанции), второй — гоняет конвейер по кадрам со вставленным
предметом. Все сценарии одного бэга обрабатываются в одном проходе по файлу:
чтение данных дороже самой обработки.

    python tools/make_benchmark.py --out artifacts/benchmark.npz --memory artifacts/mushroom_body.npz
"""
from __future__ import annotations

import argparse
import json
import time

import numpy as np

import _bootstrap as B  # noqa: F401
from flyguard.bag import Bag, find_bags
from flyguard.mushroom_body import MushroomBody
from flyguard.pipeline import FlyGuard, Params
from flyguard.synth import Placement, catalogue, inject

HOLDOUT = "doubleT_obstacle"        # там уже есть настоящий объект


def ego_track(bag: Bag, params: Params, limit: int | None):
    """Первый проход: пройденный путь на каждом кадре (разметка по дистанции)."""
    fg = FlyGuard(params, memory=None)
    s, stamps = [], []
    total = 0.0
    for _, pc in bag.frames(stop=limit):
        res = fg.process(pc)
        if res is None:
            s.append(None); stamps.append(pc.stamp); continue
        total += res.ego.ds if res.ego else 0.0
        s.append(total); stamps.append(pc.stamp)
    return s, stamps


def run_bag(bag_path, params: Params, memory, limit: int, d_start: float,
            laterals: tuple[float, ...], seed: int, readout=None) -> list[dict]:
    bag = Bag(bag_path)
    s_track, _ = ego_track(bag, params, limit)
    have = [x for x in s_track if x is not None]
    if len(have) < 20:
        return []
    travel = have[-1] - have[0]

    cat = catalogue()
    scen = [(name, lat) for name in cat for lat in laterals]
    pipes = [FlyGuard(params, memory=memory, readout=readout) for _ in scen]
    rng = np.random.default_rng(seed)
    records = [[] for _ in scen]

    # решётка и поза нужны для вставки — берутся из отдельного «чистого» конвейера
    guide = FlyGuard(params, memory=None)

    for k, (_, pc) in enumerate(bag.frames(stop=limit)):
        gres = guide.process(pc)
        if gres is None or s_track[k] is None:
            continue
        s_now = s_track[k] - have[0]
        for i, (name, lat) in enumerate(scen):
            d_true = d_start - s_now
            if d_true < 6.0:
                continue
            # Предмет лежит НА ПУТИ, а путь в кривой уходит вбок: на 150 м при
            # радиусе 1300 м это 8.6 м. Если ставить его в поперечных координатах
            # сенсора, он окажется в стене, а не в габарите.
            u_obj = float(gres.corridor.centre(np.array([d_true], np.float32))[0]) + lat
            pc2, lab = inject(pc, guide.layout_full, gres.plane, cat[name],
                              Placement(d=d_true, u=u_obj), rng=rng)
            res = pipes[i].process(pc2)
            if res is None:
                continue
            tol = max(3.0, 0.12 * d_true)
            hit = any(abs(o.distance - d_true) < tol for o in res.decision.objects)
            fp = sum(1 for o in res.decision.objects if abs(o.distance - d_true) >= tol)
            records[i].append((d_true, int(hit), fp, lab["hit_rays"]))

    out = []
    for (name, lat), rec in zip(scen, records):
        if not rec:
            continue
        a = np.array(rec, np.float32)
        out.append(dict(bag=bag.path.name, obj=name, lateral=lat,
                        d=a[:, 0].tolist(), hit=a[:, 1].tolist(),
                        fp=a[:, 2].tolist(), rays=a[:, 3].tolist(),
                        travel=float(travel)))
    return out


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--root", default=str(B.DATA / "for_hackathon"))
    ap.add_argument("--memory")
    ap.add_argument("--mbon", default="",
                    help="обученное считывание MBON; как и память, оно "
                         "видело эти бэги: полигон меряет дальность, "
                         "а не обобщаемость")
    ap.add_argument("--out", default=str(B.ARTIFACTS / "benchmark.json"))
    ap.add_argument("--limit", type=int, default=250)
    ap.add_argument("--d-start", type=float, default=200.0)
    ap.add_argument("--laterals", default="0.0,0.9")
    ap.add_argument("--seed", type=int, default=12345)
    args = ap.parse_args()

    B.ARTIFACTS.mkdir(parents=True, exist_ok=True)
    memory = MushroomBody.load(args.memory) if args.memory else None
    readout = None
    if args.mbon:
        from flyguard.mbon_readout import MbonReadout
        readout = MbonReadout.load(args.mbon)
        print(f"считывание MBON: {args.mbon}")
    params = Params()
    laterals = tuple(float(x) for x in args.laterals.split(","))

    all_rec = []
    for p in find_bags(args.root):
        if p.name == HOLDOUT:
            continue
        t0 = time.time()
        rec = run_bag(p, params, memory, args.limit, args.d_start, laterals,
                      args.seed, readout=readout)
        all_rec.extend(rec)
        n = sum(len(r["d"]) for r in rec)
        print(f"  {p.name:42s} сценариев {len(rec):3d}, наблюдений {n:6d}, "
              f"{time.time()-t0:6.1f} с", flush=True)

    with open(args.out, "w", encoding="utf-8") as f:
        json.dump(all_rec, f, ensure_ascii=False)
    print(f"сохранено: {args.out} ({len(all_rec)} сценариев)")


if __name__ == "__main__":
    main()