124 lines
5.4 KiB
Python
124 lines
5.4 KiB
Python
"""Диагностические картинки одного кадра: дальностный образ, срезы, вид сверху.
|
||
|
||
python tools/viz_frame.py --bag data/for_hackathon/doubleT_obstacle --frame 40
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
|
||
import matplotlib
|
||
matplotlib.use("Agg")
|
||
import matplotlib.pyplot as plt # noqa: E402
|
||
import numpy as np # noqa: E402
|
||
|
||
import _bootstrap as B # noqa: F401,E402
|
||
from flyguard.bag import Bag # noqa: E402
|
||
from flyguard.retina import ScanLayout # noqa: E402
|
||
|
||
|
||
def load(bag_path: str, frame: int, n_calib: int = 12):
|
||
bag = Bag(bag_path)
|
||
clouds = [pc for _, pc in bag.frames(start=2, stop=2 + n_calib)]
|
||
layout = ScanLayout.calibrate(clouds)
|
||
_, pc = next(bag.frames(start=frame, stop=frame + 1))
|
||
return bag, layout, layout.project(pc)
|
||
|
||
|
||
def main() -> None:
|
||
ap = argparse.ArgumentParser(description=__doc__)
|
||
ap.add_argument("--bag", required=True)
|
||
ap.add_argument("--frame", type=int, default=40)
|
||
ap.add_argument("--fov", type=float, default=30.0, help="полусектор по азимуту, °")
|
||
ap.add_argument("--out", default=None)
|
||
args = ap.parse_args()
|
||
|
||
bag, layout, img = load(args.bag, args.frame)
|
||
cols = layout.column_slice(args.fov)
|
||
lay = layout.sub(cols)
|
||
sub = img.crop(cols)
|
||
|
||
xyz = sub.xyz(lay)
|
||
x, y, z = xyz[..., 0], xyz[..., 1], xyz[..., 2]
|
||
d = -y # вперёд
|
||
v = sub.valid
|
||
|
||
# уровень пола по ближней зоне под сенсором
|
||
near = v & (d > 4) & (d < 25) & (np.abs(x) < 2.0)
|
||
z_rail = float(np.percentile(z[near], 2)) if near.sum() > 200 else np.nan
|
||
h = z - z_rail
|
||
|
||
fig = plt.figure(figsize=(19, 12), dpi=110)
|
||
fig.suptitle(f"{bag.path.name} кадр {args.frame} "
|
||
f"пол z={z_rail:.2f} м валидных {v.mean():.0%}", fontsize=13)
|
||
|
||
# 1. дальностный образ
|
||
ax = fig.add_subplot(3, 2, 1)
|
||
disp = np.where(v, sub.r_near, np.nan)
|
||
im = ax.imshow(disp, aspect="auto", cmap="turbo", vmin=0, vmax=120,
|
||
extent=[lay.az_grid_deg[0], lay.az_grid_deg[-1],
|
||
lay.el_deg[-1], lay.el_deg[0]])
|
||
ax.set_title("дальностный образ (выпрямленный), м")
|
||
ax.set_xlabel("азимут, °"); ax.set_ylabel("элевация, °")
|
||
fig.colorbar(im, ax=ax, fraction=0.03)
|
||
|
||
# 2. интенсивность
|
||
ax = fig.add_subplot(3, 2, 2)
|
||
im = ax.imshow(np.where(v, sub.inten, np.nan), aspect="auto", cmap="magma",
|
||
vmin=0, vmax=60,
|
||
extent=[lay.az_grid_deg[0], lay.az_grid_deg[-1],
|
||
lay.el_deg[-1], lay.el_deg[0]])
|
||
ax.set_title("интенсивность")
|
||
ax.set_xlabel("азимут, °"); ax.set_ylabel("элевация, °")
|
||
fig.colorbar(im, ax=ax, fraction=0.03)
|
||
|
||
# 3. вид сверху
|
||
ax = fig.add_subplot(3, 2, 3)
|
||
m = v & (d > 0) & (d < 200)
|
||
ax.scatter(d[m], x[m], s=0.05, c=h[m], cmap="viridis", vmin=-1, vmax=4, linewidths=0)
|
||
for s in (-1.5, 1.5):
|
||
ax.axhline(s, color="r", lw=0.8, ls="--")
|
||
ax.set_xlim(0, 200); ax.set_ylim(-8, 8)
|
||
ax.set_title("вид сверху (цвет — высота над головкой рельса)")
|
||
ax.set_xlabel("расстояние вперёд, м"); ax.set_ylabel("поперёк, м")
|
||
|
||
# 4. вид сбоку
|
||
ax = fig.add_subplot(3, 2, 4)
|
||
ax.scatter(d[m], h[m], s=0.05, c=np.abs(x[m]), cmap="plasma", vmin=0, vmax=4, linewidths=0)
|
||
ax.axhline(0, color="k", lw=0.8)
|
||
ax.axhline(2.1, color="r", lw=0.8, ls="--")
|
||
ax.set_xlim(0, 200); ax.set_ylim(-3, 6)
|
||
ax.set_title("вид сбоку (цвет — |поперёк|)")
|
||
ax.set_xlabel("расстояние вперёд, м"); ax.set_ylabel("высота над рельсом, м")
|
||
|
||
# 5. поперечные сечения
|
||
ax = fig.add_subplot(3, 2, 5)
|
||
for lo, hi, c in [(5, 15, "tab:blue"), (30, 45, "tab:orange"),
|
||
(60, 80, "tab:green"), (100, 130, "tab:red")]:
|
||
s = v & (d >= lo) & (d < hi)
|
||
ax.scatter(x[s], h[s], s=0.3, alpha=0.4, color=c, label=f"{lo}–{hi} м", linewidths=0)
|
||
ax.add_patch(plt.Rectangle((-1.5, 0), 3.0, 2.1, fill=False, ec="r", lw=1.5))
|
||
ax.set_xlim(-6, 6); ax.set_ylim(-2, 6); ax.set_aspect("equal")
|
||
ax.legend(fontsize=8, markerscale=8)
|
||
ax.set_title("поперечные сечения (красное — габарит 3.0 × 2.1 м)")
|
||
ax.set_xlabel("поперёк, м"); ax.set_ylabel("высота, м")
|
||
|
||
# 6. профиль дальности по лучам вдоль пути
|
||
ax = fig.add_subplot(3, 2, 6)
|
||
bins = np.arange(0, 205, 2.0)
|
||
for lo, hi, name in [(-1.0, 1.0, "|az|<1°"), (-5, 5, "|az|<5°"), (-20, 20, "|az|<20°")]:
|
||
sel = (lay.az_grid_deg >= lo) & (lay.az_grid_deg <= hi)
|
||
dd = d[:, sel][v[:, sel]]
|
||
ax.hist(dd, bins=bins, histtype="step", label=name, log=True)
|
||
ax.set_title("распределение дальностей по секторам")
|
||
ax.set_xlabel("расстояние вперёд, м"); ax.set_ylabel("число лучей")
|
||
ax.legend(fontsize=8)
|
||
|
||
out = args.out or (B.FIGURES / f"frame_{bag.path.name}_{args.frame:04d}.png")
|
||
B.FIGURES.mkdir(parents=True, exist_ok=True)
|
||
fig.tight_layout(rect=(0, 0, 1, 0.97))
|
||
fig.savefig(out)
|
||
print("сохранено:", out)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|