502 lines
25 KiB
Python
502 lines
25 KiB
Python
"""Видеокарта для плотных стадий конвейера: сетчатка, ламина, кластеризация.
|
||
|
||
Три стадии работают с целым образом 128 × 600 лучей и занимают больше
|
||
половины кадра: раскладка точек по решётке, ламина (центр минус окружение на
|
||
трёх масштабах) и кластеризация лучей с учётом глубины. Они и перенесены.
|
||
Остальное — плоскость пути, ось, оценка движения, треки, решение — остаётся
|
||
на процессоре: там мелкие массивы и ветвистая логика, и копирование на
|
||
видеокарту стоило бы дороже самого расчёта.
|
||
|
||
Каждая функция повторяет процессорную версию операция в операцию, в тех же
|
||
типах. Раскладка по порядку точек, ламина и кластеризация дают результат,
|
||
совпадающий с процессорным бит в бит (тесты и сверка на записях,
|
||
EXPERIMENTS п. 21):
|
||
|
||
* сетчатка — только выборки, умножения и корень, всё в float32 без слияния
|
||
операций;
|
||
* ламина — равномерные окна считаются через накопленные суммы в float64 и
|
||
округляются в float32 после каждой оси, ровно как `scipy.ndimage.
|
||
uniform_filter`: сумма чисел float32 такого диапазона в float64 точна, и
|
||
результат не зависит от порядка сложения;
|
||
* кластеризация — целочисленный граф; компоненты нумеруются по наименьшему
|
||
лучу, как у `scipy.sparse.csgraph.connected_components`.
|
||
|
||
Исключение — раскладка по углам точек (кадры с нарушенным порядком, как в
|
||
синтетике организаторов): арктангенс видеокарты и процессора может разойтись
|
||
в последнем знаке, а numpy сортирует ячейки неустойчиво. Расходятся считанные
|
||
лучи на кадр.
|
||
|
||
Любой сбой видеокарты — нет драйвера, не хватило памяти, ошибка ядра — ловит
|
||
конвейер: до конца работы он считает на процессоре (`device.
|
||
notify_cuda_error`), а кадр досчитывается там же, решение не теряется.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import warnings
|
||
|
||
import numpy as np
|
||
|
||
from .lamina import SCALES, LaminaOutput
|
||
from .retina import DEG, ORDER_TOL_DEG, RangeImage, ScanLayout
|
||
|
||
|
||
class GpuStages:
|
||
"""Плотные стадии на видеокарте; держит таблицы решётки и образ кадра."""
|
||
|
||
def __init__(self, device: str = "cuda"):
|
||
import torch
|
||
self.torch = torch
|
||
self.dev = torch.device(device)
|
||
torch.zeros(1, device=self.dev) # поднять контекст сразу, а не на первом кадре
|
||
self._tab_key = None
|
||
self._tab: dict = {}
|
||
self._win: dict = {}
|
||
self._off: dict = {}
|
||
self._box: dict = {}
|
||
self._graphs = True # графы CUDA для ламины; при отказе — по одной операции
|
||
self._lam_key = None
|
||
# образ текущего кадра: ламина и кластеризация берут его отсюда, не копируя заново
|
||
self.r = None
|
||
self.valid = None
|
||
self._r_host = None
|
||
|
||
def warmup(self, shape: tuple[int, int] = (128, 600)) -> None:
|
||
"""Прогнать все три стадии на выдуманном кадре.
|
||
|
||
Первый запуск каждого ядра CUDA подгружает его код: без прогрева первый
|
||
кадр записи шёл 1–1.7 с, очередь подписки переполнялась, и терялись
|
||
первые 6–8 кадров (замерено в контейнере). Раскладка прогревается на
|
||
маленькой решётке обоими путями — по порядку точек и по углам.
|
||
"""
|
||
from .cdr import PointCloud2
|
||
t = self.torch
|
||
rng = np.random.default_rng(0)
|
||
n, w, e = 8, 64, 2
|
||
lay = ScanLayout(np.linspace(10.0, -20.0, n), -0.1, 3.2, np.zeros(n, np.int64),
|
||
np.zeros(n), w, e)
|
||
xyz = lay.dirs.transpose(1, 0, 2)[:, None, :, :] * np.float32(20.0) # столбец·эхо·кольцо
|
||
xyz = np.broadcast_to(xyz, (w, e, n, 3)).reshape(-1, 3)
|
||
pts = np.zeros(xyz.shape[0], dtype=[("x", "<f4"), ("y", "<f4"), ("z", "<f4"),
|
||
("intensity", "<f4")])
|
||
pts["x"], pts["y"], pts["z"] = xyz[:, 0], xyz[:, 1], xyz[:, 2]
|
||
pts["intensity"] = 10.0
|
||
cloud = PointCloud2(0.0, "", 1, pts.size, pts.dtype.itemsize, True, pts)
|
||
self.project(lay, cloud, slice(8, w - 8))
|
||
mixed = PointCloud2(0.0, "", 1, pts.size, pts.dtype.itemsize, True,
|
||
pts[rng.permutation(pts.size)])
|
||
self.project(lay, mixed, slice(8, w - 8))
|
||
|
||
r = (4.0 + 150.0 * rng.random(shape)).astype(np.float32)
|
||
self.r = t.as_tensor(r, device=self.dev)
|
||
self.valid = t.ones(shape, dtype=t.bool, device=self.dev)
|
||
self._r_host = r
|
||
self.lamina()
|
||
self.cluster_by_depth(rng.random(shape) < 0.3, r)
|
||
self.r = self.valid = self._r_host = None
|
||
t.cuda.synchronize(self.dev)
|
||
|
||
# ------------------------------------------------------------------ таблицы
|
||
|
||
def _tables(self, L: ScanLayout) -> dict:
|
||
key = (id(L), L.n_points, L.n_az)
|
||
if key != self._tab_key:
|
||
t, d = self.torch, self.dev
|
||
|
||
def up(a, dtype):
|
||
return t.as_tensor(np.ascontiguousarray(a), dtype=dtype, device=d)
|
||
|
||
self._tab = {
|
||
"sin_el": up(L._sin_el, t.float32),
|
||
"el_asc": up(L._el_asc, t.float64),
|
||
"el_order": up(L._el_order, t.int64),
|
||
"el_deg": up(L.el_deg, t.float64),
|
||
"el_step": up(L.el_step_deg, t.float64),
|
||
"resid": up(L.az_resid_deg, t.float64),
|
||
}
|
||
self._win = {}
|
||
self._tab_key = key
|
||
return self._tab
|
||
|
||
def _window(self, L: ScanLayout, start: int, stop: int) -> tuple:
|
||
"""Сырые столбцы, нужные сектору, и карта выпрямления внутри них."""
|
||
key = (start, stop)
|
||
if key not in self._win:
|
||
w = L.n_az
|
||
lo = start + int(L.col_shift.min())
|
||
hi = stop + int(L.col_shift.max())
|
||
g = L.gather[:, start:stop]
|
||
ok = L.gather_ok[:, start:stop]
|
||
if L.wrap:
|
||
raw = np.arange(lo, hi) % w
|
||
g = (g - lo) % w
|
||
else:
|
||
lo, hi = max(lo, 0), min(hi, w)
|
||
raw = np.arange(lo, hi)
|
||
g = g - lo
|
||
ok = ok & (g >= 0) & (g < (hi - lo))
|
||
g = np.clip(g, 0, hi - lo - 1)
|
||
# столбцы окна идут подряд, кроме стыка кругового скана
|
||
cut = np.flatnonzero(np.diff(raw) != 1) + 1
|
||
runs = [(int(a[0]), int(a[-1]) + 1) for a in np.split(raw, cut) if a.size]
|
||
t = self.torch
|
||
self._win[key] = (runs, t.as_tensor(g, dtype=t.int64, device=self.dev),
|
||
t.as_tensor(ok, device=self.dev))
|
||
return self._win[key]
|
||
|
||
# ------------------------------------------------------------------ загрузка точек
|
||
|
||
def _upload(self, pts: np.ndarray, rows: list[tuple[int, int]]):
|
||
"""Строки структурированного массива → байты на видеокарте, (N, шаг точки)."""
|
||
t = self.torch
|
||
step = pts.dtype.itemsize
|
||
parts = []
|
||
for a, b in rows:
|
||
raw = np.ascontiguousarray(pts[a:b]).view(np.uint8)
|
||
with warnings.catch_warnings():
|
||
# буфер сообщения только для чтения, а тензор мы не пишем
|
||
warnings.simplefilter("ignore", UserWarning)
|
||
parts.append(t.from_numpy(raw).to(self.dev, non_blocking=False))
|
||
buf = parts[0] if len(parts) == 1 else t.cat(parts)
|
||
return buf.view(-1, step)
|
||
|
||
def _field(self, buf, pts: np.ndarray, name: str):
|
||
dt, off = pts.dtype.fields[name][:2]
|
||
if dt.kind != "f" or dt.itemsize != 4 or dt.byteorder == ">":
|
||
raise TypeError(f"поле {name}: ожидался float32, а не {dt}")
|
||
return buf[:, off:off + 4].contiguous().view(self.torch.float32).reshape(-1)
|
||
|
||
# ------------------------------------------------------------------ сетчатка
|
||
|
||
def project(self, L: ScanLayout, pc, cols: slice | None) -> RangeImage:
|
||
"""То же, что `ScanLayout.project`, на видеокарте."""
|
||
w = L.n_az
|
||
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)
|
||
self._tables(L)
|
||
img = None
|
||
if L.indexed and pc.n_points == L.n_points:
|
||
img = self._project_indexed(L, pc, start, stop)
|
||
if img is None:
|
||
L.n_geometric += 1
|
||
img = self._project_geometric(L, pc, start, stop)
|
||
self._r_host = img.r_near
|
||
return img
|
||
|
||
def _project_indexed(self, L: ScanLayout, pc, start: int, stop: int):
|
||
t = self.torch
|
||
tab = self._tab
|
||
n, e = L.n_rings, L.n_echo
|
||
runs, g, ok = self._window(L, start, stop)
|
||
pts = pc.points
|
||
per_col = e * n
|
||
buf = self._upload(pts, [(a * per_col, b * per_col) for a, b in runs])
|
||
ncol = sum(b - a for a, b in runs)
|
||
|
||
def cube(name: str):
|
||
return self._field(buf, pts, name).reshape(ncol, e, n).permute(2, 0, 1)
|
||
|
||
x, y, z = cube("x"), cube("y"), cube("z")
|
||
good = (x != 0) | (y != 0) | (z != 0)
|
||
r = t.sqrt(x * x + y * y + z * z)
|
||
good &= t.isfinite(r)
|
||
r = t.where(good, r, t.zeros((), dtype=r.dtype, device=self.dev))
|
||
|
||
dev = t.abs(z - r * tab["sin_el"][:, None, None])
|
||
tol = r * np.float32(ORDER_TOL_DEG * DEG) + np.float32(1e-3)
|
||
if bool(t.any(dev > tol)):
|
||
return None
|
||
|
||
# яркость — не обязательное поле, без неё нули (как в retina.py)
|
||
has_i = "intensity" in pts.dtype.names
|
||
if e == 1:
|
||
r_near = r[..., 0]
|
||
r_far = r[..., 0]
|
||
it = cube("intensity")[..., 0] if has_i else t.zeros_like(r_near)
|
||
valid = good[..., 0]
|
||
else:
|
||
inten = cube("intensity") if has_i else t.zeros_like(r)
|
||
inf = t.full((), float("inf"), dtype=r.dtype, device=self.dev)
|
||
near_i = t.argmin(t.where(good, r, inf), dim=-1, keepdim=True)
|
||
far_i = t.argmax(r, dim=-1, keepdim=True)
|
||
r_near = t.gather(r, -1, near_i)[..., 0]
|
||
r_far = t.gather(r, -1, far_i)[..., 0]
|
||
it = t.gather(inten, -1, near_i)[..., 0]
|
||
valid = good.any(dim=-1)
|
||
|
||
r_near = t.gather(r_near, 1, g)
|
||
r_far = t.gather(r_far, 1, g)
|
||
it = t.gather(it, 1, g)
|
||
valid = t.gather(valid, 1, g) & ok
|
||
zero = t.zeros((), dtype=r_near.dtype, device=self.dev)
|
||
r_near = t.where(valid, r_near, zero)
|
||
r_far = t.where(valid, r_far, zero)
|
||
return self._finish(pc.stamp, r_near, r_far, it, valid)
|
||
|
||
def _project_geometric(self, L: ScanLayout, pc, start: int, stop: int):
|
||
t = self.torch
|
||
tab = self._tab
|
||
n, wid = L.n_rings, stop - start
|
||
pts = pc.points
|
||
buf = self._upload(pts, [(0, pts.shape[0])])
|
||
x, y, z = (self._field(buf, pts, k) for k in ("x", "y", "z"))
|
||
good = (((x != 0) | (y != 0) | (z != 0)) & t.isfinite(x) & t.isfinite(y)
|
||
& t.isfinite(z))
|
||
idx = t.nonzero(good, as_tuple=True)[0]
|
||
x, y, z = x[idx], y[idx], z[idx]
|
||
|
||
step = L.az_step_deg
|
||
az = t.rad2deg(t.atan2(x, -y))
|
||
jf = (az - np.float32(L.az0_deg)) / np.float32(step)
|
||
if L.wrap:
|
||
jf = t.remainder(jf, L.n_az)
|
||
margin = float(np.abs(L.az_resid_deg).max()) / abs(step) + 1.0
|
||
sel = t.nonzero((jf > start - margin) & (jf < stop - 1 + margin), as_tuple=True)[0]
|
||
idx, x, y, z, az = idx[sel], x[sel], y[sel], z[sel], az[sel]
|
||
|
||
r = t.sqrt(x * x + y * y + z * z)
|
||
el = t.rad2deg(t.asin(t.clamp(z / t.clamp_min(r, np.float32(1e-6)), -1.0, 1.0)))
|
||
el64 = el.double()
|
||
asc = tab["el_asc"]
|
||
k = t.clamp(t.searchsorted(asc, el64), 1, n - 1)
|
||
k = k - ((el64 - asc[k - 1]) < (asc[k] - el64)).long()
|
||
h = tab["el_order"][k]
|
||
ok = t.abs(el64 - tab["el_deg"][h]) <= t.clamp_min(tab["el_step"][h], 0.2)
|
||
j = t.round(((az.double() - tab["resid"][h]) - L.az0_deg) / step).long()
|
||
if L.wrap:
|
||
j = t.remainder(j, L.n_az)
|
||
ok &= (j >= start) & (j < stop)
|
||
sel = t.nonzero(ok, as_tuple=True)[0]
|
||
|
||
r_near = t.zeros(n * wid, dtype=t.float32, device=self.dev)
|
||
r_far = t.zeros_like(r_near)
|
||
it = t.zeros_like(r_near)
|
||
valid = t.zeros(n * wid, dtype=t.bool, device=self.dev)
|
||
if sel.numel():
|
||
cell = h[sel] * wid + (j[sel] - start)
|
||
rr = r[sel]
|
||
mm = t.clamp_max(rr * 1000.0, float((1 << 20) - 1)).long()
|
||
order = t.argsort(cell * (1 << 20) + mm, stable=True)
|
||
cs, rs = cell[order], rr[order]
|
||
new = t.ones(cs.numel(), dtype=t.bool, device=self.dev)
|
||
new[1:] = cs[1:] != cs[:-1]
|
||
first = t.nonzero(new, as_tuple=True)[0]
|
||
last = t.cat([first[1:] - 1, first.new_tensor([cs.numel() - 1])])
|
||
r_near[cs[first]] = rs[first]
|
||
r_far[cs[last]] = rs[last]
|
||
valid[cs[first]] = True
|
||
if "intensity" in pts.dtype.names:
|
||
src = self._field(buf, pts, "intensity")[idx[sel]][order]
|
||
it[cs[first]] = src[first]
|
||
return self._finish(pc.stamp, r_near.reshape(n, wid), r_far.reshape(n, wid),
|
||
it.reshape(n, wid), valid.reshape(n, wid))
|
||
|
||
def _finish(self, stamp: float, r_near, r_far, it, valid) -> RangeImage:
|
||
self.r, self.valid = r_near, valid
|
||
host = [a.cpu().numpy() for a in (r_near, r_far, it, valid)]
|
||
return RangeImage(stamp, *host)
|
||
|
||
# ------------------------------------------------------------------ ламина
|
||
|
||
def _box_axis(self, X, sizes: list[int], dim: int):
|
||
"""Равномерное окно по одной оси для пачки размеров сразу.
|
||
|
||
`X` — (K, B, H, W) float32, k-й срез фильтруется окном `sizes[k]`
|
||
вдоль `dim` (2 или 3) с продолжением краевым значением. Суммы окон
|
||
берутся из накопленной суммы в float64: все частичные суммы чисел
|
||
float32 такого диапазона в float64 точны, поэтому сумма окна точна при
|
||
любом порядке сложения. Деление на размер и округление в float32 —
|
||
как в `scipy.ndimage.uniform_filter1d`.
|
||
"""
|
||
t = self.torch
|
||
Xd = X.double()
|
||
n = Xd.shape[dim]
|
||
lo, hi, n_lo, n_hi, size = self._box_tables(tuple(sizes), n, dim)
|
||
cs = t.cumsum(Xd, dim=dim)
|
||
cs = t.cat([t.zeros_like(cs.narrow(dim, 0, 1)), cs], dim=dim)
|
||
|
||
def along(a):
|
||
return a.expand(*Xd.shape[:dim], n, *Xd.shape[dim + 1:])
|
||
|
||
s = (t.gather(cs, dim, along(hi)) - t.gather(cs, dim, along(lo))
|
||
+ along(n_lo) * Xd.narrow(dim, 0, 1) + along(n_hi) * Xd.narrow(dim, n - 1, 1))
|
||
return (s / size).float()
|
||
|
||
def _box_tables(self, sizes: tuple, n: int, dim: int):
|
||
"""Границы окон по оси: считаются один раз (и до записи графа CUDA)."""
|
||
key = (sizes, n, dim)
|
||
if key not in self._box:
|
||
t = self.torch
|
||
p = t.as_tensor([s // 2 for s in sizes], dtype=t.int64, device=self.dev)
|
||
i = t.arange(n, device=self.dev)
|
||
shape = [len(sizes), 1, 1, 1]
|
||
shape[dim] = n
|
||
lo = (i[None, :] - p[:, None]).clamp(min=0).reshape(shape)
|
||
hi = ((i[None, :] + p[:, None]).clamp(max=n - 1) + 1).reshape(shape)
|
||
n_lo = (p[:, None] - i[None, :]).clamp(min=0).double().reshape(shape)
|
||
n_hi = (i[None, :] + p[:, None] - (n - 1)).clamp(min=0).double().reshape(shape)
|
||
size = t.as_tensor(sizes, dtype=t.float64, device=self.dev).reshape(-1, 1, 1, 1)
|
||
self._box[key] = (lo, hi, n_lo, n_hi, size)
|
||
return self._box[key]
|
||
|
||
def _boxes(self, X, sizes: list[tuple[int, int]]):
|
||
"""`scipy.ndimage.uniform_filter(x, size, mode="nearest")` для пачки.
|
||
|
||
`X` — (B, H, W) float32; выход (len(sizes), B, H, W): по осям по
|
||
очереди, с округлением в float32 после каждой, как в scipy.
|
||
"""
|
||
K = len(sizes)
|
||
Y = self._box_axis(X[None].expand(K, *X.shape), [s[0] for s in sizes], 2)
|
||
return self._box_axis(Y, [s[1] for s in sizes], 3)
|
||
|
||
def lamina(self, r_max: float = 300.0) -> LaminaOutput:
|
||
"""То же, что `lamina._process_cpu`, по образу текущего кадра.
|
||
|
||
Сама видеокарта считает ламину за доли миллисекунды, а запуск сотни
|
||
мелких операций из Python стоит в несколько раз дороже. Поэтому ламина
|
||
записывается в граф CUDA один раз на размер образа и дальше
|
||
запускается одной командой. Не поддерживает граф драйвер — считаем
|
||
теми же операциями по одной.
|
||
"""
|
||
t = self.torch
|
||
key = (tuple(self.r.shape), float(r_max))
|
||
out = None
|
||
if self._graphs:
|
||
try:
|
||
if self._lam_key != key:
|
||
self._capture_lamina(key, r_max)
|
||
self._lam_in[0].copy_(self.r)
|
||
self._lam_in[1].copy_(self.valid)
|
||
self._lam_graph.replay()
|
||
out = self._lam_out
|
||
except Exception:
|
||
self._graphs = False
|
||
self._lam_key = None
|
||
if out is None:
|
||
out = self._lamina_core(self.r, self.valid, r_max)
|
||
f = out[0].cpu().numpy()
|
||
return LaminaOutput(disp=f[0], on=f[1], off=f[2], on_scale=out[1].cpu().numpy(),
|
||
surround=f[3], hole=f[4])
|
||
|
||
def _capture_lamina(self, key: tuple, r_max: float) -> None:
|
||
t = self.torch
|
||
self._lam_in = (t.empty_like(self.r), t.empty_like(self.valid))
|
||
self._lam_in[0].copy_(self.r)
|
||
self._lam_in[1].copy_(self.valid)
|
||
side = t.cuda.Stream(self.dev)
|
||
side.wait_stream(t.cuda.current_stream(self.dev))
|
||
with t.cuda.stream(side):
|
||
for _ in range(2): # прогрев: таблицы окон и кэш памяти
|
||
self._lamina_core(*self._lam_in, r_max)
|
||
t.cuda.current_stream(self.dev).wait_stream(side)
|
||
g = t.cuda.CUDAGraph()
|
||
with t.cuda.graph(g):
|
||
self._lam_out = self._lamina_core(*self._lam_in, r_max)
|
||
self._lam_graph, self._lam_key = g, key
|
||
|
||
def _lamina_core(self, r, valid, r_max: float):
|
||
"""Ламина целиком на видеокарте: (disp, on, off, surround, hole) и on_scale."""
|
||
t = self.torch
|
||
v = valid.float()
|
||
zero = t.zeros((), dtype=t.float32, device=self.dev)
|
||
disp = t.where(valid & (r > 0.05), 1.0 / r, zero) * v
|
||
|
||
# все окна кадра одной пачкой: на каждом масштабе центр и окружение,
|
||
# для диспаритета и для маски эха, плюс окно «дыры»
|
||
sizes = []
|
||
for r_in, r_out in SCALES:
|
||
sizes += [(2 * r_in + 1, 2 * r_in + 1), (2 * r_out + 1, 4 * r_out + 1)]
|
||
sizes.append((5, 15))
|
||
box = self._boxes(t.stack([disp, v]), sizes)
|
||
|
||
on = t.zeros_like(disp)
|
||
off = t.zeros_like(disp)
|
||
on_scale = t.zeros(disp.shape, dtype=t.int8, device=self.dev)
|
||
surround_mid = None
|
||
for k, (r_in, r_out) in enumerate(SCALES):
|
||
s_in, s_out = sizes[2 * k], sizes[2 * k + 1]
|
||
n_in = s_in[0] * s_in[1]
|
||
n_out = s_out[0] * s_out[1]
|
||
num = box[2 * k + 1, 0] * n_out - box[2 * k, 0] * n_in
|
||
den = box[2 * k + 1, 1] * n_out - box[2 * k, 1] * n_in
|
||
sur = t.where(den > 0.5, num / den, zero)
|
||
enough = den > 8.0
|
||
c = t.where(enough, disp - sur, zero)
|
||
pos = t.clamp_min(c, 0.0) * v
|
||
neg = t.clamp_min(-(disp - sur), 0.0) * enough
|
||
better = pos > on
|
||
on = t.where(better, pos, on)
|
||
on_scale = t.where(better, t.full((), k, dtype=t.int8, device=self.dev), on_scale)
|
||
off = t.maximum(off, neg)
|
||
if k == 1:
|
||
surround_mid = sur
|
||
hole = 1.0 - box[-1, 1]
|
||
on = t.clamp(on, 0.0, 1.0 / max(r_max, 1.0) * 1e4)
|
||
return t.stack([disp, on, off, surround_mid, hole]), on_scale
|
||
|
||
# ------------------------------------------------------------------ кластеризация
|
||
|
||
def _offsets(self, col_reach: int, row_reach: int):
|
||
key = (col_reach, row_reach)
|
||
if key not in self._off:
|
||
pairs = [(dr, dc) for dr in range(0, row_reach + 1)
|
||
for dc in range(-col_reach, col_reach + 1) if not (dr == 0 and dc <= 0)]
|
||
p = self.torch.tensor(pairs, dtype=self.torch.int64, device=self.dev)
|
||
self._off[key] = (p[:, 0:1], p[:, 1:2])
|
||
return self._off[key]
|
||
|
||
def cluster_by_depth(self, mask: np.ndarray, r: np.ndarray, *,
|
||
rel_tol: float = 0.06, abs_tol: float = 0.35,
|
||
col_reach: int = 3, row_reach: int = 2):
|
||
"""То же, что `lobula.cluster_by_depth`, с теми же номерами компонент.
|
||
|
||
Рёбра строятся сразу для всех соседей, компоненты — подвешиванием к
|
||
меньшему номеру со сжатием путей: у каждой компоненты корнем остаётся
|
||
её наименьший луч, и нумерация по корням совпадает со scipy.
|
||
"""
|
||
t = self.torch
|
||
h, w = mask.shape
|
||
if r is self._r_host and self.r is not None:
|
||
rt = self.r.reshape(-1)
|
||
else:
|
||
rt = t.as_tensor(np.ascontiguousarray(r), dtype=t.float32,
|
||
device=self.dev).reshape(-1)
|
||
m = t.as_tensor(np.ascontiguousarray(mask), device=self.dev).reshape(-1)
|
||
idx = t.nonzero(m, as_tuple=True)[0]
|
||
n = idx.numel()
|
||
if n == 0:
|
||
return np.zeros(mask.shape, np.int32), 0
|
||
lut = t.full((h * w,), -1, dtype=t.int64, device=self.dev)
|
||
lut[idx] = t.arange(n, device=self.dev)
|
||
rows = t.div(idx, w, rounding_mode="floor")
|
||
cols = idx - rows * w
|
||
ri = rt[idx]
|
||
dr, dc = self._offsets(col_reach, row_reach)
|
||
nr = rows[None, :] + dr
|
||
nc = cols[None, :] + dc
|
||
ok = (nr < h) & (nc >= 0) & (nc < w)
|
||
nb = lut[t.where(ok, nr * w + nc, t.zeros_like(nr))]
|
||
ok &= nb >= 0
|
||
rb = ri[nb.clamp(min=0)]
|
||
ra = ri[None, :].expand_as(rb)
|
||
ok &= t.abs(ra - rb) <= (abs_tol + rel_tol * t.minimum(ra, rb))
|
||
s = t.arange(n, device=self.dev)[None, :].expand_as(nb)[ok]
|
||
d = nb[ok]
|
||
|
||
parent = t.arange(n, device=self.dev)
|
||
while True:
|
||
for _ in range(4):
|
||
ps, pd = parent[s], parent[d]
|
||
parent.scatter_reduce_(0, t.maximum(ps, pd), t.minimum(ps, pd), reduce="amin")
|
||
parent = parent[parent]
|
||
parent = parent[parent]
|
||
if not bool((parent[s] != parent[d]).any()):
|
||
break
|
||
while True:
|
||
nxt = parent[parent]
|
||
if bool(t.equal(nxt, parent)):
|
||
break
|
||
parent = nxt
|
||
roots, comp = t.unique(parent, return_inverse=True)
|
||
labels = t.zeros(h * w, dtype=t.int32, device=self.dev)
|
||
labels[idx] = (comp + 1).to(t.int32)
|
||
return labels.reshape(h, w).cpu().numpy(), int(roots.numel())
|