forked from Dan4ick/Lidar_Muxa
369 lines
18 KiB
Python
369 lines
18 KiB
Python
"""Схема мозга дрозофилы, подсвеченная живой активностью конвейера.
|
||
|
||
Каждая стадия FlyGuard соответствует конкретному нейропилю, поэтому состояние
|
||
системы показывается не графиками, а анатомической схемой: где сейчас «горит»,
|
||
там и идёт работа. На демонстрации это замыкает цепочку из ТЗ —
|
||
тоннель → облако точек → мозг мухи → «препятствие, 87 м».
|
||
|
||
Вид дорсальный, схематичный, но с правильным взаимным расположением: по краям
|
||
оптические доли (ламина → медулла → лобула и лобулярная пластинка), в центре
|
||
грибовидные тела (чашечка → ножка → доли) и центральный комплекс
|
||
(протоцеребральный мост, веерное тело, эллипсоидное тело, нодули), снизу —
|
||
нисходящие нейроны с гигантским волокном.
|
||
|
||
Отрисовка стоит 2–4 мс и вынесена под флаг `brain_view`, чтобы не попадать
|
||
в замеры задержки основного тракта.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import array
|
||
|
||
import numpy as np
|
||
|
||
try:
|
||
import cv2
|
||
except ImportError: # без OpenCV вид просто выключен
|
||
cv2 = None
|
||
|
||
# Встроенные шрифты OpenCV (Hershey) знают только латиницу: любая кириллица
|
||
# выходит рядом вопросительных знаков. Подписи здесь русские и составляют
|
||
# половину смысла схемы, поэтому текст рисуется через PIL шрифтом DejaVu.
|
||
try:
|
||
from PIL import Image, ImageDraw, ImageFont
|
||
except ImportError:
|
||
Image = ImageDraw = ImageFont = None
|
||
|
||
_FONT_PATHS = (
|
||
"/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
|
||
"/usr/share/fonts/truetype/dejavu/DejaVuSansCondensed.ttf",
|
||
"/usr/share/fonts/TTF/DejaVuSans.ttf",
|
||
"C:/Windows/Fonts/segoeui.ttf",
|
||
"C:/Windows/Fonts/arial.ttf",
|
||
)
|
||
_font_cache: dict[int, object] = {}
|
||
|
||
|
||
def _font(px: int):
|
||
"""Шрифт нужного кегля; None — если ни одного файла не нашлось."""
|
||
if ImageFont is None:
|
||
return None
|
||
if px not in _font_cache:
|
||
f = None
|
||
for path in _FONT_PATHS:
|
||
try:
|
||
f = ImageFont.truetype(path, px)
|
||
break
|
||
except OSError:
|
||
continue
|
||
_font_cache[px] = f or ImageFont.load_default()
|
||
return _font_cache[px]
|
||
|
||
|
||
def _px(scale: float) -> int:
|
||
"""Кегль PIL, дающий примерно ту же высоту, что и Hershey с этим масштабом."""
|
||
return max(9, int(round(scale * 30)))
|
||
|
||
|
||
def text_size(txt: str, scale: float) -> tuple[int, int]:
|
||
f = _font(_px(scale))
|
||
if f is None:
|
||
(tw, th), _ = cv2.getTextSize(txt, cv2.FONT_HERSHEY_SIMPLEX, scale, 1)
|
||
return tw, th
|
||
box = f.getbbox(txt)
|
||
return box[2] - box[0], box[3] - box[1]
|
||
|
||
|
||
def put_text(img, txt: str, org, scale: float, colour, weight: int = 1) -> None:
|
||
"""Надпись с кириллицей. `org` — левый нижний угол, как у cv2.putText.
|
||
|
||
Рисуется по вырезанному прямоугольнику, а не по всему полотну: PIL копирует
|
||
изображение при каждом преобразовании, и на полном кадре полтора десятка
|
||
подписей стоили 90 мс — больше, чем весь конвейер.
|
||
"""
|
||
f = _font(_px(scale))
|
||
if f is None or Image is None:
|
||
cv2.putText(img, txt, org, cv2.FONT_HERSHEY_SIMPLEX, scale, colour,
|
||
weight, cv2.LINE_AA)
|
||
return
|
||
x, y = int(org[0]), int(org[1])
|
||
box = f.getbbox(txt)
|
||
pad = 2
|
||
# Чернила лягут в x … x+(box[2]-box[0]) по горизонтали и
|
||
# y-(box[3]-box[1]) … y по вертикали — вырез берётся ровно по ним.
|
||
x0 = max(x - pad, 0)
|
||
y0 = max(y - (box[3] - box[1]) - pad, 0)
|
||
x1 = min(x + (box[2] - box[0]) + pad + 1, img.shape[1])
|
||
y1 = min(y + pad + 1, img.shape[0])
|
||
if x1 <= x0 or y1 <= y0:
|
||
return
|
||
crop = img[y0:y1, x0:x1]
|
||
pil = Image.fromarray(crop[:, :, ::-1])
|
||
ImageDraw.Draw(pil).text((x - box[0] - x0, y - box[3] - y0), txt, font=f,
|
||
fill=tuple(int(c) for c in colour[::-1]))
|
||
crop[:, :, :] = np.asarray(pil)[:, :, ::-1]
|
||
|
||
W, H = 1180, 620
|
||
BG = (22, 20, 26)
|
||
INK = (150, 150, 165)
|
||
|
||
# шкала «покой → работа → тревога» в BGR
|
||
_STOPS = (
|
||
(0.00, (58, 48, 44)),
|
||
(0.30, (120, 96, 52)),
|
||
(0.55, (168, 178, 64)),
|
||
(0.78, (72, 196, 240)),
|
||
(1.00, (66, 66, 248)),
|
||
)
|
||
|
||
|
||
def _colour(a: float) -> tuple[int, int, int]:
|
||
a = float(np.clip(a, 0.0, 1.0))
|
||
for (x0, c0), (x1, c1) in zip(_STOPS, _STOPS[1:]):
|
||
if a <= x1:
|
||
t = (a - x0) / max(x1 - x0, 1e-6)
|
||
return tuple(int(c0[i] + (c1[i] - c0[i]) * t) for i in range(3))
|
||
return _STOPS[-1][1]
|
||
|
||
|
||
def _annulus(cx, cy, r_out, r_in, a0, a1, squash=1.0):
|
||
"""Кольцевой сектор — форма слоя оптической доли."""
|
||
# ellipse2Poly требует положительный шаг, поэтому внутренняя дуга строится
|
||
# в прямом направлении и разворачивается уже массивом
|
||
outer = cv2.ellipse2Poly((cx, cy), (r_out, int(r_out * squash)), 0, a0, a1, 4)
|
||
inner = cv2.ellipse2Poly((cx, cy), (r_in, int(r_in * squash)), 0, a0, a1, 4)
|
||
return np.vstack([outer, inner[::-1]])
|
||
|
||
|
||
class Region:
|
||
__slots__ = ("key", "ru", "pts", "anchor", "label_at")
|
||
|
||
def __init__(self, key, ru, pts, label_at=None):
|
||
self.key = key
|
||
self.ru = ru
|
||
self.pts = np.asarray(pts, np.int32)
|
||
self.anchor = tuple(self.pts.mean(axis=0).astype(int))
|
||
self.label_at = label_at
|
||
|
||
|
||
def _build_regions() -> list[Region]:
|
||
r: list[Region] = []
|
||
for side, (cx, a0, a1, lx) in enumerate(((250, 118, 242, 18), (W - 250, -62, 62, W - 250))):
|
||
cy = 300
|
||
mirror = side == 1
|
||
spec = [
|
||
("lamina", "ламина L1 / L2", 218, 188),
|
||
("medulla", "медулла T4 / T5", 182, 132),
|
||
("lobula", "лобула LC11", 126, 88),
|
||
("lobula_plate", "лоб. пластинка LPTC / LPLC2", 82, 52),
|
||
]
|
||
for i, (key, ru, ro, ri) in enumerate(spec):
|
||
pts = _annulus(cx, cy, ro, ri, a0 + i * 3, a1 - i * 3, squash=1.06)
|
||
ly = 96 + i * 26
|
||
r.append(Region(key, ru, pts, (30 if not mirror else W - 330, ly)))
|
||
|
||
cx = W // 2
|
||
|
||
# грибовидные тела: чашечка сзади → ножка вперёд → вертикальная и медиальная
|
||
# доли. Они лежат ЛАТЕРАЛЬНЕЕ центрального комплекса, не поверх него.
|
||
for mirror in (False, True):
|
||
s = -1 if mirror else 1
|
||
bx = cx + s * 150
|
||
calyx = cv2.ellipse2Poly((bx, 132), (46, 33), 0, 0, 360, 6)
|
||
ped = np.array([[bx - s * 16, 160], [bx + s * 18, 158],
|
||
[bx + s * 46, 318], [bx + s * 18, 322]], np.int32)
|
||
lobes = np.array([[bx + s * 16, 316], [bx + s * 50, 314], [bx + s * 56, 404],
|
||
[bx + s * 26, 408], [bx + s * 26, 372], [bx + s * 16, 372]],
|
||
np.int32)
|
||
r.append(Region("mb_calyx", "чашечка PN → KC", calyx,
|
||
(cx - 330, 86) if not mirror else (cx + 210, 86)))
|
||
r.append(Region("mb_ped", "ножка · APL", ped, None))
|
||
r.append(Region("mb_lobes", "доли KC → MBON (новизна)", lobes,
|
||
(cx - 400, 432) if not mirror else (cx + 230, 432)))
|
||
|
||
pb = _annulus(cx, 232, 104, 86, 200, 340, squash=0.7)
|
||
fb = np.array([[cx - 74, 262], [cx + 74, 262], [cx + 58, 314], [cx - 58, 314]], np.int32)
|
||
eb = _annulus(cx, 372, 46, 26, 0, 360)
|
||
no_l = cv2.ellipse2Poly((cx - 52, 420), (18, 13), 0, 0, 360, 8)
|
||
no_r = cv2.ellipse2Poly((cx + 52, 420), (18, 13), 0, 0, 360, 8)
|
||
dn = np.array([[cx - 22, 444], [cx + 22, 444], [cx + 40, 566], [cx - 40, 566]], np.int32)
|
||
|
||
r += [
|
||
Region("pb", "протоцеребральный мост", pb, (cx - 96, 176)),
|
||
Region("fb", "веерное тело", fb, (cx - 40, 342)),
|
||
Region("eb", "эллипсоидное тело — кольцевой аттрактор", eb, (cx - 300, 372)),
|
||
Region("no", "нодули", no_l, None),
|
||
Region("no", "нодули", no_r, (cx + 96, 428)),
|
||
Region("dn", "нисходящие нейроны · гигантское волокно", dn, (cx - 300, 524)),
|
||
]
|
||
return r
|
||
|
||
|
||
class BrainView:
|
||
"""Отрисовка схемы с активностью."""
|
||
|
||
def __init__(self, labels: bool = True):
|
||
self.enabled = cv2 is not None
|
||
self.labels = labels
|
||
self.regions = _build_regions() if self.enabled else []
|
||
self._base = self._draw_base() if self.enabled else None
|
||
self._smooth: dict[str, float] = {}
|
||
|
||
# ------------------------------------------------------------------ статика
|
||
|
||
def _draw_base(self) -> np.ndarray:
|
||
img = np.full((H, W, 3), BG, np.uint8)
|
||
put_text(img, "FlyGuard — схемы мозга Drosophila melanogaster в реальном времени",
|
||
(26, 36), 0.66, (205, 205, 215))
|
||
cv2.line(img, (26, 48), (W - 26, 48), (52, 52, 62), 1)
|
||
return img
|
||
|
||
# ------------------------------------------------------------------ активность
|
||
|
||
@staticmethod
|
||
def activity(res) -> dict[str, float]:
|
||
"""Свести состояние кадра к активности по нейропилям, 0…1."""
|
||
lam = res.lam
|
||
d = res.decision
|
||
on = float(np.mean(lam.on)) * 130.0 if lam is not None else 0.0
|
||
off = float(np.mean(lam.off)) * 130.0 if lam is not None else 0.0
|
||
hole = float(np.mean(lam.hole)) if lam is not None else 0.0
|
||
speed = abs(res.ego.speed) / 20.0 if res.ego else 0.0
|
||
loom = (float(np.percentile(res.loom, 99.9)) * 25.0
|
||
if res.loom is not None and res.loom.size else 0.0)
|
||
nov = max((c.novelty for c in res.candidates), default=0.0)
|
||
# Дофамин: у мухи нейроны PPL1/PAM управляют пластичностью синапсов
|
||
# KC→MBON — ровно тем, чем учится наша память тоннеля. Осмысленная
|
||
# величина здесь — сила сигнала обучения: насколько сцена в целом
|
||
# незнакома памяти. В отличие от `nov` (самый незнакомый предмет)
|
||
# это средняя незнакомость всего, что видно, и она заранее говорит,
|
||
# что участок новый и ложных тревог будет больше.
|
||
dopa = (float(np.mean([c.novelty for c in res.candidates]))
|
||
if res.candidates else 0.0)
|
||
n_cand = min(len(res.candidates) / 7.0, 1.0)
|
||
evid = max((o.confidence for o in d.objects), default=0.0)
|
||
|
||
valid = float(lam.disp.astype(bool).mean()) if lam is not None else 0.0
|
||
|
||
return {
|
||
# Сетчатка «горит» тем сильнее, чем больше лучей вернулось: это и
|
||
# есть освещённость омматидиальной решётки.
|
||
"retina": np.clip(0.15 + 0.85 * valid, 0, 1),
|
||
"looming": np.clip(loom, 0, 1),
|
||
"dopamine": np.clip(dopa, 0, 1),
|
||
# Привыкание: доля синапсов KC→MBON, уже подавленных повторами
|
||
# внутри этого проезда. Ноль в начале участка, растёт по мере
|
||
# того, как однообразная обстановка перестаёт быть новостью.
|
||
"habituation": np.clip(getattr(res, "hab", 0.0), 0, 1),
|
||
"lamina": np.clip(0.55 * on + 0.45 * off, 0, 1),
|
||
"medulla": np.clip(0.25 + 0.75 * speed, 0, 1),
|
||
"lobula": np.clip(n_cand, 0, 1),
|
||
"lobula_plate": np.clip(max(speed, loom), 0, 1),
|
||
"mb_calyx": np.clip(n_cand, 0, 1),
|
||
"mb_ped": np.clip(0.4 * n_cand, 0, 1),
|
||
"mb_lobes": np.clip(nov, 0, 1),
|
||
"pb": np.clip(speed, 0, 1),
|
||
"fb": np.clip(hole, 0, 1),
|
||
"eb": np.clip(evid, 0, 1),
|
||
"no": np.clip(abs(res.ego.yaw_deg) / 1.5 if res.ego else 0.0, 0, 1),
|
||
"dn": 1.0 if d.emergency else (0.7 if d.detected else 0.06),
|
||
}
|
||
|
||
# ------------------------------------------------------------------ кадр
|
||
|
||
def render(self, res) -> np.ndarray | None:
|
||
if not self.enabled:
|
||
return None
|
||
for k, v in self.activity(res).items(): # сглаживание: мозг не мигает
|
||
self._smooth[k] = 0.6 * self._smooth.get(k, v) + 0.4 * v
|
||
|
||
img = self._base.copy()
|
||
for reg in self.regions:
|
||
a = self._smooth.get(reg.key, 0.0)
|
||
cv2.fillPoly(img, [reg.pts], _colour(a))
|
||
cv2.polylines(img, [reg.pts], True, (46, 44, 54), 1, cv2.LINE_AA)
|
||
|
||
self._draw_bump(img, res)
|
||
self._draw_flow(img)
|
||
if self.labels:
|
||
self._draw_labels(img)
|
||
self._draw_readout(img, res)
|
||
return img
|
||
|
||
def _draw_labels(self, img) -> None:
|
||
seen = set()
|
||
for reg in self.regions:
|
||
if reg.label_at is None or reg.ru in seen:
|
||
continue
|
||
seen.add(reg.ru)
|
||
x, y = reg.label_at
|
||
tw, th = text_size(reg.ru, 0.44)
|
||
# подложка под текстом: подписи неизбежно ложатся поверх нейропилей,
|
||
# и без неё их не прочитать
|
||
box = img[max(y - th - 6, 0):y + 6, max(x - 6, 0):x + tw + 6]
|
||
if box.size:
|
||
box[:] = (box * 0.25).astype(np.uint8)
|
||
cv2.line(img, (x + tw // 2, y - th - 6 if y > reg.anchor[1] else y + 6),
|
||
reg.anchor, (74, 72, 86), 1, cv2.LINE_AA)
|
||
put_text(img, reg.ru, (x, y), 0.44, INK)
|
||
|
||
def _draw_bump(self, img, res) -> None:
|
||
"""Бугор активности в эллипсоидном теле — направление на ближайший объект."""
|
||
cx, cy, r = W // 2, 380, 36
|
||
objs = res.decision.objects
|
||
if not objs:
|
||
return
|
||
o = objs[0]
|
||
ang = float(np.clip(np.arctan2(o.lateral, max(o.distance, 1.0)) * 7.0, -3.0, 3.0))
|
||
px = int(cx + r * np.sin(ang))
|
||
py = int(cy - r * np.cos(ang))
|
||
cv2.circle(img, (px, py), 11, _colour(min(o.confidence + 0.35, 1.0)), -1, cv2.LINE_AA)
|
||
cv2.circle(img, (px, py), 11, (230, 230, 240), 1, cv2.LINE_AA)
|
||
|
||
def _draw_flow(self, img) -> None:
|
||
cx = W // 2
|
||
for a, b in (((330, 300), (410, 260)), ((W - 330, 300), (W - 410, 260)),
|
||
((cx - 60, 336), (cx - 30, 356)), ((cx + 60, 336), (cx + 30, 356)),
|
||
((cx, 420), (cx, 428))):
|
||
cv2.arrowedLine(img, a, b, (86, 84, 98), 1, cv2.LINE_AA, tipLength=0.3)
|
||
|
||
def _draw_readout(self, img, res) -> None:
|
||
d = res.decision
|
||
rows = [
|
||
("скорость", f"{res.ego.kmh:.1f} км/ч" if res.ego else "—"),
|
||
("кандидатов", str(len(res.candidates))),
|
||
("новизна MBON", f"{max((c.novelty for c in res.candidates), default=0.0):.2f}"),
|
||
("треков", str(len(d.objects))),
|
||
("обработка", f"{res.total_ms:.1f} мс"),
|
||
]
|
||
y0 = H - 150
|
||
for i, (k, v) in enumerate(rows):
|
||
put_text(img, k, (28, y0 + i * 24), 0.46, (118, 116, 128))
|
||
put_text(img, v, (170, y0 + i * 24), 0.5, (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, (W - 470, H - 90), 0.95, colour, 2)
|
||
if sub:
|
||
put_text(img, sub, (W - 470, H - 58), 0.56, colour)
|
||
|
||
# ------------------------------------------------------------------ ROS
|
||
|
||
@staticmethod
|
||
def to_msg(img: np.ndarray, stamp, frame_id: str):
|
||
from sensor_msgs.msg import Image
|
||
msg = Image()
|
||
msg.header.stamp = stamp
|
||
msg.header.frame_id = frame_id
|
||
msg.height, msg.width = img.shape[:2]
|
||
msg.encoding = "bgr8"
|
||
msg.is_bigendian = 0
|
||
msg.step = img.shape[1] * 3
|
||
# array('B'), а не bytes: из bytes rclpy проверяет каждый байт на Python,
|
||
# и картинка в 2.8 МБ стоила 142 мс на кадр против 2.4 мс (замерено).
|
||
msg.data = array.array("B", img.tobytes())
|
||
return msg
|