sort.py

from scipy.optimize import linear_sum_assignment
from filterpy.kalman import KalmanFilter
import numpy as np


# --------------- IOU FUNCTION ------------------------------------------------


def iou_calculation(boxA, boxB):  # Calculate bounding-box overlap.

    # Calculate top-left and bottom-right corners of the intersection.
    top_left = (max(boxA[0], boxB[0]), max(boxA[1], boxB[1]))
    bottom_right = (min(boxA[2], boxB[2]), min(boxA[3], boxB[3]))

    # Calculate intersection area.
    intersection_width = max(
        0, bottom_right[0] - top_left[0]
    )  # Clamp negative overlap to zero.
    intersection_height = max(
        0, bottom_right[1] - top_left[1]
    )  # Clamp negative overlap to zero.

    intersection_area = intersection_width * intersection_height

    # Calculate union area:
    # union = area(box1) + area(box2) - intersection area
    # area = width * height
    boxA_area = (boxA[2] - boxA[0]) * (boxA[3] - boxA[1])
    boxB_area = (boxB[2] - boxB[0]) * (boxB[3] - boxB[1])

    union_area = boxA_area + boxB_area - intersection_area

    # Calculate IoU.
    iou = intersection_area / union_area if union_area > 0 else 0

    return iou


# -------------- TRACKER CLASS ------------------------------------------------


class Tracker:  # Handles the prediction and update phases of tracking.
    count = 0  # Global ID counter.

    def __init__(self, bbox):
        # Initialize tracker from the first detection.
        # bbox format: [x1, y1, x2, y2]

        # Create a Kalman filter.
        # 7 state variables: [u, v, s, r, u', v', s']
        # 4 measurements: [u, v, s, r]
        self.kf = KalmanFilter(dim_x=7, dim_z=4)

        # Convert bbox from [x1, y1, x2, y2] to [u, v, s, r].
        x1, y1, x2, y2 = bbox
        w = x2 - x1  # Width.
        h = y2 - y1  # Height.

        u = x1 + w / 2.0  # Center x.
        v = y1 + h / 2.0  # Center y.
        s = w * h  # Area.
        r = w / float(h)  # Aspect ratio.

        # Initialize state with zero velocity.
        self.kf.x = np.array([[u], [v], [s], [r], [0], [0], [0]])

        # Configure the state transition matrix.
        # F describes how the state changes between time steps.
        self.kf.F = np.array(
            [
                [1, 0, 0, 0, 1, 0, 0],  # u = u + u_velocity
                [0, 1, 0, 0, 0, 1, 0],  # v = v + v_velocity
                [0, 0, 1, 0, 0, 0, 1],  # s = s + s_velocity
                [0, 0, 0, 1, 0, 0, 0],  # r = r
                [0, 0, 0, 0, 1, 0, 0],  # u_velocity
                [0, 0, 0, 0, 0, 1, 0],  # v_velocity
                [0, 0, 0, 0, 0, 0, 1],  # s_velocity
            ],
            dtype=float,
        )

        # Configure the measurement matrix.
        # Only [u, v, s, r] are directly measured.
        self.kf.H = np.array(  # type: ignore
            [
                [1, 0, 0, 0, 0, 0, 0],  # Measure u.
                [0, 1, 0, 0, 0, 0, 0],  # Measure v.
                [0, 0, 1, 0, 0, 0, 0],  # Measure s.
                [0, 0, 0, 1, 0, 0, 0],  # Measure r.
            ],
            dtype=float,
        )

        # Measurement noise: copied from the SORT paper.
        self.kf.R[2:, 2:] *= 10.0  # s and r are noisier.

        # Initial uncertainty: velocities are not known.
        self.kf.P[4:, 4:] *= 1000.0  # High uncertainty in velocities.
        self.kf.P *= 10.0

        # Process noise: model uncertainty in object motion.
        self.kf.Q[-1, -1] *= 0.01  # Aspect ratio changes slowly.
        self.kf.Q[4:, 4:] *= 0.01  # Velocities are relatively stable.

        # Tracking metadata.
        self.id = Tracker.count  # Assign a unique ID.
        Tracker.count += 1  # Increment the global ID counter.

        self.time_since_update = 0  # Number of frames since the last successful update.
        self.hits = 0  # Number of successful detections.
        self.age = 0  # Number of frames since the tracker was created.
        self.hit_streak = 0  # Number of consecutive successful detections.

    def predict(self):
        self.kf.predict()
        self.age += 1

        # Reset the hit streak if the object was missed in a previous frame.
        if self.time_since_update > 0:
            self.hit_streak = 0

        self.time_since_update += 1

        return self.get_state()

    def get_state(self):
        """
        Convert the Kalman filter state representation
        [u, v, s, r] into [x1, y1, x2, y2].
        """

        u, v, s, r = self.kf.x[:4].flatten()

        # Clamp invalid area values to a small positive number.
        if s <= 0:
            s = 1e-4

            # Update the Kalman state to maintain numerical stability.
            self.kf.x[2, 0] = s

        # Convert area and aspect ratio back to width and height.
        w = np.sqrt(s * r)  # Width from area and aspect ratio.
        h = s / w  # Height from area and width.

        # Calculate bounding-box corners from the center point.
        x1 = u - w / 2
        y1 = v - h / 2
        x2 = u + w / 2
        y2 = v + h / 2

        return np.array([x1, y1, x2, y2])

    def update(self, bbox):
        """
        Update the tracker with a new detection.

        The detection is converted from [x1, y1, x2, y2]
        to the Kalman filter measurement format [u, v, s, r].
        """

        # Convert bbox to [u, v, s, r].
        x1, y1, x2, y2 = bbox
        w = x2 - x1  # Width.
        h = y2 - y1  # Height.

        u = x1 + w / 2.0  # Center x.
        v = y1 + h / 2.0  # Center y.
        s = w * h  # Area.
        r = w / float(h)  # Aspect ratio.

        # Update the Kalman filter with the new measurement.
        z = np.array([[u], [v], [s], [r]])
        self.kf.update(z)

        # Update tracking statistics.
        self.hits += 1
        self.hit_streak += 1
        self.time_since_update = 0


# -------------- TRACKER MANAGER CLASS ----------------------------------------


class TrackerManager:
    def __init__(self, max_age=1, min_hits=3, iou_threshold=0.3):
        # Initialize tracker manager.
        self.tracks = []  # List of active tracks.
        self.max_age = max_age
        self.min_hits = min_hits
        self.iou_threshold = iou_threshold

    def update(self, detections):
        """
        Update all tracks for the current frame.

        Steps:
        1. Predict all existing tracks.
        2. Compute the IoU matrix.
        3. Perform data association using the Hungarian algorithm.
        4. Update matched tracks.
        5. Handle unmatched tracks.
        6. Create new tracks for unmatched detections.
        7. Delete expired tracks.
        8. Output confirmed tracks.
        """

        # 1. Predict all existing tracks.
        predictions = []

        for t in self.tracks:
            t.predict()
            predictions.append(t.get_state())

        # 2. Compute the IoU matrix.
        iou_matrix = np.zeros((len(predictions), len(detections)))

        for i, pred in enumerate(predictions):
            for j, det in enumerate(detections):
                iou_matrix[i, j] = iou_calculation(pred, det)

        # Convert IoU scores to costs because the Hungarian algorithm
        # minimizes cost.
        cost_matrix = 1 - iou_matrix

        # 3. Perform data association using the Hungarian algorithm.
        row_idx, col_idx = linear_sum_assignment(cost_matrix)

        # Hungarian assignment can produce weak matches,
        # so matches below the IoU threshold are rejected.
        matches = []

        unmatched_tracks = set(range(len(self.tracks)))
        unmatched_detections = set(range(len(detections)))

        for r, c in zip(row_idx, col_idx):
            if iou_matrix[r, c] >= self.iou_threshold:
                matches.append((r, c))
                unmatched_tracks.remove(r)
                unmatched_detections.remove(c)

        # 4. Update matched tracks.
        for track_idx, det_idx in matches:
            self.tracks[track_idx].update(detections[det_idx])

        # 5. Handle unmatched tracks.
        # No update is applied to these tracks. Their Kalman prediction
        # remains active, and time_since_update was incremented during predict().

        # 6. Create new tracks for unmatched detections.
        for det_idx in unmatched_detections:
            new_track = Tracker(detections[det_idx])
            self.tracks.append(new_track)

        # 7. Delete tracks that have exceeded the allowed number of missed frames.
        self.tracks = [t for t in self.tracks if t.time_since_update <= self.max_age]

        # 8. Output confirmed tracks.
        outputs = []

        for t in self.tracks:
            # Only output tracks that have received enough successful detections.
            if t.hits >= self.min_hits:
                bbox = t.get_state()
                outputs.append(np.concatenate([bbox, [t.id]]))

        return np.array(outputs)  # [x1, y1, x2, y2, track_id]