124 lines
4.8 KiB
Python
124 lines
4.8 KiB
Python
"""Сравнение способов оценить осевую линию пути: свод тоннеля против полотна.
|
||
|
||
Свод виден далеко, но в кривой уходит из поля зрения; полотно видно только
|
||
вблизи, зато прямо задаёт положение пути. Скрипт печатает и рисует оба профиля,
|
||
чтобы выбрать источник и границы применимости.
|
||
|
||
python tools/compare_centreline.py --bag data/for_hackathon/roundT_doubleT
|
||
"""
|
||
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.geometry import TrackFrame, fit_rail_plane # noqa: E402
|
||
from flyguard.retina import ScanLayout # noqa: E402
|
||
|
||
|
||
def slices(d_lo=8.0, d_hi=200.0, n=30):
|
||
e = np.geomspace(d_lo, d_hi, n + 1)
|
||
return list(zip(e[:-1], e[1:]))
|
||
|
||
|
||
def bore_centre(tf):
|
||
out = []
|
||
for lo, hi in slices():
|
||
m = tf.valid & (tf.d >= lo) & (tf.d < hi) & (tf.h > 0.6) & (tf.h < 3.2)
|
||
n = int(m.sum())
|
||
if n < 60:
|
||
continue
|
||
u = tf.u[m]
|
||
a, b = np.percentile(u, (3, 97))
|
||
if b - a < 1.5:
|
||
continue
|
||
out.append((0.5 * (lo + hi), 0.5 * (a + b), n))
|
||
return np.array(out) if out else np.zeros((0, 3))
|
||
|
||
|
||
def bed_centre(tf, track=0.0):
|
||
"""Центр полотна: полоса вокруг уровня рельсов, узкая по ширине."""
|
||
out = []
|
||
ref = track
|
||
for lo, hi in slices(8.0, 140.0, 26):
|
||
m = (tf.valid & (tf.d >= lo) & (tf.d < hi)
|
||
& (tf.h > -0.20) & (tf.h < 0.45) & (np.abs(tf.u - ref) < 2.6))
|
||
n = int(m.sum())
|
||
if n < 40:
|
||
continue
|
||
u = tf.u[m]
|
||
a, b = np.percentile(u, (5, 95))
|
||
if not (0.8 < b - a < 5.0):
|
||
continue
|
||
c = 0.5 * (a + b)
|
||
out.append((0.5 * (lo + hi), c, n))
|
||
ref = c # следующий срез ищем рядом с найденным
|
||
return np.array(out) if out else np.zeros((0, 3))
|
||
|
||
|
||
def fit_arc(d, u, w):
|
||
"""u ≈ c1·d + c2·d² (без свободного члена: под поездом ось проходит через 0)."""
|
||
A = np.stack([d, d * d], axis=1) * w[:, None]
|
||
c, *_ = np.linalg.lstsq(A, u * w, rcond=None)
|
||
return c
|
||
|
||
|
||
def main() -> None:
|
||
ap = argparse.ArgumentParser(description=__doc__)
|
||
ap.add_argument("--bag", required=True)
|
||
ap.add_argument("--frames", type=int, default=8)
|
||
ap.add_argument("--stride", type=int, default=20)
|
||
args = ap.parse_args()
|
||
|
||
bag = Bag(args.bag)
|
||
layout = ScanLayout.calibrate([pc for _, pc in bag.frames(start=2, stop=14)])
|
||
cols = layout.column_slice(40.0)
|
||
lay = layout.sub(cols)
|
||
|
||
fig, axes = plt.subplots(1, 2, figsize=(13, 5), dpi=110)
|
||
plane = None
|
||
radii_b, radii_d = [], []
|
||
for _, pc in bag.frames(stop=args.frames * args.stride, stride=args.stride):
|
||
img = layout.project(pc).crop(cols)
|
||
plane = fit_rail_plane(img, lay, prev=plane)
|
||
tf = TrackFrame(img, lay, plane)
|
||
|
||
bo = bore_centre(tf)
|
||
be = bed_centre(tf)
|
||
dd = np.linspace(0, 200, 100)
|
||
if bo.shape[0] >= 4:
|
||
u0 = bo[0, 1]
|
||
c = fit_arc(bo[:, 0], bo[:, 1] - u0, np.sqrt(bo[:, 2]))
|
||
axes[0].plot(bo[:, 0], bo[:, 1] - u0, ".", ms=4)
|
||
axes[0].plot(dd, c[0] * dd + c[1] * dd ** 2, lw=1, alpha=0.6)
|
||
radii_b.append(1 / (2 * c[1]) if c[1] else np.inf)
|
||
if be.shape[0] >= 4:
|
||
c = fit_arc(be[:, 0], be[:, 1], np.sqrt(be[:, 2]))
|
||
axes[1].plot(be[:, 0], be[:, 1], ".", ms=4)
|
||
axes[1].plot(dd, c[0] * dd + c[1] * dd ** 2, lw=1, alpha=0.6)
|
||
radii_d.append(1 / (2 * c[1]) if c[1] else np.inf)
|
||
|
||
for ax, t, r in ((axes[0], "центр свода тоннеля", radii_b),
|
||
(axes[1], "центр полотна пути", radii_d)):
|
||
ax.axhline(1.7, color="r", ls="--", lw=0.8); ax.axhline(-1.7, color="r", ls="--", lw=0.8)
|
||
med = np.median(r) if r else np.nan
|
||
spread = np.std(r) if len(r) > 1 else np.nan
|
||
ax.set_title(f"{t}\nR медиана {med:.0f} м, разброс {spread:.0f} м, оценок {len(r)}")
|
||
ax.set_xlabel("вперёд, м"); ax.set_ylabel("смещение оси, м")
|
||
ax.set_ylim(-14, 14); ax.set_xlim(0, 200)
|
||
print(f"{t}: R = " + ", ".join(f"{x:.0f}" for x in r))
|
||
|
||
B.FIGURES.mkdir(parents=True, exist_ok=True)
|
||
out = B.FIGURES / f"centreline_{bag.path.name}.png"
|
||
fig.suptitle(bag.path.name)
|
||
fig.tight_layout(); fig.savefig(out)
|
||
print("сохранено:", out)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|