Link to this sectionReference for ultralytics/trackers/bot_sort.py#

class BOTrack(STrack):
    """An extended version of the STrack class for YOLO, adding object tracking features.

    This class extends the STrack class to include additional functionalities for object tracking, such as feature
    smoothing, Kalman filter prediction, and reactivation of tracks.

    Attributes:
        shared_kalman (KalmanFilterXYWH): A shared Kalman filter for all instances of BOTrack.
        smooth_feat (np.ndarray): Smoothed feature vector.
        curr_feat (np.ndarray): Current feature vector.
        alpha (float): Smoothing factor for the exponential moving average of features.
        mean (np.ndarray): The mean state of the Kalman filter.
        covariance (np.ndarray): The covariance matrix of the Kalman filter.

    Methods:
        update_features: Update features vector and smooth it using exponential moving average.
        predict: Predict the mean and covariance using Kalman filter.
        re_activate: Reactivate a track with updated features and optionally new ID.
        update: Update the track with new detection and frame ID.
        tlwh: Property that gets the current position in tlwh format `(top left x, top left y, width, height)`.
        multi_predict: Predict the mean and covariance of multiple object tracks using shared Kalman filter.
        convert_coords: Convert tlwh bounding box coordinates to xywh format.
        tlwh_to_xywh: Convert bounding box to xywh format `(center x, center y, width, height)`.

    Examples:
        Create a BOTrack instance and update its features
        >>> bo_track = BOTrack(xywh=np.array([100, 50, 80, 40, 0]), score=0.9, cls=1, feat=np.random.rand(128))
        >>> bo_track.predict()
        >>> new_track = BOTrack(xywh=np.array([110, 60, 80, 40, 0]), score=0.85, cls=1, feat=np.random.rand(128))
        >>> bo_track.update(new_track, frame_id=2)
    """

    shared_kalman = KalmanFilterXYWH()

    def __init__(self, xywh: np.ndarray, score: float, cls: int, feat: np.ndarray | None = None):
        """Initialize a BOTrack object with feature-smoothing state and a Kalman filter.

        Args:
            xywh (np.ndarray): Bounding box in `(x, y, w, h, idx)` or `(x, y, w, h, angle, idx)` format, where (x, y) is
                the center, (w, h) are width and height, and `idx` is the detection index.
            score (float): Confidence score of the detection.
            cls (int): Class ID of the detected object.
            feat (np.ndarray, optional): Feature vector associated with the detection.
        """
        super().__init__(xywh, score, cls)

        self.smooth_feat = None
        self.curr_feat = None
        self.alpha = 0.9
        if feat is not None:
            self.update_features(feat)

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

def convert_coords(self, tlwh: np.ndarray) -> np.ndarray:
    """Convert tlwh bounding box coordinates to xywh format."""
    return self.tlwh_to_xywh(tlwh)

@staticmethod
def multi_predict(stracks: list[BOTrack]) -> None:
    """Predict the mean and covariance for multiple object tracks using a shared Kalman filter."""
    if not stracks:
        return
    multi_mean = np.asarray([st.mean.copy() for st in stracks])
    multi_covariance = np.asarray([st.covariance for st in stracks])
    for i, st in enumerate(stracks):
        if st.state != TrackState.Tracked:
            multi_mean[i][6] = 0
            multi_mean[i][7] = 0
    multi_mean, multi_covariance = BOTrack.shared_kalman.multi_predict(multi_mean, multi_covariance)
    for i, (mean, cov) in enumerate(zip(multi_mean, multi_covariance)):
        stracks[i].mean = mean
        stracks[i].covariance = cov

def predict(self) -> None:
    """Predict the object's future state using the Kalman filter to update its mean and covariance."""
    mean_state = self.mean.copy()
    if self.state != TrackState.Tracked:
        mean_state[6] = 0
        mean_state[7] = 0

    self.mean, self.covariance = self.kalman_filter.predict(mean_state, self.covariance)

def re_activate(self, new_track: BOTrack, frame_id: int, new_id: bool = False) -> None:
    """Reactivate a track with updated features and optionally assign a new ID."""
    if new_track.curr_feat is not None:
        self.update_features(new_track.curr_feat)
    super().re_activate(new_track, frame_id, new_id)

@staticmethod
def tlwh_to_xywh(tlwh: np.ndarray) -> np.ndarray:
    """Convert bounding box from tlwh (top-left-width-height) to xywh (center-x-center-y-width-height) format."""
    ret = np.asarray(tlwh).copy()
    ret[:2] += ret[2:] / 2
    return ret

def update(self, new_track: BOTrack, frame_id: int) -> None:
    """Update the track with new detection information and the current frame ID."""
    if new_track.curr_feat is not None:
        self.update_features(new_track.curr_feat)
    super().update(new_track, frame_id)

def update_features(self, feat: np.ndarray) -> None:
    """Update the current feature and its exponential-moving-average smoothed feature."""
    curr, smooth = smooth_feature(feat, self.smooth_feat, self.alpha)
    if curr is not None:
        self.curr_feat, self.smooth_feat = curr, smooth

class BOTSORT(BYTETracker):
    """An extended version of the BYTETracker class for YOLO, designed for object tracking with ReID and GMC algorithm.

    Attributes:
        proximity_thresh (float): Threshold for spatial proximity (IoU) between tracks and detections.
        appearance_thresh (float): Threshold for appearance similarity (ReID embeddings) between tracks and detections.
        encoder (Any): Object to handle ReID embeddings, set to None if ReID is not enabled.
        gmc (GMC): An instance of the GMC algorithm for data association.
        args (Any): Parsed command-line arguments containing tracking parameters.

    Methods:
        get_kalmanfilter: Return an instance of KalmanFilterXYWH for object tracking.
        init_track: Initialize track with detection results and optional image for ReID.
        get_dists: Get distances between tracks and detections using IoU and (optionally) ReID.
        multi_predict: Predict the mean and covariance of multiple object tracks using a shared Kalman filter.
        reset: Reset the BOTSORT tracker to its initial state.

    Examples:
        Initialize BOTSORT and process detections
        >>> bot_sort = BOTSORT(args)
        >>> bot_sort.init_track(results, img)
        >>> bot_sort.multi_predict(tracks)

    Notes:
        The class is designed to work with a YOLO object detection model and supports ReID only if enabled via args.
    """

    def __init__(self, args: Any):
        """Initialize BOTSORT object with ReID module and GMC algorithm.

        Args:
            args (Any): Parsed command-line arguments containing tracking parameters.
        """
        super().__init__(args)
        self.gmc = GMC(method=args.gmc_method)

        # ReID module
        self.proximity_thresh = args.proximity_thresh
        self.appearance_thresh = args.appearance_thresh
        self.encoder = build_encoder(args.with_reid, args.model)

def get_dists(self, tracks: list[BOTrack], detections: list[BOTrack]) -> np.ndarray:
    """Calculate distances between tracks and detections using IoU and optionally ReID embeddings."""
    dists = matching.iou_distance(tracks, detections)
    dists_mask = dists > (1 - self.proximity_thresh)

    if self.args.fuse_score:
        dists = matching.fuse_score(dists, detections)

    if self.args.with_reid and self.encoder is not None:
        emb_dists = matching.embedding_distance(tracks, detections) / 2.0
        emb_dists[emb_dists > (1 - self.appearance_thresh)] = 1.0
        emb_dists[dists_mask] = 1.0
        dists = np.minimum(dists, emb_dists)
    return dists

def get_kalmanfilter(self) -> KalmanFilterXYWH:
    """Return an instance of KalmanFilterXYWH for predicting and updating object states in the tracking process."""
    return KalmanFilterXYWH()

def init_track(self, results, img: np.ndarray | None = None) -> list[BOTrack]:
    """Initialize object tracks using detection bounding boxes, scores, class labels, and optional ReID features."""
    if len(results) == 0:
        return []
    bboxes = parse_bboxes(results)
    if self.args.with_reid and self.encoder is not None:
        features_keep = self.encoder(img, bboxes)
        return [BOTrack(xywh, s, c, f) for (xywh, s, c, f) in zip(bboxes, results.conf, results.cls, features_keep)]
    return [BOTrack(xywh, s, c) for (xywh, s, c) in zip(bboxes, results.conf, results.cls)]

def multi_predict(self, tracks: list[BOTrack]) -> None:
    """Predict the mean and covariance of multiple object tracks using a shared Kalman filter."""
    BOTrack.multi_predict(tracks)

def reset(self) -> None:
    """Reset the BOTSORT tracker to its initial state, clearing all tracked objects and internal states."""
    super().reset()
    self.gmc.reset_params()