forked from Dan4ick/Lidar_Muxa
364 lines
20 KiB
Python
364 lines
20 KiB
Python
"""Мозг как облако из 139 255 нейронов, подсвеченное живой активностью.
|
||
|
||
Отличие от `brain_view` — не в красоте, а в том, чем нарисовано. Там схема,
|
||
нарисованная руками; здесь каждая точка — **настоящий нейрон из коннектома**
|
||
FlyWire FAFB v783, стоящий на своём анатомическом месте, а его яркость берётся
|
||
от той стадии FlyGuard, которой он соответствует.
|
||
|
||
Симуляции нет и не нужно. Мембранные потенциалы 139 тысяч клеток никто не
|
||
интегрирует: коннектом даёт анатомию и принадлежность клеток, конвейер даёт
|
||
активность по стадиям, а здесь одно накладывается на другое. Стоимость кадра —
|
||
взвешенная сумма десяти заранее посчитанных карт плотности.
|
||
|
||
Привязка держится на именах типов клеток, а не на догадках. FlyGuard с самого
|
||
начала собран из конкретных схем, и все они в выгрузке есть поимённо:
|
||
|
||
LC11 127 нейронов детектор мелкого объекта → lobula.py
|
||
LPLC2 210 детектор надвигания → medulla.py
|
||
HS/VS 22 широкопольный поток → medulla.py
|
||
T4/T5 12245 элементарные детекторы движения
|
||
KC 5177 клетки Кеньона, разрежённый код → mushroom_body.py
|
||
MBON 96 выход памяти, новизна
|
||
APL 2 глобальное торможение
|
||
DNp01 2 гигантское волокно, экстренное решение → descending.py
|
||
|
||
Атлас собирается `tools/build_brain_atlas.py` и лежит в `data/brain_atlas.npz`
|
||
(605 КБ). Данные FlyWire — CC-BY 4.0, Dorkenwald et al. и Schlegel et al.,
|
||
Nature 2024.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
from collections import deque
|
||
from pathlib import Path
|
||
|
||
import numpy as np
|
||
|
||
try:
|
||
import cv2
|
||
except ImportError:
|
||
cv2 = None
|
||
|
||
from .brain_view import BG, INK, _colour, put_text, text_size
|
||
|
||
ATLAS = Path(__file__).with_name("data") / "brain_atlas.npz"
|
||
|
||
# Какая стадия конвейера освещает какую группу нейронов. Ключи справа — из
|
||
# `BrainView.activity`; так обе картинки всегда показывают одно и то же.
|
||
STAGE_SOURCE: dict[str, tuple[str, ...]] = {
|
||
"retina": ("retina",),
|
||
"lamina": ("lamina",),
|
||
"medulla": ("medulla",),
|
||
"lptc": ("lobula_plate",),
|
||
"looming": ("looming",),
|
||
"lobula": ("lobula",),
|
||
"mushroom": ("mb_lobes", "mb_calyx"),
|
||
"central": ("eb", "pb", "fb"),
|
||
"descending": ("dn",),
|
||
"other": (),
|
||
}
|
||
OTHER_BASE = 0.045 # остальной мозг: виден, но не отвлекает
|
||
TRACE_N = 100 # длина ритмограммы: 10 с при 10 Гц
|
||
|
||
RU = {
|
||
"retina": "сетчатка R1–R8",
|
||
"lamina": "ламина L1 / L2",
|
||
"medulla": "медулла T4 / T5",
|
||
"lptc": "лоб. пластинка HS / VS",
|
||
"looming": "LPLC2 надвигание",
|
||
"lobula": "лобула LC11",
|
||
"mushroom": "грибовидное тело KC → MBON",
|
||
"central": "центральный комплекс",
|
||
"descending": "нисходящие DNp01",
|
||
}
|
||
|
||
|
||
class NeuronCloud:
|
||
"""Облако нейронов в анатомических координатах, яркость — активность стадии."""
|
||
|
||
def __init__(self, labels: bool = True, gain: float = 1.0, scale: int = 1):
|
||
self.enabled = cv2 is not None and ATLAS.exists()
|
||
self.labels = labels
|
||
self.gain = gain
|
||
# Масштаб отрисовки: 1 — для узла (1180 × 620), 2–3 — для экрана и
|
||
# видео в 2K/4K. Сомы берутся из атласа своего размера
|
||
# (brain_atlas_x2.npz …) — это настоящие координаты, а не растянутая
|
||
# картинка; шрифты, линии и отступы растут вместе с масштабом.
|
||
self.s = max(1, int(scale))
|
||
self._smooth: dict[str, float] = {}
|
||
self._trace: deque = deque(maxlen=TRACE_N)
|
||
if not self.enabled:
|
||
return
|
||
|
||
path = ATLAS.with_name(f"brain_atlas_x{self.s}.npz") if self.s > 1 else ATLAS
|
||
d = np.load(path if path.exists() else ATLAS, allow_pickle=False)
|
||
self.h, self.w = (int(v) for v in d["shape"])
|
||
stages = [str(s) for s in d["stages"]]
|
||
px, py = d["px"].astype(np.int64), d["py"].astype(np.int64)
|
||
stage, side, depth = d["stage"], d["side"], d["depth"]
|
||
if not path.exists(): # атласа нужного размера нет — растянуть обычный
|
||
px, py, self.h, self.w = px * self.s, py * self.s, self.h * self.s, self.w * self.s
|
||
|
||
ok = (px >= 0) & (px < self.w) & (py >= 0) & (py < self.h)
|
||
flat = py * self.w + px
|
||
# Глубина: дальние сомы тусклее, иначе фронтальная проекция выглядит плоской
|
||
shade = (0.55 + 0.45 * (1.0 - depth / 255.0)).astype(np.float32)
|
||
|
||
# Сомы занимают шестую часть холста; считать по всему полотну незачем.
|
||
# Карты сразу строятся по занятым пикселям: это убирает из каждого
|
||
# кадра 80 % работы, а в 4K — ещё и сотни мегабайт промежуточных карт.
|
||
self.support = np.unique(flat[ok])
|
||
at = np.searchsorted(self.support, flat)
|
||
|
||
# Две карты на стадию — левое и правое полушарие: так видно, с какой
|
||
# стороны пришёл объект, ровно как у мухи с её двумя оптическими долями.
|
||
self.keys: list[str] = []
|
||
maps = []
|
||
for i, name in enumerate(stages):
|
||
for sd, tag in ((0, "L"), (1, "R"), (2, "C")):
|
||
m = ok & (stage == i) & (side == sd)
|
||
if m.sum() == 0:
|
||
continue
|
||
self.keys.append(f"{name}|{tag}")
|
||
maps.append(np.bincount(at[m], weights=shade[m],
|
||
minlength=self.support.size).astype(np.float32))
|
||
self.maps = np.ascontiguousarray(np.stack(maps)) # (K, занятые пиксели)
|
||
self.stage_of = [k.split("|")[0] for k in self.keys]
|
||
self.side_of = [k.split("|")[1] for k in self.keys]
|
||
|
||
# Якоря подписей — центры тяжести стадии в своём полушарии. Общий
|
||
# центр тяжести у всех стадий приходится на середину мозга и ничего
|
||
# не показывает: стадии парные.
|
||
self.anchor: dict[str, tuple[tuple[int, int], tuple[int, int]]] = {}
|
||
for i, name in enumerate(stages):
|
||
if name not in RU:
|
||
continue
|
||
a = []
|
||
for sd in (0, 1):
|
||
m = ok & (stage == i) & (side == sd)
|
||
m = m if m.sum() > 30 else (ok & (stage == i))
|
||
a.append((int(px[m].mean()), int(py[m].mean())))
|
||
self.anchor[name] = tuple(a)
|
||
|
||
self._lut = self._build_lut()
|
||
self._base = self._draw_base()
|
||
|
||
# ------------------------------------------------------------------ статика
|
||
|
||
@staticmethod
|
||
def _build_lut() -> np.ndarray:
|
||
"""Палитра «покой → работа → тревога», та же, что у схемы."""
|
||
lut = np.zeros((256, 3), np.uint8)
|
||
for i in range(256):
|
||
lut[i] = _colour(i / 255.0)
|
||
return lut
|
||
|
||
def _draw_base(self) -> np.ndarray:
|
||
s = self.s
|
||
img = np.full((self.h, self.w, 3), BG, np.uint8)
|
||
put_text(img, "FlyGuard — 139 255 нейронов коннектома FlyWire, "
|
||
"подсвеченных активностью конвейера",
|
||
(26 * s, 34 * s), 0.6 * s, (205, 205, 215))
|
||
cv2.line(img, (26 * s, 46 * s), (self.w - 26 * s, 46 * s), (52, 52, 62), s)
|
||
return img
|
||
|
||
# ------------------------------------------------------------------ кадр
|
||
|
||
def weights(self, act: dict[str, float], bias: float) -> np.ndarray:
|
||
"""Яркость каждой полукарты. `bias` < 0 — объект слева, > 0 — справа."""
|
||
w = np.empty(len(self.keys), np.float32)
|
||
for i, (stage, side) in enumerate(zip(self.stage_of, self.side_of)):
|
||
src = STAGE_SOURCE.get(stage, ())
|
||
v = max((self._smooth.get(k, 0.0) for k in src), default=0.0)
|
||
v = OTHER_BASE if not src else max(v, OTHER_BASE * 0.7)
|
||
if side in ("L", "R") and stage in ("retina", "lamina", "medulla",
|
||
"lptc", "looming", "lobula"):
|
||
# У мухи зрительный мир перекрещен: предмет справа сильнее
|
||
# освещает левую долю. Для демонстрации важно, что асимметрия
|
||
# видна и соответствует стороне объекта.
|
||
v *= 1.0 + 0.10 * (bias if side == "L" else -bias)
|
||
w[i] = np.clip(v, 0.0, 1.0)
|
||
return w
|
||
|
||
def render(self, res) -> np.ndarray | None:
|
||
if not self.enabled:
|
||
return None
|
||
from .brain_view import BrainView
|
||
|
||
for k, v in BrainView.activity(res).items():
|
||
self._smooth[k] = 0.6 * self._smooth.get(k, v) + 0.4 * v
|
||
|
||
# Ритмограмма: у дрозофилы ритмы мозга меряют электродом, из наших
|
||
# данных они не выводятся. Поэтому здесь честный след собственного
|
||
# состояния — сколько работает зрительный тракт, сколько памяти и
|
||
# когда сработало решение.
|
||
vis = float(np.mean([self._smooth.get(k, 0.0) for k in
|
||
("lamina", "medulla", "lobula", "lobula_plate")]))
|
||
self._trace.append((float(vis),
|
||
float(self._smooth.get("dopamine", 0.0)),
|
||
bool(res.decision.detected),
|
||
float(res.total_ms),
|
||
float(self._smooth.get("habituation", 0.0))))
|
||
|
||
lat = [c.u for c in res.candidates] or [0.0]
|
||
bias = float(np.clip(np.mean(lat) / 1.6, -1.0, 1.0))
|
||
|
||
vals = self.weights(self.activity_smoothed(), bias) @ self.maps
|
||
img = self._base.copy()
|
||
if self.s == 1:
|
||
idx = np.clip(vals * (255.0 * self.gain), 0, 255).astype(np.uint8)
|
||
img.reshape(-1, 3)[self.support] = self._lut[idx]
|
||
# Точки одиночные и на большом холсте теряются: лёгкое размытие делает
|
||
# облако читаемым, не превращая его в заливку.
|
||
cv2.GaussianBlur(img, (3, 3), 0, dst=img)
|
||
else:
|
||
self._paint_fine(img, vals)
|
||
|
||
if self.labels:
|
||
self._draw_labels(img)
|
||
self._draw_rhythm(img)
|
||
self._draw_readout(img, res)
|
||
return img
|
||
|
||
def _paint_fine(self, img, vals) -> None:
|
||
"""Облако в высоком разрешении: каждая сома на своём месте.
|
||
|
||
В масштабе 1 пиксель холста собирает несколько сом, и яркость — их
|
||
сумма. Здесь сома — отдельная точка, поэтому сумма берётся по
|
||
окрестности той же площади (гауссово окно, нормированное так, что
|
||
одиночная сома сохраняет свою яркость), а края точек смешиваются с
|
||
фоном по доле покрытия. Облако выглядит как в узле, только чётче.
|
||
"""
|
||
sig = 0.55 * self.s
|
||
norm = 2.0 * np.pi * sig * sig
|
||
field = np.zeros((2, self.h * self.w), np.float32)
|
||
field[0, self.support] = vals
|
||
field[1, self.support] = 1.0
|
||
v = cv2.GaussianBlur(field[0].reshape(self.h, self.w), (0, 0), sig) * norm
|
||
a = np.clip(cv2.GaussianBlur(field[1].reshape(self.h, self.w), (0, 0), sig)
|
||
* norm, 0.0, 1.0)
|
||
idx = np.clip(v * (255.0 * self.gain), 0, 255).astype(np.uint8)
|
||
col = self._lut[idx]
|
||
img[:] = cv2.blendLinear(col, img, a, 1.0 - a)
|
||
|
||
def activity_smoothed(self) -> dict[str, float]:
|
||
return self._smooth
|
||
|
||
def _draw_labels(self, img) -> None:
|
||
"""Подписи столбиком у ближнего края с выноской к центру тяжести стадии.
|
||
|
||
Центры тяжести половины стадий лежат в середине мозга и налезают друг
|
||
на друга, поэтому подписи разводятся по вертикали, а привязка остаётся
|
||
честной — линия ведёт ровно в центр тяжести.
|
||
"""
|
||
# Зрительный путь подписывается слева, память и решение — справа:
|
||
# так выноски не пересекают мозг и идут в своё же полушарие.
|
||
s = self.s
|
||
left_side = ("retina", "lamina", "medulla", "lptc")
|
||
cols = {True: 76 * s, False: 76 * s}
|
||
for name in ("retina", "lamina", "medulla", "lptc",
|
||
"looming", "lobula", "mushroom", "central", "descending"):
|
||
if name not in self.anchor:
|
||
continue
|
||
left = name in left_side
|
||
ax, ay = self.anchor[name][0 if left else 1]
|
||
txt = RU[name]
|
||
tw, th = text_size(txt, 0.42 * s)
|
||
x = 26 * s if left else self.w - tw - 26 * s
|
||
y = cols[left]
|
||
cols[left] += 27 * s
|
||
box = img[max(y - th - 5 * s, 0):y + 5 * s, max(x - 6 * s, 0):x + tw + 6 * s]
|
||
if box.size:
|
||
box[:] = (box * 0.2).astype(np.uint8)
|
||
put_text(img, txt, (x, y), 0.42 * s, INK)
|
||
tip = (x + tw + 7 * s, y - th // 2) if left else (x - 7 * s, y - th // 2)
|
||
cv2.line(img, tip, (ax, ay), (82, 80, 96), s, cv2.LINE_AA)
|
||
cv2.circle(img, (ax, ay), 3 * s, (162, 160, 178), s, cv2.LINE_AA)
|
||
|
||
def _draw_rhythm(self, img) -> None:
|
||
"""Ритмограмма: след активности и дофамина за последние десять секунд.
|
||
|
||
Это не ритм мозга дрозофилы — тот меряют электродом, и из лидарных
|
||
данных он не выводится. Это след собственного состояния системы:
|
||
зрительный тракт, сигнал обучения и отметки моментов решения.
|
||
"""
|
||
if len(self._trace) < 3:
|
||
return
|
||
s = self.s
|
||
x0, y0, w, hgt = 250 * s, self.h - 96 * s, 380 * s, 56 * s
|
||
cv2.rectangle(img, (x0, y0), (x0 + w, y0 + hgt), (40, 39, 48), -1)
|
||
cv2.rectangle(img, (x0, y0), (x0 + w, y0 + hgt), (60, 58, 70), s)
|
||
put_text(img, "ритм · 10 с", (x0 + 4 * s, y0 - 6 * s), 0.37 * s, (128, 126, 140))
|
||
|
||
n = len(self._trace)
|
||
xs = x0 + (np.arange(n) * (w - 6 * s) / max(TRACE_N - 1, 1)).astype(int) + 3 * s
|
||
tr = np.array([(a, b, float(c), e) for a, b, c, _, e in self._trace],
|
||
np.float32)
|
||
|
||
lx = x0 + 78 * s
|
||
traces = [(0, (176, 196, 92), "зрение"), (1, (86, 178, 240), "дофамин")]
|
||
if tr[:, 3].max() > 0.0:
|
||
traces.append((3, (150, 120, 226), "привыкание"))
|
||
for col, colour, name in traces:
|
||
ys = (y0 + hgt - 8 * s - tr[:, col] * (hgt - 16 * s)).astype(int)
|
||
cv2.polylines(img, [np.stack([xs, ys], 1)], False, colour, s, cv2.LINE_AA)
|
||
cv2.line(img, (lx, y0 - 10 * s), (lx + 11 * s, y0 - 10 * s), colour, 2 * s,
|
||
cv2.LINE_AA)
|
||
put_text(img, name, (lx + 15 * s, y0 - 6 * s), 0.35 * s, (132, 130, 144))
|
||
lx += 15 * s + text_size(name, 0.35 * s)[0] + 16 * s
|
||
|
||
# Полоска решений внизу: где система говорила «препятствие».
|
||
hot = tr[:, 2] > 0.5
|
||
if hot.any():
|
||
for i in np.flatnonzero(hot):
|
||
cv2.line(img, (int(xs[i]), y0 + hgt - 5 * s), (int(xs[i]), y0 + hgt - 2 * s),
|
||
(66, 66, 248), s)
|
||
cv2.line(img, (lx, y0 - 10 * s), (lx + 11 * s, y0 - 10 * s), (66, 66, 248), 2 * s)
|
||
put_text(img, "решение", (lx + 15 * s, y0 - 6 * s), 0.35 * s, (132, 130, 144))
|
||
|
||
def _draw_readout(self, img, res) -> None:
|
||
d = res.decision
|
||
hab = float(getattr(res, "hab", 0.0))
|
||
rows = [("скорость", f"{abs(d.speed) * 3.6:.1f} км/ч"),
|
||
("кандидатов", f"{len(res.candidates)}"),
|
||
("новизна MBON", f"{max((c.novelty for c in res.candidates), default=0.0):.2f}"),
|
||
("дофамин PPL1/PAM", f"{self._smooth.get('dopamine', 0.0):.2f}")]
|
||
# привыкание по умолчанию выключено — строка появляется, только если оно
|
||
# включено и уже что-то накопило
|
||
if hab > 0.0:
|
||
rows.append(("привыкание", f"{hab:.2f}"))
|
||
rows.append(("обработка", f"{res.total_ms:.1f} мс"))
|
||
s = self.s
|
||
y0 = self.h - 126 * s - 23 * s * (len(rows) - 5)
|
||
# Столбик лежит поверх облака, и серые подписи в нём тонули в точках
|
||
# оптической доли — подложка, как под выносками стадий.
|
||
vw = max(text_size(v, 0.46 * s)[0] for _, v in rows)
|
||
right = min(196 * s + vw + 10 * s, 244 * s) # не заходить на ритмограмму
|
||
box = img[y0 - 18 * s:y0 + 23 * s * (len(rows) - 1) + 8 * s, 18 * s:right]
|
||
box[:] = (box * 0.25).astype(np.uint8)
|
||
for i, (k, v) in enumerate(rows):
|
||
put_text(img, k, (28 * s, y0 + i * 23 * s), 0.44 * s, (118, 116, 128))
|
||
put_text(img, v, (196 * s, y0 + i * 23 * s), 0.46 * s, (206, 206, 216))
|
||
|
||
if d.detected:
|
||
txt = f"ПРЕПЯТСТВИЕ {d.distance:.0f} м"
|
||
sub = ("ЭКСТРЕННОЕ ТОРМОЖЕНИЕ" if d.emergency
|
||
else f"уверенность {d.confidence:.2f}")
|
||
colour = (66, 66, 248) if d.emergency else (72, 196, 240)
|
||
else:
|
||
txt, sub, colour = "ПУТЬ СВОБОДЕН", "", (120, 214, 130)
|
||
put_text(img, txt, (self.w - 470 * s, self.h - 62 * s), 0.9 * s, colour, 2)
|
||
if sub:
|
||
put_text(img, sub, (self.w - 470 * s, self.h - 32 * s), 0.54 * s, colour)
|
||
|
||
x = self.w - 300 * s
|
||
for lvl, name in ((0.10, "покой"), (0.45, "работа"), (0.85, "тревога")):
|
||
cv2.circle(img, (x, 34 * s), 5 * s, _colour(lvl), -1, cv2.LINE_AA)
|
||
put_text(img, name, (x + 12 * s, 39 * s), 0.4 * s, (150, 148, 162))
|
||
x += 14 * s + text_size(name, 0.4 * s)[0] + 22 * s
|
||
|
||
# ------------------------------------------------------------------ ROS
|
||
|
||
@staticmethod
|
||
def to_msg(img: np.ndarray, stamp, frame_id: str):
|
||
from .brain_view import BrainView
|
||
return BrainView.to_msg(img, stamp, frame_id)
|