feat(edge): Track model + IoU tracker with lost-track reporting
This commit is contained in:
parent
991178a896
commit
02ca524bf6
3 changed files with 124 additions and 0 deletions
12
edge/ports.py
Normal file
12
edge/ports.py
Normal file
|
|
@ -0,0 +1,12 @@
|
|||
from typing import Protocol
|
||||
|
||||
|
||||
class SourcePort(Protocol):
|
||||
def read(self) -> tuple[bool, object]:
|
||||
"""(ok, frame). ok=False면 프레임 없음."""
|
||||
...
|
||||
|
||||
|
||||
class PosePort(Protocol):
|
||||
def infer(self, frame) -> list: # list[Person]
|
||||
...
|
||||
74
edge/tracking.py
Normal file
74
edge/tracking.py
Normal file
|
|
@ -0,0 +1,74 @@
|
|||
from collections import deque
|
||||
from dataclasses import dataclass
|
||||
|
||||
from edge.models import Person, BBox, Keypoint
|
||||
|
||||
|
||||
@dataclass
|
||||
class TrackFrame:
|
||||
t: float
|
||||
bbox: BBox
|
||||
keypoints: list # list[Keypoint]
|
||||
|
||||
|
||||
@dataclass
|
||||
class Track:
|
||||
track_id: int
|
||||
frames: deque
|
||||
|
||||
@property
|
||||
def latest(self) -> TrackFrame:
|
||||
return self.frames[-1]
|
||||
|
||||
|
||||
def _iou(a: BBox, b: BBox) -> float:
|
||||
ix1, iy1 = max(a.x1, b.x1), max(a.y1, b.y1)
|
||||
ix2, iy2 = min(a.x2, b.x2), min(a.y2, b.y2)
|
||||
iw, ih = max(0.0, ix2 - ix1), max(0.0, iy2 - iy1)
|
||||
inter = iw * ih
|
||||
union = a.width * a.height + b.width * b.height - inter
|
||||
return inter / union if union else 0.0
|
||||
|
||||
|
||||
class IouTracker:
|
||||
"""경량 IoU 그리디 매칭 추적기. 실장비는 ByteTrack 어댑터로 교체 가능(동일 TrackerPort)."""
|
||||
|
||||
def __init__(self, iou_thr: float = 0.3, max_age: float = 1.0, history: int = 90):
|
||||
self.iou_thr = iou_thr
|
||||
self.max_age = max_age
|
||||
self.history = history
|
||||
self._tracks: dict[int, Track] = {}
|
||||
self._last_seen: dict[int, float] = {}
|
||||
self._lost: list[Track] = []
|
||||
self._next_id = 1
|
||||
|
||||
def update(self, persons: list[Person], t: float) -> list[Track]:
|
||||
self._lost = []
|
||||
unmatched = list(persons)
|
||||
# 기존 트랙과 그리디 IoU 매칭
|
||||
for tid, track in list(self._tracks.items()):
|
||||
best, best_iou = None, self.iou_thr
|
||||
for p in unmatched:
|
||||
iou = _iou(track.latest.bbox, p.bbox)
|
||||
if iou >= best_iou:
|
||||
best, best_iou = p, iou
|
||||
if best is not None:
|
||||
unmatched.remove(best)
|
||||
track.frames.append(TrackFrame(t, best.bbox, best.keypoints))
|
||||
self._last_seen[tid] = t
|
||||
# 신규 트랙
|
||||
for p in unmatched:
|
||||
tid = self._next_id
|
||||
self._next_id += 1
|
||||
self._tracks[tid] = Track(tid, deque([TrackFrame(t, p.bbox, p.keypoints)],
|
||||
maxlen=self.history))
|
||||
self._last_seen[tid] = t
|
||||
# 소실 판정
|
||||
for tid in list(self._tracks.keys()):
|
||||
if t - self._last_seen[tid] > self.max_age:
|
||||
self._lost.append(self._tracks.pop(tid))
|
||||
self._last_seen.pop(tid, None)
|
||||
return list(self._tracks.values())
|
||||
|
||||
def lost(self) -> list[Track]:
|
||||
return list(self._lost)
|
||||
38
tests/test_tracking.py
Normal file
38
tests/test_tracking.py
Normal file
|
|
@ -0,0 +1,38 @@
|
|||
from edge.models import Person, BBox, Keypoint
|
||||
from edge.tracking import IouTracker
|
||||
|
||||
|
||||
def _person(x1, y1, x2, y2):
|
||||
kps = [Keypoint((x1 + x2) / 2, (y1 + y2) / 2, 0.9) for _ in range(17)]
|
||||
return Person(BBox(x1, y1, x2, y2), kps)
|
||||
|
||||
|
||||
def test_new_person_gets_track_id():
|
||||
tr = IouTracker()
|
||||
tracks = tr.update([_person(0, 0, 10, 20)], t=0.0)
|
||||
assert len(tracks) == 1
|
||||
assert tracks[0].track_id == 1
|
||||
|
||||
|
||||
def test_overlapping_person_keeps_same_id():
|
||||
tr = IouTracker(iou_thr=0.3)
|
||||
tr.update([_person(0, 0, 10, 20)], t=0.0)
|
||||
tracks = tr.update([_person(1, 1, 11, 21)], t=0.1) # high overlap
|
||||
assert tracks[0].track_id == 1
|
||||
|
||||
|
||||
def test_disappeared_track_is_reported_lost():
|
||||
tr = IouTracker(max_age=0.2)
|
||||
tr.update([_person(0, 0, 10, 20)], t=0.0)
|
||||
tr.update([], t=0.5) # nothing detected, past max_age
|
||||
lost = tr.lost()
|
||||
assert len(lost) == 1
|
||||
assert lost[0].track_id == 1
|
||||
|
||||
|
||||
def test_history_accumulates_frames():
|
||||
tr = IouTracker()
|
||||
tr.update([_person(0, 0, 10, 20)], t=0.0)
|
||||
tracks = tr.update([_person(0, 0, 10, 20)], t=0.1)
|
||||
assert len(tracks[0].frames) == 2
|
||||
assert tracks[0].latest.t == 0.1
|
||||
Loading…
Reference in a new issue