b8e95eccbe
Ретина, ламина, медулла, лобула, грибовидное тело, веерное тело, центральный комплекс, нисходящие нейроны. Обучение памяти тоннеля и считывания MBON, оценка leave-one-bag-out, полигон дальности, 24 теста. Реальный объект на 55 м — 98.9 % кадров, ложных 7.5 трека на км, кадр обрабатывается за 33 мс на CPU.
351 lines
18 KiB
Python
351 lines
18 KiB
Python
"""RETINA — омматидиальная решётка.
|
||
|
||
Фасеточный глаз дрозофилы — регулярная решётка омматидиев, каждый смотрит в свою
|
||
фиксированную сторону. Вращающийся лидар устроен так же: пара (кольцо, столбец)
|
||
задаёт направление луча. Поэтому облако точек сразу переводится в *ретинотопический*
|
||
дальностный образ `(кольцо, азимут)`, и вся дальнейшая обработка идёт в этой
|
||
решётке — как в зрительной системе мухи, а не в неупорядоченном облаке.
|
||
|
||
Три особенности конкретного сенсора, измеренные по данным (см. `docs/ALGORITHM.md`):
|
||
|
||
1. **Раскладка различается между бэгами**: 3600 азимутов на 360° против 1200 на
|
||
100°. Решётка калибруется по самим данным, ничего не захардкожено.
|
||
2. **Двойное эхо**: соседние столбцы делят один азимут. Когда эхо одно, оба слота
|
||
содержат одно значение; когда два — ближнее несёт объект, дальнее фон за ним.
|
||
Реально различаются ~3 % лучей, и это именно тонкие предметы и кромки.
|
||
3. **Скос решётки**: у каждого лазерного канала свой постоянный азимутальный сдвиг,
|
||
разброс достигает **15.5°** (≈155 столбцов). В сыром виде «столбец» не является
|
||
направлением: соседние кольца одного столбца смотрят в стороны, разнесённые на
|
||
градусы. Поэтому образ **выпрямляется** целочисленным сдвигом строк; остаточная
|
||
ошибка < половины шага азимута и учитывается в таблице направлений.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import warnings
|
||
from contextlib import contextmanager
|
||
from dataclasses import dataclass
|
||
from pathlib import Path
|
||
|
||
import numpy as np
|
||
|
||
from .cdr import PointCloud2
|
||
|
||
DEG = np.pi / 180.0
|
||
|
||
|
||
@dataclass
|
||
class RangeImage:
|
||
"""Выпрямленный дальностный образ в координатах `(кольцо, азимут)`."""
|
||
|
||
stamp: float
|
||
r_near: np.ndarray # (H, W) float32 — ближнее эхо, 0 = нет эха
|
||
r_far: np.ndarray # (H, W) float32 — дальнее эхо, 0 = нет эха
|
||
inten: np.ndarray # (H, W) float32 — интенсивность ближнего эха
|
||
valid: np.ndarray # (H, W) bool
|
||
|
||
@property
|
||
def shape(self) -> tuple[int, int]:
|
||
return self.r_near.shape
|
||
|
||
def xyz(self, layout: "ScanLayout") -> np.ndarray:
|
||
"""Декартовы координаты ближнего эха: (H, W, 3)."""
|
||
return layout.dirs * self.r_near[..., None]
|
||
|
||
def crop(self, cols: slice) -> "RangeImage":
|
||
return RangeImage(self.stamp, self.r_near[:, cols], self.r_far[:, cols],
|
||
self.inten[:, cols], self.valid[:, cols])
|
||
|
||
|
||
class ScanLayout:
|
||
"""Калиброванная и выпрямленная решётка лучей.
|
||
|
||
Азимут отсчитывается от направления движения (вперёд = −Y в системе сенсора),
|
||
положительный вправо; элевация — вверх от горизонтали сенсора.
|
||
"""
|
||
|
||
def __init__(self, el_deg: np.ndarray, az_step_deg: float, az0_deg: float,
|
||
col_shift: np.ndarray, az_resid_deg: np.ndarray,
|
||
n_az: int, n_echo: int, wrap: bool = False):
|
||
self.el_deg = np.asarray(el_deg, np.float64) # (H,)
|
||
self.az_step_deg = float(az_step_deg) # шаг на азимутальный индекс
|
||
self.az0_deg = float(az0_deg)
|
||
self.col_shift = np.asarray(col_shift, np.int64) # (H,) выпрямление
|
||
self.az_resid_deg = np.asarray(az_resid_deg, np.float64) # (H,) остаток < шага/2
|
||
self.n_az = int(n_az)
|
||
self.n_echo = int(n_echo)
|
||
self.wrap = bool(wrap) # круговой скан: края смыкаются
|
||
self.n_rings = self.el_deg.size
|
||
|
||
self.az_grid_deg = self.az0_deg + self.az_step_deg * np.arange(self.n_az)
|
||
|
||
# карта выборки для выпрямления: out[h, j] = raw[h, j + shift[h]]
|
||
j = np.arange(self.n_az)[None, :]
|
||
src = j + self.col_shift[:, None]
|
||
if self.wrap:
|
||
# у кругового скана «выехавшие» столбцы приходят с другого края
|
||
self.gather = np.mod(src, self.n_az).astype(np.intp)
|
||
self.gather_ok = np.ones_like(self.gather, bool)
|
||
else:
|
||
self.gather_ok = (src >= 0) & (src < self.n_az)
|
||
self.gather = np.clip(src, 0, self.n_az - 1).astype(np.intp)
|
||
|
||
self.az_full_deg = self.az_grid_deg[None, :] + self.az_resid_deg[:, None]
|
||
self.dirs = self._unit_dirs().astype(np.float32) # (H, W, 3)
|
||
# угловой шаг по кольцам: у решётки из одного кольца градиента нет
|
||
self.el_step_deg = (np.abs(np.gradient(self.el_deg)) if self.n_rings > 1
|
||
else np.full(self.n_rings, 0.125))
|
||
|
||
# ------------------------------------------------------------------ геометрия
|
||
|
||
def _unit_dirs(self) -> np.ndarray:
|
||
az = self.az_full_deg * DEG
|
||
el = (self.el_deg[:, None] * DEG) * np.ones_like(az)
|
||
c = np.cos(el)
|
||
return np.stack([c * np.sin(az), -c * np.cos(az), np.sin(el)], axis=-1)
|
||
|
||
def column_slice(self, half_fov_deg: float) -> slice:
|
||
"""Непрерывный диапазон столбцов внутри ±half_fov по азимуту."""
|
||
inside = np.flatnonzero(np.abs(self.az_grid_deg) <= half_fov_deg)
|
||
if inside.size == 0:
|
||
return slice(0, self.n_az)
|
||
return slice(int(inside[0]), int(inside[-1]) + 1)
|
||
|
||
def sub(self, cols: slice) -> "ScanLayout":
|
||
"""Урезанная по азимуту копия решётки (для обработки только переднего сектора)."""
|
||
start = cols.start or 0
|
||
stop = cols.stop if cols.stop is not None else self.n_az
|
||
out = ScanLayout.__new__(ScanLayout)
|
||
out.el_deg = self.el_deg
|
||
out.az_step_deg = self.az_step_deg
|
||
out.az0_deg = float(self.az_grid_deg[start])
|
||
out.col_shift = self.col_shift
|
||
out.az_resid_deg = self.az_resid_deg
|
||
out.n_az = stop - start
|
||
out.n_echo = self.n_echo
|
||
out.wrap = False # вырезанный сектор больше не смыкается
|
||
out.n_rings = self.n_rings
|
||
out.az_grid_deg = self.az_grid_deg[start:stop]
|
||
out.gather_ok = self.gather_ok[:, start:stop]
|
||
out.gather = self.gather[:, start:stop]
|
||
out.az_full_deg = self.az_full_deg[:, start:stop]
|
||
out.dirs = np.ascontiguousarray(self.dirs[:, start:stop])
|
||
out.el_step_deg = self.el_step_deg
|
||
return out
|
||
|
||
# ------------------------------------------------------------------ проекция
|
||
|
||
def project(self, pc: PointCloud2, cols: slice | None = None) -> RangeImage:
|
||
"""Облако точек → выпрямленный дальностный образ.
|
||
|
||
`cols` задаёт нужный диапазон **выходных** столбцов. Тяжёлая
|
||
арифметика (корень по 900 тыс. точек) выполняется только над теми
|
||
сырыми столбцами, которые в этот диапазон попадут с учётом скоса
|
||
каналов, — на круговом скане это экономит почти всё время стадии.
|
||
"""
|
||
n, w, e = self.n_rings, self.n_az, self.n_echo
|
||
start = 0 if cols is None else (cols.start or 0)
|
||
stop = w if cols is None else (cols.stop if cols.stop is not None else w)
|
||
|
||
lo = start + int(self.col_shift.min())
|
||
hi = stop + int(self.col_shift.max())
|
||
if self.wrap:
|
||
raw_cols = np.arange(lo, hi) % w
|
||
else:
|
||
lo = max(lo, 0)
|
||
hi = min(hi, w)
|
||
raw_cols = None
|
||
|
||
pts = pc.points
|
||
|
||
def cube(name: str) -> np.ndarray:
|
||
a = pts[name].reshape(w, e, n)
|
||
a = a[raw_cols] if raw_cols is not None else a[lo:hi]
|
||
return a.transpose(2, 0, 1)
|
||
|
||
x, y, z = cube("x"), cube("y"), cube("z")
|
||
good = (x != 0) | (y != 0) | (z != 0)
|
||
r = np.sqrt(x * x + y * y + z * z, dtype=np.float32)
|
||
r *= good
|
||
|
||
if e == 1:
|
||
r_near = r[..., 0]
|
||
r_far = r[..., 0]
|
||
it = cube("intensity")[..., 0]
|
||
valid = good[..., 0]
|
||
else:
|
||
inten = cube("intensity")
|
||
near_i = np.argmin(np.where(good, r, np.float32(np.inf)), axis=-1)[..., None]
|
||
far_i = np.argmax(r, axis=-1)[..., None]
|
||
r_near = np.take_along_axis(r, near_i, -1)[..., 0]
|
||
r_far = np.take_along_axis(r, far_i, -1)[..., 0]
|
||
it = np.take_along_axis(inten, near_i, -1)[..., 0]
|
||
valid = good.any(axis=-1)
|
||
|
||
# выпрямление скоса каналов, с поправкой на смещение окна
|
||
g = self.gather[:, start:stop]
|
||
ok = self.gather_ok[:, start:stop]
|
||
if raw_cols is not None:
|
||
g = (g - lo) % w
|
||
else:
|
||
g = g - lo
|
||
ok = ok & (g >= 0) & (g < (hi - lo))
|
||
np.clip(g, 0, hi - lo - 1, out=g)
|
||
|
||
r_near = np.take_along_axis(r_near, g, 1)
|
||
r_far = np.take_along_axis(r_far, g, 1)
|
||
it = np.take_along_axis(it, g, 1)
|
||
valid = np.take_along_axis(valid, g, 1) & ok
|
||
|
||
r_near = np.where(valid, r_near, np.float32(0.0))
|
||
r_far = np.where(valid, r_far, np.float32(0.0))
|
||
return RangeImage(pc.stamp, np.ascontiguousarray(r_near),
|
||
np.ascontiguousarray(r_far), np.ascontiguousarray(it),
|
||
np.ascontiguousarray(valid))
|
||
|
||
# ------------------------------------------------------------------ калибровка
|
||
|
||
@staticmethod
|
||
def calibrate(clouds: list[PointCloud2], n_rings: int | None = None) -> "ScanLayout":
|
||
"""Восстановить решётку по нескольким кадрам.
|
||
|
||
Определяются: число колец и эх, элевация каждого кольца, шаг развёртки,
|
||
азимутальный сдвиг каждого канала и целочисленное выпрямление образа.
|
||
"""
|
||
if not clouds:
|
||
raise ValueError("нужен хотя бы один кадр для калибровки")
|
||
pc0 = clouds[0]
|
||
if n_rings is None:
|
||
n_rings = int(pc0.points["ring"].max()) + 1
|
||
n_cols, rem = divmod(pc0.n_points, n_rings)
|
||
if rem:
|
||
raise ValueError(f"{pc0.n_points} точек не делится на {n_rings} колец")
|
||
|
||
# Направление луча задано сенсором и в каждом кадре одно и то же:
|
||
# кадры нужны только чтобы закрыть лучи, не вернувшие эхо. Поэтому
|
||
# берётся первое конечное значение, а не медиана по стопке кадров —
|
||
# та стоила 1.4 с из 2.8 с всей калибровки и ничего не уточняла:
|
||
# ниже и шаг развёртки, и сдвиг канала берутся медианой по тысячам
|
||
# столбцов, где шум одного отсчёта всё равно усредняется.
|
||
az = el = None
|
||
for pc in clouds:
|
||
x = pc.points["x"].reshape(n_cols, n_rings).T.astype(np.float32)
|
||
y = pc.points["y"].reshape(n_cols, n_rings).T.astype(np.float32)
|
||
z = pc.points["z"].reshape(n_cols, n_rings).T.astype(np.float32)
|
||
ok = (x != 0) | (y != 0) | (z != 0)
|
||
with _quiet():
|
||
a = np.where(ok, np.degrees(np.arctan2(x, -y)), np.nan)
|
||
e = np.where(ok, np.degrees(np.arctan2(z, np.hypot(x, y))), np.nan)
|
||
if az is None:
|
||
az, el = a, e
|
||
continue
|
||
gap = np.isnan(az)
|
||
if not gap.any():
|
||
break
|
||
az[gap] = a[gap]
|
||
el[gap] = e[gap]
|
||
az = az.astype(np.float64)
|
||
el = el.astype(np.float64)
|
||
|
||
n_echo = _detect_echoes(az)
|
||
n_az = n_cols // n_echo
|
||
if n_echo > 1:
|
||
with _quiet():
|
||
az = np.nanmean(az.reshape(n_rings, n_az, n_echo), axis=2)
|
||
el = np.nanmean(el.reshape(n_rings, n_az, n_echo), axis=2)
|
||
|
||
# общий шаг развёртки: медиана по кольцам от робастной оценки наклона
|
||
slopes = []
|
||
for h in range(n_rings):
|
||
row = az[h]
|
||
idx = np.flatnonzero(np.isfinite(row))
|
||
if idx.size < 50:
|
||
continue
|
||
d = np.diff(np.unwrap(np.radians(row[idx]))) / np.diff(idx)
|
||
slopes.append(np.median(np.degrees(d)))
|
||
if not slopes:
|
||
raise ValueError("недостаточно валидных лучей для калибровки развёртки")
|
||
step = float(np.median(slopes))
|
||
|
||
# смещение каждого канала относительно общей развёртки
|
||
j = np.arange(n_az, dtype=np.float64)
|
||
base = step * j
|
||
with _quiet():
|
||
c_ring = np.nanmedian(_wrap180(az - base[None, :]), axis=1) # (H,)
|
||
c_ring = _fill_linear(c_ring)
|
||
c0 = float(np.median(c_ring))
|
||
|
||
# круговой скан: развёртка покрывает полные 360°
|
||
wrap = abs(step) * n_az > 350.0
|
||
|
||
# начало отсчёта выбирается так, чтобы «вперёд» (азимут 0) был в середине
|
||
# образа — иначе шов ±180° разрезал бы рабочий сектор пополам
|
||
if wrap:
|
||
k = int(np.rint(-c0 / step)) - n_az // 2
|
||
c0 = _wrap180(c0 + step * k)
|
||
# разница берётся по кратчайшей дуге: иначе шов ±180° даёт сдвиг в пол-оборота
|
||
shift = np.rint(_wrap180(c0 - c_ring) / step).astype(np.int64)
|
||
resid = _wrap180(c_ring + step * shift - c0)
|
||
|
||
with _quiet():
|
||
el_ring = _fill_linear(np.nanmedian(el, axis=1))
|
||
return ScanLayout(el_ring, step, float(c0), shift, resid, n_az, n_echo, wrap)
|
||
|
||
# ------------------------------------------------------------------ сериализация
|
||
|
||
def save(self, path: str | Path) -> None:
|
||
np.savez_compressed(path, el_deg=self.el_deg, az_step_deg=self.az_step_deg,
|
||
az0_deg=self.az0_deg, col_shift=self.col_shift,
|
||
az_resid_deg=self.az_resid_deg, n_az=self.n_az,
|
||
n_echo=self.n_echo, wrap=self.wrap)
|
||
|
||
@staticmethod
|
||
def load(path: str | Path) -> "ScanLayout":
|
||
d = np.load(path)
|
||
return ScanLayout(d["el_deg"], float(d["az_step_deg"]), float(d["az0_deg"]),
|
||
d["col_shift"], d["az_resid_deg"], int(d["n_az"]),
|
||
int(d["n_echo"]), bool(d["wrap"]) if "wrap" in d else False)
|
||
|
||
def __repr__(self) -> str:
|
||
return (f"ScanLayout(колец={self.n_rings}, азимутов={self.n_az}, эх={self.n_echo}, "
|
||
f"сектор {self.az_grid_deg.min():.1f}°…{self.az_grid_deg.max():.1f}°, "
|
||
f"шаг {abs(self.az_step_deg):.3f}°, "
|
||
f"элевация {self.el_deg.min():.1f}°…{self.el_deg.max():.1f}°, "
|
||
f"скос каналов {np.ptp(self.col_shift)} стлб)")
|
||
|
||
|
||
# ---------------------------------------------------------------------- вспомогательное
|
||
|
||
@contextmanager
|
||
def _quiet():
|
||
"""Пустые срезы и деление на ноль при калибровке — норма, а не ошибка."""
|
||
with warnings.catch_warnings(), np.errstate(all="ignore"):
|
||
warnings.simplefilter("ignore", RuntimeWarning)
|
||
yield
|
||
|
||
|
||
def _wrap180(a):
|
||
"""Привести угол(ы) в градусах к полуинтервалу (−180, 180]."""
|
||
return -((-np.asarray(a, np.float64) + 180.0) % 360.0 - 180.0)
|
||
|
||
|
||
def _detect_echoes(az: np.ndarray) -> int:
|
||
"""Двойное эхо: соседние столбцы делят азимут (проверяется по каждому кольцу)."""
|
||
if az.shape[1] < 4:
|
||
return 1
|
||
a, b = az[:, 0::2], az[:, 1::2]
|
||
m = np.isfinite(a) & np.isfinite(b)
|
||
if m.sum() < 100:
|
||
return 1
|
||
return 2 if np.mean(np.abs(a[m] - b[m]) < 1e-3) > 0.8 else 1
|
||
|
||
|
||
def _fill_linear(a: np.ndarray) -> np.ndarray:
|
||
"""Линейно достроить NaN-пропуски по индексу (сетка равномерная)."""
|
||
a = np.asarray(a, np.float64).copy()
|
||
bad = ~np.isfinite(a)
|
||
if bad.all():
|
||
raise ValueError("нет ни одного валидного угла для калибровки")
|
||
if bad.any():
|
||
idx = np.arange(a.size)
|
||
a[bad] = np.interp(idx[bad], idx[~bad], a[~bad])
|
||
return a
|