utils/tracker.py

import logging
import numpy as np
from scipy.optimize import linear_sum_assignment
import scipy.linalg


def _xyah_from_xyxy(xyxy):
    """Convert bounding box (x1, y1, x2, y2) to (center_x, center_y, aspect_ratio, height)."""
    bbox = np.asarray(xyxy).copy()
    cx = (bbox[0] + bbox[2]) / 2.0
    cy = (bbox[1] + bbox[3]) / 2.0
    w = bbox[2] - bbox[0]
    h = bbox[3] - bbox[1]

    ret = np.zeros(4, dtype=bbox.dtype)
    ret[0] = cx
    ret[1] = cy
    ret[2] = w / h
    ret[3] = h
    return ret


class KalmanFilter:
    """
    A simple Kalman Filter for tracking bounding boxes in image space.
    The 8-dimensional state space is (x, y, a, h, vx, vy, va, vh), where
    x, y is the center position, a is the aspect ratio, and h is the height.
    """

    def __init__(self):
        ndim, dt = 4, 1.0

        # Create Kalman filter model matrices.
        self._motion_mat = np.eye(2 * ndim, 2 * ndim)
        for i in range(ndim):
            self._motion_mat[i, ndim + i] = dt
        self._update_mat = np.eye(ndim, 2 * ndim)

        # Motion and observation uncertainty are chosen relative to the current
        # state estimate. These weights control the amount of uncertainty in
        # the model. This is a bit heuristic.
        self._std_weight_position = 1.0 / 20
        self._std_weight_velocity = 1.0 / 160

    def initiate(self, measurement):
        """Create track from unassociated measurement.

        Parameters
        ----------
        measurement : ndarray
            Bounding box coordinates (x1, y1, x2, y2) with confidence score.

        Returns
        -------
        (mean, covariance)
            Returns the mean vector (8 dimensional) and covariance matrix (8x8)
            of the new track.
        """
        mean_pos = _xyah_from_xyxy(measurement)
        mean = np.r_[mean_pos, np.zeros_like(mean_pos)]

        std = [
            2 * self._std_weight_position * mean_pos[3],
            2 * self._std_weight_position * mean_pos[3],
            1e-2,
            2 * self._std_weight_position * mean_pos[3],
            10 * self._std_weight_velocity * mean_pos[3],
            10 * self._std_weight_velocity * mean_pos[3],
            1e-5,
            10 * self._std_weight_velocity * mean_pos[3],
        ]
        covariance = np.diag(np.square(std))
        return mean, covariance

    def predict(self, mean, covariance):
        """Run Kalman filter prediction step.

        Parameters
        ----------
        mean : ndarray
            The 8 dimensional mean vector of the object state at the previous
            time step.
        covariance : ndarray
            The 8x8 dimensional covariance matrix of the object state at the
            previous time step.

        Returns
        -------
        (mean, covariance)
            Returns the mean vector and covariance matrix of the predicted
            state.
        """
        std_pos = [
            self._std_weight_position * mean[3],
            self._std_weight_position * mean[3],
            1e-2,
            self._std_weight_position * mean[3],
        ]

        std_vel = [
            self._std_weight_velocity * mean[3],
            self._std_weight_velocity * mean[3],
            1e-5,
            self._std_weight_velocity * mean[3],
        ]

        motion_cov = np.diag(np.square(np.r_[std_pos, std_vel]))
        mean = np.dot(self._motion_mat, mean)
        covariance = (
            np.linalg.multi_dot((self._motion_mat, covariance, self._motion_mat.T))
            + motion_cov
        )
        return mean, covariance

    def project(self, mean, covariance):
        """Project state distribution to measurement space.

        Parameters
        ----------
        mean : ndarray
            The state's mean vector (8 dimensional).
        covariance : ndarray
            The state's covariance matrix (8x8 dimensional).

        Returns
        -------
        (mean, covariance)
            Returns the projected mean and covariance matrix of the given state
            estimate.
        """
        std = [
            self._std_weight_position * mean[3],
            self._std_weight_position * mean[3],
            1e-1,
            self._std_weight_position * mean[3],
        ]

        innovation_cov = np.diag(np.square(std))
        mean = np.dot(self._update_mat, mean)
        covariance = np.linalg.multi_dot((self._update_mat, covariance, self._update_mat.T))
        return mean, covariance + innovation_cov

    def update(self, mean, covariance, measurement):
        """Run Kalman filter correction step.

        Parameters
        ----------
        mean : ndarray
            The predicted state's mean vector (8 dimensional).
        covariance : ndarray
            The state's covariance matrix (8x8 dimensional).
        measurement : ndarray
            The 4 dimensional measurement vector (x, y, a, h), where (x, y)
            is the center position, a the aspect ratio, and h the height.

        Returns
        -------
        (mean, covariance)
            Returns the measurement-corrected state distribution.
        """
        projected_mean, projected_cov = self.project(mean, covariance)
        chol_factor, lower = scipy.linalg.cho_factor(
            projected_cov, lower=True, check_finite=False
        )
        kalman_gain = scipy.linalg.cho_solve(
            (chol_factor, lower),
            np.dot(covariance, self._update_mat.T).T,
            check_finite=False,
        ).T
        innovation = measurement - projected_mean
        new_mean = mean + np.dot(innovation, kalman_gain.T)
        new_covariance = covariance - np.linalg.multi_dot(
            (kalman_gain, projected_cov, kalman_gain.T)
        )
        return new_mean, new_covariance

    def gating_distance(self, mean, covariance, measurements, only_position=False, metric="mahalanobis"):
        """Compute gating distance between state distribution and measurements."""
        mean, covariance = self.project(mean, covariance)
        if only_position:
            mean, covariance = mean[:2], covariance[:2, :2]
            measurements = measurements[:, :2]

        d = measurements - mean
        if metric == "gaussian":
            return np.sum(d * d, axis=1)
        elif metric == "mahalanobis":
            cholesky_factor = np.linalg.cholesky(covariance)
            z = scipy.linalg.solve_triangular(
                cholesky_factor, d.T, lower=True, check_finite=False, overwrite_b=True
            )
            squared_maha = np.sum(z * z, axis=0)
            return squared_maha
        else:
            raise ValueError("invalid distance metric")


class STrack:
    """
    Single object track. Wrapper around KalmanFilter state.
    """

    def __init__(self, tlwh, score, label):
        self._tlwh = np.asarray(self._tlwh_from_xyxy(tlwh), dtype=np.float32)
        self.is_activated = False
        self.track_id = 0
        self.state = 1 # 1: New, 2: Tracked, 3: Lost, 4: Removed

        self.score = score
        self.label = label
        self.start_frame = 0
        self.frame_id = 0
        self.time_since_update = 0
        self.hits = 0

        # Multi-frame history
        self.history = []

        # Kalman Filter
        self.kalman_filter = None
        self.mean = None
        self.covariance = None


    def _tlwh_from_xyxy(self, xyxy):
        """Convert xyxy to tlwh."""
        w = xyxy[2] - xyxy[0]
        h = xyxy[3] - xyxy[1]
        return [xyxy[0], xyxy[1], w, h]

    def _xyxy_from_tlwh(self, tlwh):
        """Convert tlwh to xyxy."""
        x1 = tlwh[0]
        y1 = tlwh[1]
        x2 = x1 + tlwh[2]
        y2 = y1 + tlwh[3]
        return [x1, y1, x2, y2]

    @property
    def tlwh(self):
        """Get current position in bounding box format `(top left x, top left y,
        width, height)`.
        """
        if self.mean is None:
            return self._tlwh.copy()
        ret = self.mean[:4].copy()
        ret[2] *= ret[3]
        ret[:2] -= ret[2:] / 2
        return ret

    @property
    def tlbr(self):
        """Get current position in bounding box format `(min x, min y, max x,
        max y)`.
        """
        ret = self.tlwh.copy()
        ret[2:] += ret[:2]
        return ret

    def activate(self, kalman_filter, frame_id):
        """Start a new track (tentative until min_hits reached)."""
        self.kalman_filter = kalman_filter
        self.track_id = self.next_id()
        self.mean, self.covariance = self.kalman_filter.initiate(self.tlbr)

        self.state = 2 # Tracked
        self.frame_id = frame_id
        self.start_frame = frame_id
        self.is_activated = False  # Tentative until min_hits reached
        self.hits = 1

    def re_activate(self, new_track, frame_id, new_id=False):
        """Reactivate a lost track with a new detection."""
        self.mean, self.covariance = self.kalman_filter.update(
            self.mean, self.covariance, _xyah_from_xyxy(new_track.tlbr)
        )
        self.time_since_update = 0
        self.state = 2 # Tracked
        self.frame_id = frame_id
        self.score = new_track.score

        if new_id:
            self.track_id = self.next_id()

    def update(self, new_track, frame_id, min_hits=3):
        """Update a tracked object with a new detection."""
        self.frame_id = frame_id
        self.time_since_update = 0
        self.score = new_track.score
        self.hits += 1

        self.mean, self.covariance = self.kalman_filter.update(
            self.mean, self.covariance, _xyah_from_xyxy(new_track.tlbr)
        )
        self.state = 2 # Tracked
        if self.hits >= min_hits:
            self.is_activated = True  # Confirmed after N consecutive hits

    def predict(self):
        """Propagate tracking state distribution one time step forward."""
        if self.mean is None: return
        self.mean, self.covariance = self.kalman_filter.predict(self.mean, self.covariance)

    @staticmethod
    def next_id():
        # Global counter
        if not hasattr(STrack, "_count"):
            STrack._count = 0
        STrack._count += 1
        return STrack._count

    @staticmethod
    def reset_count():
        """Reset track ID counter so the next video starts at T01."""
        STrack._count = 0


class ByteTracker:
    def __init__(self, track_high_thresh=0.3, track_low_thresh=0.1,
                 new_track_thresh=0.3, track_buffer=60, match_thresh=0.8,
                 frame_rate=30, min_hits=3):
        STrack.reset_count()
        self.track_high_thresh = track_high_thresh
        self.track_low_thresh = track_low_thresh
        self.new_track_thresh = new_track_thresh
        self.track_buffer = track_buffer
        self.match_thresh = match_thresh
        self.min_hits = min_hits
        self.frame_id = 0

        self.tracked_stracks = []  # Type: List[STrack]
        self.lost_stracks = []     # Type: List[STrack]
        self.removed_stracks = []  # Type: List[STrack]

        self.kalman_filter = KalmanFilter()

    def update(self, detections_list):
        """
        Update the tracker with a list of detections.

        Args:
            detections_list: List of dicts, each having:
                - bbox: [x1, y1, x2, y2]
                - score: float
                - label: str
                - (optional) other keys preserved

        Returns:
            List of dicts with 'track_id' added/updated.
        """
        self.frame_id += 1

        activated_stracks = []
        refind_stracks = []
        lost_stracks = []
        removed_stracks = []

        # Split into high and low confidence
        detections = []
        detections_second = []

        for d in detections_list:
            score = d['score']
            if score < self.track_low_thresh:
                continue  # Background noise -- discard entirely

            t = STrack(d['bbox'], score, d['label'])
            t.original_data = d  # Link back

            if score >= self.track_high_thresh:
                detections.append(t)
            else:
                detections_second.append(t)

        # 1. Prediction
        unconfirmed = []
        tracked_stracks = []  # Type: List[STrack]
        for track in self.tracked_stracks:
            if not track.is_activated:
                unconfirmed.append(track)
            else:
                tracked_stracks.append(track)

        strack_pool = join_stracks(tracked_stracks, self.lost_stracks)
        # Predict the current location with KF
        STrack.multi_predict(strack_pool, self.kalman_filter)

        # 2. First association (High score)
        dists = iou_distance(strack_pool, detections)
        dists = fuse_score(dists, detections)
        matches, u_track, u_detection = linear_assignment(dists, thresh=self.match_thresh)

        for itracked, idet in matches:
            track = strack_pool[itracked]
            det = detections[idet]
            if track.state == 2:
                track.update(det, self.frame_id, min_hits=self.min_hits)
                activated_stracks.append(track)
            else:
                track.re_activate(det, self.frame_id, new_id=False)
                refind_stracks.append(track)

        # 3. Second association (Low score)
        # Match unmatched tracks to low score detections
        r_tracked_stracks = [strack_pool[i] for i in u_track if strack_pool[i].state == 2]
        dists = iou_distance(r_tracked_stracks, detections_second)
        matches, u_track, u_detection_second = linear_assignment(dists, thresh=0.5)

        for itracked, idet in matches:
            track = r_tracked_stracks[itracked]
            det = detections_second[idet]
            if track.state == 2:
                track.update(det, self.frame_id, min_hits=self.min_hits)
                activated_stracks.append(track)
            else:
                track.re_activate(det, self.frame_id, new_id=False)
                refind_stracks.append(track)

        for it in u_track:
            track = r_tracked_stracks[it]
            if not track.state == 3: # If not already lost
                track.state = 3 # Lost
                lost_stracks.append(track)

        # 3.5 Associate unconfirmed tracks with remaining unmatched detections
        if unconfirmed and u_detection:
            remaining_dets = [detections[i] for i in u_detection]
            dists = iou_distance(unconfirmed, remaining_dets)
            matches_unc, u_unconfirmed, u_det_remaining = linear_assignment(dists, thresh=0.7)

            for itracked, idet in matches_unc:
                track = unconfirmed[itracked]
                det = remaining_dets[idet]
                track.update(det, self.frame_id, min_hits=self.min_hits)
                activated_stracks.append(track)

            # Update u_detection to only contain indices not matched to unconfirmed
            matched_det_indices = set(u_detection[idet] for _, idet in matches_unc) if len(matches_unc) > 0 else set()
            u_detection = [i for i in u_detection if i not in matched_det_indices]

            # Unconfirmed tracks that didn't match -> remove (too noisy to keep)
            for it in u_unconfirmed:
                track = unconfirmed[it]
                track.state = 4  # Removed
                removed_stracks.append(track)

        elif unconfirmed:
            # No detections left to match -- remove all unconfirmed
            for track in unconfirmed:
                track.state = 4  # Removed
                removed_stracks.append(track)

        # 4. Init new tracks from unmatched high score detections
        # Note: Unmatched low score detections are ignored (noise)
        unmatched_dets = [detections[i] for i in u_detection]
        rejected_by_thresh = 0
        for track in unmatched_dets:
            if track.score < self.new_track_thresh:
                rejected_by_thresh += 1
                continue  # Not confident enough to start a new track

            track.activate(self.kalman_filter, self.frame_id)
            activated_stracks.append(track)

        if rejected_by_thresh > 0 and self.frame_id <= 5:
            logging.warning(
                "Tracker frame %d: %d detections rejected by new_track_thresh=%.2f (scores: %s)",
                self.frame_id, rejected_by_thresh, self.new_track_thresh,
                [f"{t.score:.3f}" for t in unmatched_dets if t.score < self.new_track_thresh]
            )

        # 5. Update state
        self.tracked_stracks = [t for t in self.tracked_stracks if t.state == 2]
        self.tracked_stracks = join_stracks(self.tracked_stracks, activated_stracks)
        self.tracked_stracks = join_stracks(self.tracked_stracks, refind_stracks)
        self.lost_stracks = sub_stracks(self.lost_stracks, self.tracked_stracks)
        self.lost_stracks.extend(lost_stracks)
        self.lost_stracks = sub_stracks(self.lost_stracks, self.removed_stracks)
        self.removed_stracks.extend(removed_stracks)
        self.tracked_stracks, self.lost_stracks = remove_duplicate_stracks(self.tracked_stracks, self.lost_stracks)

        # 6. Age out lost tracks
        for track in self.lost_stracks:
            if self.frame_id - track.frame_id > self.track_buffer:
                self.removed_stracks.append(track)
        self.lost_stracks = [t for t in self.lost_stracks if self.frame_id - t.frame_id <= self.track_buffer]

        # 7. Collect all active tracks for output
        output_stracks = [t for t in self.tracked_stracks if t.is_activated]

        results = []
        for track in output_stracks:
            d_out = track.original_data.copy() if hasattr(track, 'original_data') else {}

            tracked_bbox = track.tlbr
            d_out['bbox'] = [float(x) for x in tracked_bbox]
            d_out['track_id'] = str(track.track_id)

            # Update bbox history
            track.history.append(d_out['bbox'])
            if len(track.history) > 30:
                track.history.pop(0)
            d_out['history'] = list(track.history)

            results.append(d_out)

        if self.frame_id % 30 == 0:
            logging.info(
                "Tracker frame %d: %d tracked (%d activated), %d lost, %d new this frame, output=%d",
                self.frame_id,
                len(self.tracked_stracks),
                sum(1 for t in self.tracked_stracks if t.is_activated),
                len(self.lost_stracks),
                len(activated_stracks),
                len(results),
            )

        return results


# --- Helper Functions ---

def linear_assignment(cost_matrix, thresh):
    """Linear assignment with threshold using scipy."""
    if cost_matrix.size == 0:
        return np.empty((0, 2), dtype=int), tuple(range(cost_matrix.shape[0])), tuple(range(cost_matrix.shape[1]))

    row_ind, col_ind = linear_sum_assignment(cost_matrix)

    matched_rows = set(row_ind)
    matched_cols = set(col_ind)

    matches = []
    unmatched_a = list(set(range(cost_matrix.shape[0])) - matched_rows)
    unmatched_b = list(set(range(cost_matrix.shape[1])) - matched_cols)

    for r, c in zip(row_ind, col_ind):
        if cost_matrix[r, c] <= thresh:
            matches.append((r, c))
        else:
            unmatched_a.append(r)
            unmatched_b.append(c)

    matches = np.array(matches) if len(matches) > 0 else np.empty((0, 2), dtype=int)
    return matches, unmatched_a, unmatched_b


def iou_distance(atracks, btracks):
    """Compute IOU cost matrix between tracks and detections."""
    if len(atracks) == 0 or len(btracks) == 0:
        return np.zeros((len(atracks), len(btracks)), dtype=float)

    atlbrs = [track.tlbr for track in atracks]
    btlbrs = [track.tlbr for track in btracks]

    _ious = bbox_ious(np.array(atlbrs), np.array(btlbrs))
    cost_matrix = 1 - _ious
    return cost_matrix

def bbox_ious(boxes1, boxes2):
    """IOU matrix."""
    b1_x1, b1_y1, b1_x2, b1_y2 = boxes1[:, 0], boxes1[:, 1], boxes1[:, 2], boxes1[:, 3]
    b2_x1, b2_y1, b2_x2, b2_y2 = boxes2[:, 0], boxes2[:, 1], boxes2[:, 2], boxes2[:, 3]

    inter_rect_x1 = np.maximum(b1_x1[:, None], b2_x1)
    inter_rect_y1 = np.maximum(b1_y1[:, None], b2_y1)
    inter_rect_x2 = np.minimum(b1_x2[:, None], b2_x2)
    inter_rect_y2 = np.minimum(b1_y2[:, None], b2_y2)

    inter_area = np.maximum(inter_rect_x2 - inter_rect_x1, 0) * np.maximum(inter_rect_y2 - inter_rect_y1, 0)

    b1_area = (b1_x2 - b1_x1) * (b1_y2 - b1_y1)
    b2_area = (b2_x2 - b2_x1) * (b2_y2 - b2_y1)

    iou = inter_area / (b1_area[:, None] + b2_area - inter_area + 1e-6)
    return iou


def fuse_score(cost_matrix, detections):
    """Refine cost matrix with detection scores."""
    if cost_matrix.size == 0: return cost_matrix
    iou_sim = 1 - cost_matrix
    det_scores = np.array([d.score for d in detections])
    det_scores = np.expand_dims(det_scores, axis=0).repeat(cost_matrix.shape[0], axis=0)
    fuse_sim = iou_sim * det_scores
    fuse_cost = 1 - fuse_sim
    return fuse_cost


# STrack collection helpers

def join_stracks(tlist_a, tlist_b):
    exists = {}
    res = []
    for t in tlist_a:
        exists[t.track_id] = 1
        res.append(t)
    for t in tlist_b:
        tid = t.track_id
        if not exists.get(tid, 0):
            exists[tid] = 1
            res.append(t)
    return res

def sub_stracks(tlist_a, tlist_b):
    stracks = {}
    for t in tlist_a:
        stracks[t.track_id] = t
    for t in tlist_b:
        tid = t.track_id
        if stracks.get(tid, 0):
            del stracks[tid]
    return list(stracks.values())

def remove_duplicate_stracks(stracksa, stracksb):
    pdist = iou_distance(stracksa, stracksb)
    pairs = np.where(pdist < 0.15)
    dupa, dupb = set(), set()
    for a, b in zip(pairs[0], pairs[1]):
        time_a = stracksa[a].frame_id - stracksa[a].start_frame
        time_b = stracksb[b].frame_id - stracksb[b].start_frame
        if time_a > time_b:
            dupb.add(b)
        else:
            dupa.add(a)

    res_a = [t for i, t in enumerate(stracksa) if i not in dupa]
    res_b = [t for i, t in enumerate(stracksb) if i not in dupb]
    return res_a, res_b


def _multi_predict(stracks, kalman_filter):
    for t in stracks:
        if t.state != 2:
            t.mean[7] = 0 # reset velocity h if lost
        t.mean, t.covariance = kalman_filter.predict(t.mean, t.covariance)

STrack.multi_predict = staticmethod(_multi_predict)