test_sort.py

import numpy as np

from sort import TrackerManager, iou_calculation


def test_iou_identical_boxes():
    assert (
        iou_calculation(
            [0, 0, 10, 10],
            [0, 0, 10, 10],
        )
        == 1.0
    )


def test_iou_partial_overlap():
    result = iou_calculation(
        [0, 0, 10, 10],
        [5, 5, 15, 15],
    )

    assert abs(result - 0.14285714285714285) < 1e-9


def test_iou_no_overlap():
    assert (
        iou_calculation(
            [0, 0, 10, 10],
            [20, 20, 30, 30],
        )
        == 0.0
    )


def test_single_object_keeps_same_id():
    tracker = TrackerManager(
        max_age=1,
        min_hits=1,
        iou_threshold=0.3,
    )

    frame_1 = np.array([[10, 10, 20, 20]])
    frame_2 = np.array([[12, 12, 22, 22]])

    result_1 = tracker.update(frame_1)

    assert len(result_1) == 0
    assert len(tracker.tracks) == 1

    track_id_after_creation = tracker.tracks[0].id

    result_2 = tracker.update(frame_2)

    assert len(result_2) == 1
    assert len(tracker.tracks) == 1

    track_id_after_update = tracker.tracks[0].id

    assert track_id_after_update == track_id_after_creation
    assert result_2[0, 4] == track_id_after_creation


def test_two_objects_keep_distinct_ids():
    tracker = TrackerManager(
        max_age=1,
        min_hits=1,
        iou_threshold=0.3,
    )

    frame_1 = np.array(
        [
            [10, 10, 20, 20],
            [50, 50, 60, 60],
        ]
    )

    frame_2 = np.array(
        [
            [12, 12, 22, 22],
            [48, 48, 58, 58],
        ]
    )

    result_1 = tracker.update(frame_1)
    result_2 = tracker.update(frame_2)

    assert len(result_1) == 0
    assert len(result_2) == 2

    ids = result_2[:, 4]

    assert len(set(ids)) == 2


def test_track_is_confirmed_after_min_hits():
    tracker = TrackerManager(
        max_age=1,
        min_hits=3,
        iou_threshold=0.3,
    )

    detection = np.array([[10, 10, 20, 20]])

    result_1 = tracker.update(detection)
    result_2 = tracker.update(detection)
    result_3 = tracker.update(detection)
    result_4 = tracker.update(detection)

    assert len(result_1) == 0
    assert len(result_2) == 0
    assert len(result_3) == 0
    assert len(result_4) == 1


def test_track_survives_temporary_missed_detection():
    tracker = TrackerManager(
        max_age=2,
        min_hits=1,
        iou_threshold=0.3,
    )

    frame_1 = np.array([[10, 10, 20, 20]])
    frame_2 = np.array([[12, 12, 22, 22]])
    frame_3 = np.empty((0, 4))

    result_1 = tracker.update(frame_1)
    result_2 = tracker.update(frame_2)
    result_3 = tracker.update(frame_3)

    assert len(result_1) == 0
    assert len(result_2) == 1

    assert len(tracker.tracks) == 1
    assert tracker.tracks[0].time_since_update == 1


def test_track_is_deleted_after_max_age():
    tracker = TrackerManager(
        max_age=2,
        min_hits=1,
        iou_threshold=0.3,
    )

    detected = np.array([[10, 10, 20, 20]])
    no_detections = np.empty((0, 4))

    tracker.update(detected)
    tracker.update(detected)

    assert len(tracker.tracks) == 1

    tracker.update(no_detections)
    assert len(tracker.tracks) == 1
    assert tracker.tracks[0].time_since_update == 1

    tracker.update(no_detections)
    assert len(tracker.tracks) == 1
    assert tracker.tracks[0].time_since_update == 2

    tracker.update(no_detections)
    assert len(tracker.tracks) == 0


def test_new_detection_creates_new_track():
    tracker = TrackerManager(
        max_age=1,
        min_hits=1,
        iou_threshold=0.3,
    )

    first_detection = np.array([[10, 10, 20, 20]])
    second_detection = np.array([[50, 50, 60, 60]])

    tracker.update(first_detection)
    tracker.update(first_detection)

    assert len(tracker.tracks) == 1

    tracker.update(
        np.array(
            [
                [12, 12, 22, 22],
                second_detection[0],
            ]
        )
    )

    assert len(tracker.tracks) == 2

    ids = {track.id for track in tracker.tracks}

    assert len(ids) == 2


def test_low_iou_detection_does_not_update_existing_track():
    tracker = TrackerManager(
        max_age=2,
        min_hits=1,
        iou_threshold=0.3,
    )

    first_detection = np.array([[10, 10, 20, 20]])
    far_detection = np.array([[100, 100, 110, 110]])

    tracker.update(first_detection)
    tracker.update(first_detection)

    assert len(tracker.tracks) == 1

    original_id = tracker.tracks[0].id
    original_hits = tracker.tracks[0].hits

    tracker.update(far_detection)

    assert len(tracker.tracks) == 2

    original_track = next(track for track in tracker.tracks if track.id == original_id)

    new_track = next(track for track in tracker.tracks if track.id != original_id)

    assert original_track.hits == original_hits
    assert original_track.time_since_update == 1
    assert new_track.hits == 0


def test_empty_detections_on_new_tracker():
    tracker = TrackerManager(
        max_age=1,
        min_hits=1,
        iou_threshold=0.3,
    )

    result = tracker.update(np.empty((0, 4)))

    assert result.shape == (0,)
    assert len(tracker.tracks) == 0