109 lines
4.4 KiB
Python
109 lines
4.4 KiB
Python
"""Потоковый прогон большого бэга прямо из tar-архива.
|
||
|
||
`new_data` — это 90 ГБ в 221 шарде, и распаковывать его целиком некуда. Скрипт
|
||
достаёт по одному шарду во временный каталог, прогоняет конвейер, забирает
|
||
дескрипторы кандидатов и удаляет шард. Состояние конвейера (поза, коридор,
|
||
собственное движение, треки) переносится между шардами, поэтому запись
|
||
обрабатывается как одна непрерывная поездка.
|
||
|
||
python tools/stream_new_data.py --shards 60 --out data/cache/new_data_candidates.npz
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import re
|
||
import shutil
|
||
import tarfile
|
||
import tempfile
|
||
import time
|
||
from pathlib import Path
|
||
|
||
import numpy as np
|
||
|
||
import _bootstrap as B # noqa: F401
|
||
from flyguard.bag import Bag
|
||
from flyguard.mushroom_body import describe
|
||
from flyguard.pipeline import FlyGuard, Params
|
||
|
||
SHARD_RE = re.compile(r"_(\d+)\.db3$")
|
||
|
||
|
||
def shard_members(tar_path: Path):
|
||
"""Список (номер, TarInfo) шардов, отсортированный по номеру."""
|
||
out = []
|
||
with tarfile.open(tar_path, "r:") as t:
|
||
for m in t:
|
||
mm = SHARD_RE.search(m.name)
|
||
if mm and m.isfile():
|
||
out.append((int(mm.group(1)), m.name, m.size))
|
||
out.sort()
|
||
return out
|
||
|
||
|
||
def main() -> None:
|
||
ap = argparse.ArgumentParser(description=__doc__)
|
||
ap.add_argument("--tar", default=r"C:\Games\Study\AI_Lidar\датасет\new_data")
|
||
ap.add_argument("--out", default=str(B.CACHE / "new_data_candidates.npz"))
|
||
ap.add_argument("--shards", type=int, default=0, help="0 — все")
|
||
ap.add_argument("--stride", type=int, default=1, help="прореживание кадров")
|
||
ap.add_argument("--workdir", default=None)
|
||
args = ap.parse_args()
|
||
|
||
tar_path = Path(args.tar)
|
||
B.CACHE.mkdir(parents=True, exist_ok=True)
|
||
members = shard_members(tar_path)
|
||
if args.shards:
|
||
members = members[:args.shards]
|
||
print(f"шардов к обработке: {len(members)} из архива {tar_path.name}")
|
||
|
||
work = Path(args.workdir) if args.workdir else Path(tempfile.mkdtemp(prefix="flyguard_"))
|
||
work.mkdir(parents=True, exist_ok=True)
|
||
|
||
fg = FlyGuard(Params(), memory=None)
|
||
rows, dists, frames_total = [], [], 0
|
||
t_start = time.time()
|
||
|
||
with tarfile.open(tar_path, "r:") as tar:
|
||
for k, (num, name, size) in enumerate(members):
|
||
shard_dir = work / "bag"
|
||
if shard_dir.exists():
|
||
shutil.rmtree(shard_dir)
|
||
shard_dir.mkdir(parents=True)
|
||
target = shard_dir / Path(name).name
|
||
|
||
src = tar.extractfile(name)
|
||
with open(target, "wb") as dst:
|
||
shutil.copyfileobj(src, dst, length=1 << 22)
|
||
|
||
n_here = 0
|
||
try:
|
||
bag = Bag(shard_dir)
|
||
for _, pc in bag.frames(stride=args.stride):
|
||
res = fg.process(pc)
|
||
n_here += 1
|
||
if res is None:
|
||
continue
|
||
for c in res.candidates:
|
||
rows.append(describe(c))
|
||
dists.append(c.d)
|
||
except Exception as exc: # шард может быть обрезан
|
||
print(f" шард {num}: пропущен ({exc})")
|
||
frames_total += n_here
|
||
target.unlink(missing_ok=True)
|
||
|
||
if k % 5 == 0 or k == len(members) - 1:
|
||
el = time.time() - t_start
|
||
print(f" шард {k+1}/{len(members)} (#{num}) кадров {n_here:4d} | "
|
||
f"всего кадров {frames_total:6d} кандидатов {len(rows):7d} | "
|
||
f"{el/60:5.1f} мин", flush=True)
|
||
|
||
shutil.rmtree(work, ignore_errors=True)
|
||
X = np.stack(rows).astype(np.float32) if rows else np.zeros((0, 21), np.float32)
|
||
np.savez_compressed(args.out, X=X, d=np.asarray(dists, np.float32),
|
||
frames=frames_total)
|
||
print(f"готово: {frames_total} кадров, {X.shape[0]} кандидатов → {args.out}")
|
||
print(f"время: {(time.time()-t_start)/60:.1f} мин")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|