38 lines
1.1 KiB
Python
38 lines
1.1 KiB
Python
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
|