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>
280 lines
12 KiB
Python
280 lines
12 KiB
Python
"""Трекинг нескольких объектов в кадре: ID + статистика зон за сессию.
|
||
|
||
Сопоставление кадр-к-кадру по ближайшему центроиду. На каждый трек —
|
||
свой DecisionStabilizer. LOCK в ленту/MQTT — один раз на товар: после
|
||
фиксации запоминаем «отпечаток» (центр+габариты); если depth кратковременно
|
||
пропал и объект нашёлся снова рядом — трек возрождается без нового LOCK.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from dataclasses import dataclass
|
||
from typing import Callable, Dict, List, Optional, Sequence, Tuple
|
||
|
||
from classify import ClassificationResult
|
||
from measure import ObjectMeasurement
|
||
from stabilize import DecisionStabilizer, StableDecision
|
||
|
||
|
||
@dataclass
|
||
class Track:
|
||
track_id: int
|
||
stabilizer: DecisionStabilizer
|
||
centroid: Tuple[int, int]
|
||
miss: int = 0
|
||
counted_zone: Optional[str] = None
|
||
counted_uncertain: bool = False
|
||
measurement: Optional[ObjectMeasurement] = None
|
||
decision: Optional[StableDecision] = None
|
||
reported: bool = False # уже ушёл в ленту — повторно не пишем
|
||
frozen: Optional[StableDecision] = None # снимок первого LOCK (для возрождения)
|
||
|
||
|
||
@dataclass
|
||
class LockEvent:
|
||
track_id: int
|
||
decision: StableDecision
|
||
|
||
|
||
@dataclass
|
||
class _Fingerprint:
|
||
"""Товар уже зафиксирован — не считаем повторно при перезахвате."""
|
||
track_id: int
|
||
cx: int
|
||
cy: int
|
||
length_mm: float
|
||
width_mm: float
|
||
height_mm: float
|
||
zone: str
|
||
uncertain: bool
|
||
result: ClassificationResult
|
||
age: int = 0 # кадров без детекции
|
||
|
||
|
||
def slot_key(cx: int, cy: int, L: float, W: float, H: float, zone: str) -> str:
|
||
"""Стабильный ключ товара на ленте — один физический предмет = одна запись."""
|
||
a = sorted((L, W, H))
|
||
return f"{cx // 35}_{cy // 35}_{int(a[0] // 12)}_{int(a[1] // 12)}_{int(a[2] // 8)}_{zone}"
|
||
|
||
|
||
class MultiObjectTracker:
|
||
def __init__(
|
||
self,
|
||
stabilizer_factory: Callable[[], DecisionStabilizer],
|
||
max_dist_px: int = 120,
|
||
lost_frames: int = 12,
|
||
fingerprint_ttl: int = 450, # ~15–20 с помнить «уже учтён»
|
||
) -> None:
|
||
self._factory = stabilizer_factory
|
||
self.max_dist_px = int(max_dist_px)
|
||
self.lost_frames = int(lost_frames)
|
||
self.fingerprint_ttl = int(fingerprint_ttl)
|
||
self._tracks: Dict[int, Track] = {}
|
||
self._fps: List[_Fingerprint] = []
|
||
self._seen_slots: Dict[str, _Fingerprint] = {} # уже учтённые товары (позиция+габариты)
|
||
self._next_id = 1
|
||
self.stats: Dict[str, int] = {"B": 0, "C": 0, "D": 0, "total": 0, "uncertain": 0}
|
||
|
||
def reset(self) -> None:
|
||
self._tracks.clear()
|
||
self._fps.clear()
|
||
|
||
def update(
|
||
self,
|
||
measurements: List[ObjectMeasurement],
|
||
min_mm: Sequence[float] = (10, 10, 10),
|
||
max_mm: Sequence[float] = (450, 320, 320),
|
||
) -> Tuple[List[Track], List[LockEvent]]:
|
||
free_meas = list(range(len(measurements)))
|
||
assigned: Dict[int, int] = {}
|
||
pairs = []
|
||
for tid, tr in self._tracks.items():
|
||
for mi in free_meas:
|
||
m = measurements[mi]
|
||
d2 = (tr.centroid[0] - m.centroid_px[0]) ** 2 + (tr.centroid[1] - m.centroid_px[1]) ** 2
|
||
pairs.append((d2, tid, mi))
|
||
for d2, tid, mi in sorted(pairs):
|
||
if tid in assigned or mi not in free_meas:
|
||
continue
|
||
if d2 > self.max_dist_px**2:
|
||
continue
|
||
assigned[tid] = mi
|
||
free_meas.remove(mi)
|
||
|
||
events: List[LockEvent] = []
|
||
|
||
for tid in list(self._tracks.keys()):
|
||
tr = self._tracks[tid]
|
||
# залоченный трек держим дольше — глянец даёт короткие выпадения depth
|
||
kill_after = self.lost_frames * 3 if tr.counted_zone is not None else self.lost_frames
|
||
if tid in assigned:
|
||
m = measurements[assigned[tid]]
|
||
tr.centroid = m.centroid_px
|
||
tr.miss = 0
|
||
tr.measurement = m
|
||
tr.decision = tr.stabilizer.update(m, min_mm=min_mm, max_mm=max_mm)
|
||
# после возрождения стабилизатор ещё «холодный» — держим прошлый LOCK на экране
|
||
if tr.reported and tr.frozen is not None and not (tr.decision and tr.decision.locked):
|
||
tr.decision = tr.frozen
|
||
self._account(tr, events)
|
||
if tr.reported:
|
||
self._touch_fp(tr)
|
||
else:
|
||
tr.miss += 1
|
||
tr.measurement = None
|
||
tr.decision = tr.stabilizer.update(None, min_mm=min_mm, max_mm=max_mm)
|
||
if tr.reported and tr.frozen is not None and not (tr.decision and tr.decision.locked):
|
||
tr.decision = tr.frozen
|
||
if tr.miss >= kill_after:
|
||
if tr.reported and tr.frozen is not None and tr.frozen.result is not None:
|
||
self._remember(tr)
|
||
del self._tracks[tid]
|
||
|
||
# старение отпечатков
|
||
for fp in self._fps:
|
||
fp.age += 1
|
||
self._fps = [fp for fp in self._fps if fp.age < self.fingerprint_ttl]
|
||
|
||
for mi in free_meas:
|
||
m = measurements[mi]
|
||
fp = self._match_fp(m)
|
||
if fp is not None:
|
||
# тот же товар вернулся после выпадения depth — без нового LOCK
|
||
fp.age = 0
|
||
fp.cx, fp.cy = m.centroid_px
|
||
frozen = StableDecision(
|
||
fp.result, True, 100, fp.result.circle_ratio, True, uncertain=fp.uncertain,
|
||
)
|
||
tr = Track(
|
||
track_id=fp.track_id,
|
||
stabilizer=self._factory(),
|
||
centroid=m.centroid_px,
|
||
measurement=m,
|
||
counted_zone=fp.zone,
|
||
counted_uncertain=fp.uncertain,
|
||
reported=True,
|
||
frozen=frozen,
|
||
decision=frozen,
|
||
)
|
||
self._tracks[tr.track_id] = tr
|
||
self._fps = [x for x in self._fps if x.track_id != fp.track_id]
|
||
continue
|
||
|
||
tr = Track(
|
||
track_id=self._next_id,
|
||
stabilizer=self._factory(),
|
||
centroid=m.centroid_px,
|
||
measurement=m,
|
||
)
|
||
self._next_id += 1
|
||
tr.decision = tr.stabilizer.update(m, min_mm=min_mm, max_mm=max_mm)
|
||
self._tracks[tr.track_id] = tr
|
||
self._account(tr, events)
|
||
|
||
alive = sorted(self._tracks.values(), key=lambda t: t.track_id)
|
||
return [t for t in alive if t.measurement is not None or t.miss < (
|
||
self.lost_frames * 3 if t.counted_zone else self.lost_frames
|
||
)], events
|
||
|
||
def _account(self, tr: Track, events: List[LockEvent]) -> None:
|
||
d = tr.decision
|
||
if d is None or not d.locked or d.result is None:
|
||
return
|
||
zone = d.result.category.zone
|
||
L, W, H = d.result.dims_sorted_mm
|
||
sk = slot_key(tr.centroid[0], tr.centroid[1], L, W, H, zone)
|
||
if tr.counted_zone is None:
|
||
if sk in self._seen_slots:
|
||
# тот же товар уже был в ленте/статистике — только показываем на экране
|
||
prev = self._seen_slots[sk]
|
||
tr.counted_zone = prev.zone
|
||
tr.counted_uncertain = prev.uncertain
|
||
tr.frozen = StableDecision(
|
||
prev.result, True, 100, prev.result.circle_ratio, True, uncertain=prev.uncertain,
|
||
)
|
||
tr.reported = True
|
||
tr.decision = tr.frozen
|
||
return
|
||
self.stats[zone] = self.stats.get(zone, 0) + 1
|
||
self.stats["total"] += 1
|
||
if d.uncertain:
|
||
self.stats["uncertain"] += 1
|
||
tr.counted_zone = zone
|
||
tr.counted_uncertain = d.uncertain
|
||
tr.frozen = StableDecision(
|
||
d.result, True, 100, d.result.circle_ratio, True, uncertain=d.uncertain,
|
||
)
|
||
self._seen_slots[sk] = _Fingerprint(
|
||
track_id=tr.track_id,
|
||
cx=tr.centroid[0],
|
||
cy=tr.centroid[1],
|
||
length_mm=float(L),
|
||
width_mm=float(W),
|
||
height_mm=float(H),
|
||
zone=zone,
|
||
uncertain=d.uncertain,
|
||
result=d.result,
|
||
)
|
||
if not tr.reported:
|
||
tr.reported = True
|
||
events.append(LockEvent(tr.track_id, tr.frozen))
|
||
elif tr.counted_zone != zone:
|
||
# смена зоны на экране/в счётчиках — в ленту повторно не пишем
|
||
self.stats[tr.counted_zone] = max(0, self.stats.get(tr.counted_zone, 0) - 1)
|
||
self.stats[zone] = self.stats.get(zone, 0) + 1
|
||
if tr.counted_uncertain and not d.uncertain:
|
||
self.stats["uncertain"] = max(0, self.stats["uncertain"] - 1)
|
||
tr.counted_uncertain = d.uncertain
|
||
tr.counted_zone = zone
|
||
tr.frozen = StableDecision(
|
||
d.result, True, 100, d.result.circle_ratio, True, uncertain=d.uncertain,
|
||
)
|
||
|
||
def _remember(self, tr: Track) -> None:
|
||
assert tr.frozen is not None and tr.frozen.result is not None
|
||
r = tr.frozen.result
|
||
L, W, H = r.dims_sorted_mm
|
||
self._fps = [fp for fp in self._fps if fp.track_id != tr.track_id]
|
||
self._fps.append(
|
||
_Fingerprint(
|
||
track_id=tr.track_id,
|
||
cx=tr.centroid[0],
|
||
cy=tr.centroid[1],
|
||
length_mm=float(L),
|
||
width_mm=float(W),
|
||
height_mm=float(H),
|
||
zone=tr.counted_zone or r.category.zone,
|
||
uncertain=tr.counted_uncertain,
|
||
result=r,
|
||
age=0,
|
||
)
|
||
)
|
||
|
||
def _touch_fp(self, tr: Track) -> None:
|
||
for fp in self._fps:
|
||
if fp.track_id == tr.track_id:
|
||
fp.age = 0
|
||
fp.cx, fp.cy = tr.centroid
|
||
|
||
def _match_fp(self, m: ObjectMeasurement) -> Optional[_Fingerprint]:
|
||
best: Optional[_Fingerprint] = None
|
||
best_d2 = self.max_dist_px**2
|
||
for fp in self._fps:
|
||
d2 = (fp.cx - m.centroid_px[0]) ** 2 + (fp.cy - m.centroid_px[1]) ** 2
|
||
if d2 > best_d2:
|
||
continue
|
||
if not _dims_close(fp.length_mm, fp.width_mm, fp.height_mm, m.length_mm, m.width_mm, m.height_mm):
|
||
continue
|
||
best, best_d2 = fp, d2
|
||
return best
|
||
|
||
|
||
def _dims_close(L0: float, W0: float, H0: float, L1: float, W1: float, H1: float, tol: float = 0.40) -> bool:
|
||
"""Габариты «похожи» (порядок осей уже отсортирован в classify, здесь — сырые L×W×H)."""
|
||
a = sorted((L0, W0, H0))
|
||
b = sorted((L1, W1, H1))
|
||
for x, y in zip(a, b):
|
||
if abs(x - y) / max(x, y, 1.0) > tol:
|
||
return False
|
||
return True
|