21e83be26e
Add cleaned RealSense/OpenCV CV under cv/ with secret-free example config, MQTT disabled by default, and canonical README coverage for web + CV relationship without claiming production integration. Co-authored-by: Cursor <cursoragent@cursor.com>
303 lines
13 KiB
Python
303 lines
13 KiB
Python
"""HUD демо: RGB-подложка + depth, несколько объектов с ID, кириллица через PIL."""
|
||
|
||
from __future__ import annotations
|
||
|
||
from functools import lru_cache
|
||
from typing import Dict, List, Optional, Tuple
|
||
|
||
import cv2
|
||
import numpy as np
|
||
from PIL import Image, ImageDraw, ImageFont
|
||
|
||
ZONE_COLOR = { # BGR
|
||
"B": (40, 180, 40),
|
||
"C": (40, 40, 220),
|
||
"D": (0, 165, 255),
|
||
}
|
||
UNCERTAIN_COLOR = (0, 130, 250) # оранжевый
|
||
PENDING_COLOR = (0, 255, 255) # жёлтый — идёт накопление
|
||
|
||
_FONT_CANDIDATES = [
|
||
"/usr/share/fonts/noto/NotoSans-Bold.ttf",
|
||
"/usr/share/fonts/noto/NotoSans-Regular.ttf",
|
||
"/usr/share/fonts/TTF/DejaVuSans-Bold.ttf",
|
||
"/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",
|
||
]
|
||
|
||
|
||
@lru_cache(maxsize=8)
|
||
def _font(size: int) -> ImageFont.FreeTypeFont:
|
||
for path in _FONT_CANDIDATES:
|
||
try:
|
||
return ImageFont.truetype(path, size)
|
||
except OSError:
|
||
continue
|
||
return ImageFont.load_default()
|
||
|
||
|
||
def _draw_texts(
|
||
img_bgr: np.ndarray,
|
||
texts: List[Tuple[int, int, str, Tuple[int, int, int], int]],
|
||
) -> np.ndarray:
|
||
"""texts: (x, y, строка, цвет BGR, размер). Кириллица через PIL."""
|
||
if not texts:
|
||
return img_bgr
|
||
pil = Image.fromarray(cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB))
|
||
draw = ImageDraw.Draw(pil)
|
||
for x, y, s, bgr, size in texts:
|
||
rgb = (bgr[2], bgr[1], bgr[0])
|
||
draw.text((x, y), s, font=_font(size), fill=rgb, stroke_width=2, stroke_fill=(0, 0, 0))
|
||
return cv2.cvtColor(np.array(pil), cv2.COLOR_RGB2BGR)
|
||
|
||
|
||
class DepthSmoother:
|
||
"""Временное сглаживание depth только для отображения (не для измерений)."""
|
||
|
||
def __init__(self, alpha: float = 0.25) -> None:
|
||
self.alpha = float(alpha)
|
||
self._acc: Optional[np.ndarray] = None
|
||
|
||
def update(self, depth_mm: np.ndarray) -> np.ndarray:
|
||
d = depth_mm.astype(np.float32)
|
||
if self._acc is None or self._acc.shape != d.shape:
|
||
self._acc = d.copy()
|
||
valid = d > 0
|
||
self._acc[valid] = (1.0 - self.alpha) * self._acc[valid] + self.alpha * d[valid]
|
||
out = self._acc.astype(np.uint16)
|
||
out[~valid & (self._acc <= 0)] = 0
|
||
return out
|
||
|
||
|
||
class ContourSmoother:
|
||
"""Стабильная окантовка: EMA маски по каждому треку + аппроксимация контура,
|
||
плюс «примагничивание» контура к краям объекта на RGB (снимает остаточный
|
||
параллакс depth↔color и распухание depth-маски).
|
||
|
||
Только для отрисовки — измерения идут по сырому контуру.
|
||
"""
|
||
|
||
def __init__(self, alpha: float = 0.3, snap_alpha: float = 0.35) -> None:
|
||
self.alpha = float(alpha)
|
||
self.snap_alpha = float(snap_alpha)
|
||
self._acc: Dict[int, np.ndarray] = {}
|
||
self._snap: Dict[int, Tuple[float, float, float]] = {} # tid -> (dx, dy, shrink)
|
||
|
||
def smooth(
|
||
self, track_id: int, mask: np.ndarray, edge_img: Optional[np.ndarray] = None
|
||
) -> Optional[np.ndarray]:
|
||
m = mask.astype(np.float32) / 255.0
|
||
acc = self._acc.get(track_id)
|
||
if acc is None or acc.shape != m.shape:
|
||
acc = m.copy()
|
||
else:
|
||
acc = (1.0 - self.alpha) * acc + self.alpha * m
|
||
self._acc[track_id] = acc
|
||
|
||
soft = cv2.GaussianBlur((acc * 255.0).astype(np.uint8), (11, 11), 0)
|
||
_, binm = cv2.threshold(soft, 127, 255, cv2.THRESH_BINARY)
|
||
contours, _ = cv2.findContours(binm, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
||
if not contours:
|
||
return None
|
||
contour = max(contours, key=cv2.contourArea)
|
||
|
||
pts = contour.reshape(-1, 2).astype(np.float32)
|
||
if edge_img is not None and pts.shape[0] >= 8:
|
||
pts = self._snap_to_edges(track_id, pts, edge_img)
|
||
|
||
contour = pts.reshape(-1, 1, 2).astype(np.int32)
|
||
eps = 0.008 * cv2.arcLength(contour, True)
|
||
return cv2.approxPolyDP(contour, eps, True)
|
||
|
||
def _snap_to_edges(
|
||
self, track_id: int, pts: np.ndarray, edge: np.ndarray
|
||
) -> np.ndarray:
|
||
"""Локальный поиск сдвига (±14 px) и поджатия контура, при которых под
|
||
контуром максимум RGB-краёв. Найденная поправка сглаживается по времени."""
|
||
H, W = edge.shape[:2]
|
||
xs, ys = pts[:, 0], pts[:, 1]
|
||
|
||
def score(dx: float, dy: float, f: float, c: np.ndarray) -> float:
|
||
x = np.clip((c[0] + f * (xs - c[0]) + dx).astype(np.int32), 0, W - 1)
|
||
y = np.clip((c[1] + f * (ys - c[1]) + dy).astype(np.int32), 0, H - 1)
|
||
return float(edge[y, x].mean()) - 0.6 * float(np.hypot(dx, dy))
|
||
|
||
c = pts.mean(axis=0)
|
||
best_dx, best_dy, best_s = 0.0, 0.0, score(0, 0, 1.0, c)
|
||
for dy in range(-14, 15, 2):
|
||
for dx in range(-14, 15, 2):
|
||
s = score(dx, dy, 1.0, c)
|
||
if s > best_s:
|
||
best_s, best_dx, best_dy = s, float(dx), float(dy)
|
||
best_f = 1.0
|
||
# только сдвиг и лёгкое РАСШИРЕНИЕ — поджатие (f<1) отрезало часть объекта
|
||
for f in (1.0, 1.04, 1.08):
|
||
s = score(best_dx, best_dy, f, c)
|
||
if s > best_s:
|
||
best_s, best_f = s, f
|
||
|
||
prev = self._snap.get(track_id, (0.0, 0.0, 1.0))
|
||
a = self.snap_alpha
|
||
sm = (
|
||
(1 - a) * prev[0] + a * best_dx,
|
||
(1 - a) * prev[1] + a * best_dy,
|
||
max(1.0, (1 - a) * prev[2] + a * best_f),
|
||
)
|
||
self._snap[track_id] = sm
|
||
out = pts.copy()
|
||
out[:, 0] = c[0] + sm[2] * (xs - c[0]) + sm[0]
|
||
out[:, 1] = c[1] + sm[2] * (ys - c[1]) + sm[1]
|
||
return out
|
||
|
||
def drop_missing(self, alive_ids: set) -> None:
|
||
for tid in list(self._acc.keys()):
|
||
if tid not in alive_ids:
|
||
del self._acc[tid]
|
||
self._snap.pop(tid, None)
|
||
|
||
|
||
def align_color(color_bgr: np.ndarray, dx: float, dy: float, scale: float) -> np.ndarray:
|
||
"""Совмещение RGB с depth: сдвиг+масштаб (у D415 сенсоры разнесены)."""
|
||
if abs(dx) < 0.5 and abs(dy) < 0.5 and abs(scale - 1.0) < 1e-3:
|
||
return color_bgr
|
||
h, w = color_bgr.shape[:2]
|
||
M = np.float32([
|
||
[scale, 0, dx + (1.0 - scale) * w / 2.0],
|
||
[0, scale, dy + (1.0 - scale) * h / 2.0],
|
||
])
|
||
return cv2.warpAffine(color_bgr, M, (w, h), flags=cv2.INTER_LINEAR)
|
||
|
||
|
||
def build_demo_frame(
|
||
color_bgr: np.ndarray,
|
||
depth_mm: np.ndarray,
|
||
tracks: list, # List[tracker.Track]
|
||
belt_mm: float,
|
||
stats: Optional[Dict[str, int]] = None,
|
||
confidence_pct: int = 75,
|
||
rgb_available: bool = False,
|
||
background_active: bool = False,
|
||
color_align: Tuple[float, float, float] = (0.0, 0.0, 1.0),
|
||
contour_smoother: Optional[ContourSmoother] = None,
|
||
) -> np.ndarray:
|
||
from camera import depth_colormap
|
||
|
||
# RGB для веба; если цвет недоступен — colorize(depth), НЕ чёрный экран
|
||
if rgb_available and color_bgr is not None and float(np.std(color_bgr)) > 4.0:
|
||
base = color_bgr
|
||
if base.shape[:2] != depth_mm.shape[:2]:
|
||
base = cv2.resize(base, (depth_mm.shape[1], depth_mm.shape[0]))
|
||
base = align_color(base, color_align[0], color_align[1], color_align[2])
|
||
view = base.copy()
|
||
rgb_ok = True
|
||
else:
|
||
view = depth_colormap(depth_mm)
|
||
view = cv2.medianBlur(view, 3)
|
||
base = view
|
||
rgb_ok = False
|
||
h, w = view.shape[:2]
|
||
|
||
# мягкое поле RGB-краёв для «примагничивания» контуров
|
||
edge_field: Optional[np.ndarray] = None
|
||
if rgb_ok and contour_smoother is not None and tracks:
|
||
gray = cv2.cvtColor(base, cv2.COLOR_BGR2GRAY)
|
||
gray = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8)).apply(gray)
|
||
edge_field = cv2.Canny(cv2.GaussianBlur(gray, (5, 5), 0), 40, 120)
|
||
edge_field = cv2.GaussianBlur(edge_field, (13, 13), 0)
|
||
|
||
texts: List[Tuple[int, int, str, Tuple[int, int, int], int]] = []
|
||
max_conf_pending = 0
|
||
all_locked = bool(tracks)
|
||
alive_ids = set()
|
||
|
||
for tr in tracks:
|
||
m = tr.measurement
|
||
d = tr.decision
|
||
if m is None or d is None:
|
||
continue
|
||
alive_ids.add(tr.track_id)
|
||
locked = d.locked and d.result is not None
|
||
if locked:
|
||
color = UNCERTAIN_COLOR if d.uncertain else ZONE_COLOR.get(d.result.category.zone, PENDING_COLOR)
|
||
else:
|
||
color = PENDING_COLOR
|
||
all_locked = False
|
||
max_conf_pending = max(max_conf_pending, d.confidence_pct)
|
||
|
||
# bbox / класс — по сырому контуру измерения; сглаживание только для окантовки «анализ»
|
||
raw_contour = m.contour
|
||
draw_contour = raw_contour
|
||
if not locked and contour_smoother is not None:
|
||
sm = contour_smoother.smooth(tr.track_id, m.mask, edge_img=edge_field)
|
||
if sm is not None:
|
||
draw_contour = sm
|
||
if locked:
|
||
bx, by, bw, bh = cv2.boundingRect(raw_contour)
|
||
# небольшой запас, чтобы рамка не обрезала края
|
||
pad = 4
|
||
bx, by = max(0, bx - pad), max(0, by - pad)
|
||
bw = min(w - bx, bw + 2 * pad)
|
||
bh = min(h - by, bh + 2 * pad)
|
||
cv2.rectangle(view, (bx, by), (bx + bw, by + bh), color, 2, lineType=cv2.LINE_AA)
|
||
cl = max(8, min(bw, bh) // 5)
|
||
for px, py, sx, sy in ((bx, by, 1, 1), (bx + bw, by, -1, 1),
|
||
(bx, by + bh, 1, -1), (bx + bw, by + bh, -1, -1)):
|
||
cv2.line(view, (px, py), (px + sx * cl, py), color, 4, lineType=cv2.LINE_AA)
|
||
cv2.line(view, (px, py), (px, py + sy * cl), color, 4, lineType=cv2.LINE_AA)
|
||
contour = raw_contour
|
||
else:
|
||
cv2.drawContours(view, [draw_contour], -1, color, 2, lineType=cv2.LINE_AA)
|
||
contour = draw_contour
|
||
ccx, ccy = contour.reshape(-1, 2).mean(axis=0)
|
||
cv2.circle(view, (int(ccx), int(ccy)), 4, (0, 0, 255), -1, lineType=cv2.LINE_AA)
|
||
|
||
x0, y0, _, _ = cv2.boundingRect(contour)
|
||
tx = int(np.clip(x0, 4, w - 220))
|
||
ty = int(np.clip(y0 - 46, 4, h - 46))
|
||
if locked:
|
||
label = d.result.category.short_label
|
||
if d.uncertain:
|
||
label = "НЕУВЕРЕННО → " + label
|
||
texts.append((tx, ty, f"#{tr.track_id} {label}", color, 20))
|
||
else:
|
||
texts.append((tx, ty, f"#{tr.track_id} анализ… {d.confidence_pct}%", color, 20))
|
||
dims = d.result.dims_sorted_mm if (locked and d.result) else (m.length_mm, m.width_mm, m.height_mm)
|
||
ratio = d.result.circle_ratio if (locked and d.result) else m.circle_ratio
|
||
src = " · RGB" if getattr(m, "source", "depth") == "rgb" else ""
|
||
texts.append(
|
||
(tx, ty + 24, f"{dims[0]:.0f}×{dims[1]:.0f}×{dims[2]:.0f} мм · круг {ratio:.2f}{src}", (235, 235, 235), 15)
|
||
)
|
||
|
||
if contour_smoother is not None:
|
||
contour_smoother.drop_missing(alive_ids)
|
||
|
||
view = _draw_texts(view, texts)
|
||
|
||
# надписи — в отдельной полосе НАД кадром, чтобы не закрывать камеру
|
||
top_bar = np.full((36, w, 3), 18, np.uint8)
|
||
st = stats or {}
|
||
stats_line = (
|
||
f"ГОТОВ {st.get('B', 0)} · НЕГАБАРИТ {st.get('C', 0)} · ДОУПАК {st.get('D', 0)}"
|
||
+ (f" · неувер. {st['uncertain']}" if st.get("uncertain") else "")
|
||
)
|
||
bg_tag = " · фон:карта" if background_active else ""
|
||
left = "объектов нет" if not tracks else f"объектов: {len(alive_ids)}"
|
||
top_texts = [
|
||
(10, 6, left, (150, 150, 255) if not tracks else (200, 230, 200), 17),
|
||
(max(200, w - 440), 8, f"{stats_line} | h={belt_mm:.0f}мм{bg_tag}", (200, 230, 200), 14),
|
||
]
|
||
top_bar = _draw_texts(top_bar, top_texts)
|
||
view = np.vstack([top_bar, view])
|
||
h = view.shape[0]
|
||
|
||
# прогресс уверенности внизу
|
||
bar_y = h - 8
|
||
cv2.rectangle(view, (0, bar_y), (w, h), (40, 40, 40), -1)
|
||
conf_show = 100 if (all_locked and tracks) else max_conf_pending
|
||
fill = int(w * min(1.0, conf_show / 100.0))
|
||
col = (40, 200, 40) if (all_locked and tracks) else (0, 200, 255)
|
||
cv2.rectangle(view, (0, bar_y), (fill, h), col, -1)
|
||
thr = int(w * confidence_pct / 100.0)
|
||
cv2.line(view, (thr, bar_y), (thr, h), (255, 255, 255), 1)
|
||
|
||
return view
|