def test_sort_ellipse():
tracklets = dict()
mot = trackingutils.SORTEllipse(1, 1, 0.6)
poses = np.random.rand(2, 10, 3)
trackers = mot.track(poses[..., :2])
assert trackers.shape == (2, 7)
trackingutils.fill_tracklets(tracklets, trackers, poses, imname=0)
assert all(id_ in tracklets for id_ in trackers[:, -2])
def test_tracking(real_assemblies, real_tracklets):
tracklets_ref = real_tracklets.copy()
_ = tracklets_ref.pop("header", None)
tracklets = dict()
mot_tracker = trackingutils.SORTEllipse(1, 1, 0.6)
for ind, assemblies in real_assemblies.items():
animals = np.stack([ass.data[:, :3] for ass in assemblies])
trackers = mot_tracker.track(animals[..., :2])
trackingutils.fill_tracklets(tracklets, trackers, animals, ind)
assert len(tracklets) == len(tracklets_ref)
assert [len(tracklet) for tracklet in tracklets.values()
] == [len(tracklet) for tracklet in tracklets_ref.values()]
def test_tracking_box(real_assemblies, real_tracklets):
tracklets_ref = real_tracklets.copy()
_ = tracklets_ref.pop("header", None)
tracklets = dict()
mot_tracker = trackingutils.SORTBox(1, 1, 0.1)
for ind, assemblies in real_assemblies.items():
animals = np.stack([ass.data for ass in assemblies])
bboxes = trackingutils.calc_bboxes_from_keypoints(animals)
trackers = mot_tracker.track(bboxes)
trackingutils.fill_tracklets(tracklets, trackers, animals, ind)
assert len(tracklets) == len(tracklets_ref)
assert [len(tracklet) for tracklet in tracklets.values()
] == [len(tracklet) for tracklet in tracklets_ref.values()]
assert all(t.shape[1] == 4 for tracklet in tracklets.values()
for t in tracklet.values())
def test_tracking_montblanc(
real_assemblies_montblanc,
real_tracklets_montblanc,
):
tracklets_ref = real_tracklets_montblanc.copy()
_ = tracklets_ref.pop("header", None)
tracklets = dict()
tracklets["single"] = real_assemblies_montblanc[1]
mot_tracker = trackingutils.SORTEllipse(1, 1, 0.6)
for ind, assemblies in real_assemblies_montblanc[0].items():
animals = np.stack([ass.data for ass in assemblies])
trackers = mot_tracker.track(animals[..., :2])
trackingutils.fill_tracklets(tracklets, trackers, animals, ind)
assert len(tracklets) == len(tracklets_ref)
assert [len(tracklet) for tracklet in tracklets.values()
] == [len(tracklet) for tracklet in tracklets_ref.values()]
for k, assemblies in tracklets.items():
ref = tracklets_ref[k]
for ind, data in assemblies.items():
frame = f"frame{str(ind).zfill(3)}" if k != "single" else ind
np.testing.assert_equal(data, ref[frame])
