Lidar_Muxa/flyguard/brain_atlas.py

364 lines
20 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

"""Мозг как облако из 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)