86 lines
3.4 KiB
Python
86 lines
3.4 KiB
Python
"""Отрисовка схемы мозга мухи по кадрам бэга — проверка и материал для видео.
|
||
|
||
python tools/render_brain.py --bag data/for_hackathon/doubleT_obstacle --frames 8
|
||
python tools/render_brain.py --bag ... --video docs/figures/brain.mp4 --limit 200
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
|
||
import numpy as np
|
||
|
||
import _bootstrap as B # noqa: F401
|
||
from flyguard.bag import Bag
|
||
from flyguard.brain_view import BrainView
|
||
from flyguard.mushroom_body import MushroomBody
|
||
from flyguard.pipeline import FlyGuard, Params
|
||
|
||
|
||
def main() -> None:
|
||
ap = argparse.ArgumentParser(description=__doc__)
|
||
ap.add_argument("--bag", required=True)
|
||
ap.add_argument("--memory")
|
||
ap.add_argument("--limit", type=int, default=120)
|
||
ap.add_argument("--frames", type=int, default=6, help="сколько кадров сохранить в PNG")
|
||
ap.add_argument("--video", default=None)
|
||
ap.add_argument("--style", default="hybrid",
|
||
choices=("scheme", "cloud", "hybrid"),
|
||
help="схема | облако нейронов | панели + облако")
|
||
args = ap.parse_args()
|
||
|
||
if args.style == "scheme":
|
||
view = BrainView()
|
||
elif args.style == "cloud":
|
||
from flyguard.brain_atlas import NeuronCloud
|
||
view = NeuronCloud()
|
||
else:
|
||
from flyguard.brain_hybrid import BrainHybrid
|
||
view = BrainHybrid()
|
||
if not view.enabled:
|
||
raise SystemExit("нужен OpenCV, а для облака — ещё атлас нейронов "
|
||
"(python tools/build_brain_atlas.py)")
|
||
import cv2
|
||
|
||
memory = MushroomBody.load(args.memory) if args.memory else None
|
||
fg = FlyGuard(Params(enable_looming=True), memory=memory)
|
||
bag = Bag(args.bag)
|
||
|
||
B.FIGURES.mkdir(parents=True, exist_ok=True)
|
||
writer = None
|
||
saved = 0
|
||
processed = 0
|
||
# кадры калибровки решётки не доходят до отрисовки, поэтому шаг считается
|
||
# от реально обрабатываемых, иначе при малом --limit не сохранится ничего
|
||
usable = max(args.limit - Params().calib_frames, 1)
|
||
every = max(1, usable // max(args.frames, 1))
|
||
|
||
for k, (_, pc) in enumerate(bag.frames(stop=args.limit)):
|
||
res = fg.process(pc, keep_debug=True)
|
||
if res is None:
|
||
continue
|
||
processed += 1
|
||
img = view.render(res)
|
||
if img is None:
|
||
continue
|
||
if args.video:
|
||
if writer is None:
|
||
fourcc = cv2.VideoWriter_fourcc(*"mp4v")
|
||
writer = cv2.VideoWriter(args.video, fourcc, 10.0,
|
||
(img.shape[1], img.shape[0]))
|
||
writer.write(img)
|
||
if processed % every == 0 and saved < args.frames:
|
||
out = B.FIGURES / f"brain_{args.style}_{bag.path.name}_{k:04d}.png"
|
||
cv2.imwrite(str(out), img)
|
||
saved += 1
|
||
print(f" кадр {k:4d}: {out.name} "
|
||
f"{'ПРЕПЯТСТВИЕ %.0f м' % res.decision.distance if res.decision.detected else 'путь свободен'}"
|
||
f" {res.total_ms:.1f} мс")
|
||
|
||
if writer is not None:
|
||
writer.release()
|
||
print("видео:", args.video)
|
||
print(f"обработано кадров: {processed}, сохранено картинок: {saved}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|