Skip to content

deeplabcut.core.trackingutils

Classes:

Name Description
BaseTracker

Base class for a constant-velocity Kalman filter-based tracker.

BoxTracker
EllipseFitter
SORTBox

Functions:

Name Description
reconstruct_all_ellipses

Reconstructs ellipses for multiple individuals based on their body part

BaseTracker

Base class for a constant-velocity Kalman filter-based tracker.

Source code in deeplabcut/core/trackingutils.py
class BaseTracker:
    """Base class for a constant-velocity Kalman filter-based tracker."""

    n_trackers = 0

    def __init__(self, dim, dim_z):
        self.kf = kinematic_kf(
            dim,
            1,
            dim_z=dim_z,
            order_by_dim=False,
        )
        self.id = self.__class__.n_trackers
        self.__class__.n_trackers += 1
        self.time_since_update = 0
        self.age = 0
        self.hits = 0
        self.hit_streak = 0

    def update(self, z):
        self.time_since_update = 0
        self.hits += 1
        self.hit_streak += 1
        self.kf.update(z)

    def predict(self):
        self.kf.predict()
        self.age += 1
        if self.time_since_update > 0:
            self.hit_streak = 0
        self.time_since_update += 1
        return self.state

    @property
    def state(self):
        return self.kf.x.squeeze()[: self.kf.dim_z]

    @state.setter
    def state(self, state):
        self.kf.x[: self.kf.dim_z] = state

BoxTracker

Bases: BaseTracker

Methods:

Name Description
convert_bbox_to_z

Takes a bounding box in the form [x1,y1,x2,y2] and returns z in the form

convert_x_to_bbox

Takes a bounding box in the centre form [x,y,s,r] and returns it in the form

Source code in deeplabcut/core/trackingutils.py
class BoxTracker(BaseTracker):
    def __init__(self, bbox):
        super().__init__(dim=4, dim_z=4)
        self.kf = KalmanFilter(dim_x=7, dim_z=4)
        self.kf.F = np.array(
            [
                [1, 0, 0, 0, 1, 0, 0],
                [0, 1, 0, 0, 0, 1, 0],
                [0, 0, 1, 0, 0, 0, 1],
                [0, 0, 0, 1, 0, 0, 0],
                [0, 0, 0, 0, 1, 0, 0],
                [0, 0, 0, 0, 0, 1, 0],
                [0, 0, 0, 0, 0, 0, 1],
            ]
        )
        self.kf.H = np.array(
            [
                [1, 0, 0, 0, 0, 0, 0],
                [0, 1, 0, 0, 0, 0, 0],
                [0, 0, 1, 0, 0, 0, 0],
                [0, 0, 0, 1, 0, 0, 0],
            ]
        )
        self.kf.R[2:, 2:] *= 10.0
        # Give high uncertainty to the unobservable initial velocities
        self.kf.P[4:, 4:] *= 1000.0
        self.kf.P *= 10.0
        self.kf.Q[-1, -1] *= 0.01
        self.kf.Q[4:, 4:] *= 0.01
        self.state = bbox

    def update(self, bbox):
        super().update(self.convert_bbox_to_z(bbox))

    def predict(self):
        if (self.kf.x[6] + self.kf.x[2]) <= 0:
            self.kf.x[6] *= 0.0
        return super().predict()

    @property
    def state(self):
        return self.convert_x_to_bbox(self.kf.x)[0]

    @state.setter
    def state(self, bbox):
        state = self.convert_bbox_to_z(bbox)
        super(BoxTracker, type(self)).state.fset(self, state)

    @staticmethod
    def convert_x_to_bbox(x, score=None):
        """Takes a bounding box in the centre form [x,y,s,r] and returns it in the form
        [x1,y1,x2,y2] where x1,y1 is the top left and x2,y2 is the bottom right."""
        w = np.sqrt(x[2] * x[3])
        h = x[2] / w
        if score is None:
            return np.array([x[0] - w / 2.0, x[1] - h / 2.0, x[0] + w / 2.0, x[1] + h / 2.0]).reshape((1, 4))
        else:
            return np.array([x[0] - w / 2.0, x[1] - h / 2.0, x[0] + w / 2.0, x[1] + h / 2.0, score]).reshape((1, 5))

    @staticmethod
    def convert_bbox_to_z(bbox):
        """Takes a bounding box in the form [x1,y1,x2,y2] and returns z in the form
        [x,y,s,r] where x,y is the centre of the box and s is the scale/area and r is
        the aspect ratio."""
        w = bbox[2] - bbox[0]
        h = bbox[3] - bbox[1]
        x = bbox[0] + w / 2.0
        y = bbox[1] + h / 2.0
        s = w * h  # scale is just area
        r = w / float(h)
        return np.array([x, y, s, r]).reshape((4, 1))

convert_bbox_to_z staticmethod

convert_bbox_to_z(bbox)

Takes a bounding box in the form [x1,y1,x2,y2] and returns z in the form [x,y,s,r] where x,y is the centre of the box and s is the scale/area and r is the aspect ratio.

Source code in deeplabcut/core/trackingutils.py
@staticmethod
def convert_bbox_to_z(bbox):
    """Takes a bounding box in the form [x1,y1,x2,y2] and returns z in the form
    [x,y,s,r] where x,y is the centre of the box and s is the scale/area and r is
    the aspect ratio."""
    w = bbox[2] - bbox[0]
    h = bbox[3] - bbox[1]
    x = bbox[0] + w / 2.0
    y = bbox[1] + h / 2.0
    s = w * h  # scale is just area
    r = w / float(h)
    return np.array([x, y, s, r]).reshape((4, 1))

convert_x_to_bbox staticmethod

convert_x_to_bbox(x, score=None)

Takes a bounding box in the centre form [x,y,s,r] and returns it in the form [x1,y1,x2,y2] where x1,y1 is the top left and x2,y2 is the bottom right.

Source code in deeplabcut/core/trackingutils.py
@staticmethod
def convert_x_to_bbox(x, score=None):
    """Takes a bounding box in the centre form [x,y,s,r] and returns it in the form
    [x1,y1,x2,y2] where x1,y1 is the top left and x2,y2 is the bottom right."""
    w = np.sqrt(x[2] * x[3])
    h = x[2] / w
    if score is None:
        return np.array([x[0] - w / 2.0, x[1] - h / 2.0, x[0] + w / 2.0, x[1] + h / 2.0]).reshape((1, 4))
    else:
        return np.array([x[0] - w / 2.0, x[1] - h / 2.0, x[0] + w / 2.0, x[1] + h / 2.0, score]).reshape((1, 5))

EllipseFitter

Methods:

Name Description
calc_parameters

Calculate ellipse center coordinates, semi-axes lengths, and

Source code in deeplabcut/core/trackingutils.py
class EllipseFitter:
    def __init__(self, sd=2):
        self.sd = sd
        self.x = None
        self.y = None
        self.params = None
        self._coeffs = None

    def fit(self, xy):
        self.x, self.y = xy[np.isfinite(xy).all(axis=1)].T
        if len(self.x) < 3:
            return None
        if self.sd:
            self.params = self._fit_error(self.x, self.y, self.sd)
        else:
            self._coeffs = self._fit(self.x, self.y)
            self.params = self.calc_parameters(self._coeffs)
        if not np.isnan(self.params).any():
            return Ellipse(*self.params)
        return None

    @staticmethod
    @jit(nopython=True)
    def _fit(x, y):
        """Least Squares ellipse fitting algorithm Fit an ellipse to a set of X- and
        Y-coordinates. See Halir and Flusser, 1998 for implementation details.

        :param x: ndarray, 1D trajectory
        :param y: ndarray, 1D trajectory
        :return: 1D ndarray of 6 coefficients of the general quadratic curve: ax^2 +
            2bxy + cy^2 + 2dx + 2fy + g = 0
        """
        D1 = np.vstack((x * x, x * y, y * y))
        D2 = np.vstack((x, y, np.ones_like(x)))
        S1 = D1 @ D1.T
        S2 = D1 @ D2.T
        S3 = D2 @ D2.T
        T = -np.linalg.inv(S3) @ S2.T
        temp = S1 + S2 @ T
        M = np.zeros_like(temp)
        M[0] = temp[2] * 0.5
        M[1] = -temp[1]
        M[2] = temp[0] * 0.5
        E, V = np.linalg.eig(M)
        cond = 4 * V[0] * V[2] - V[1] ** 2
        a1 = V[:, cond > 0][:, 0]
        a2 = T @ a1
        return np.hstack((a1, a2))

    @staticmethod
    @jit(nopython=True)
    def _fit_error(x, y, sd):
        """Fit a sd-sigma covariance error ellipse to the data.

        :param x: ndarray, 1D input of X coordinates
        :param y: ndarray, 1D input of Y coordinates
        :param sd: int, size of the error ellipse in 'standard deviation'
        :return: ellipse center, semi-axes length, angle to the X-axis
        """
        cov = np.cov(x, y)
        E, V = np.linalg.eigh(cov)  # Returns the eigenvalues in ascending order
        # r2 = chi2.ppf(2 * norm.cdf(sd) - 1, 2)
        # height, width = np.sqrt(E * r2)
        height, width = 2 * sd * np.sqrt(E)
        a, b = V[:, 1]
        rotation = math.atan2(b, a) % np.pi
        return [np.mean(x), np.mean(y), width, height, rotation]

    @staticmethod
    @jit(nopython=True)
    def calc_parameters(coeffs):
        """
        Calculate ellipse center coordinates, semi-axes lengths, and
        the counterclockwise angle of rotation from the x-axis to the ellipse major axis.
        Visit http://mathworld.wolfram.com/Ellipse.html
        for how to estimate ellipse parameters.

        :param coeffs: list of fitting coefficients
        :return: center: 1D ndarray, semi-axes: 1D ndarray, angle: float
        """
        # The general quadratic curve has the form:
        # ax^2 + 2bxy + cy^2 + 2dx + 2fy + g = 0
        a, b, c, d, f, g = coeffs
        b *= 0.5
        d *= 0.5
        f *= 0.5

        # Ellipse center coordinates
        x0 = (c * d - b * f) / (b * b - a * c)
        y0 = (a * f - b * d) / (b * b - a * c)

        # Semi-axes lengths
        num = 2 * (a * f * f + c * d * d + g * b * b - 2 * b * d * f - a * c * g)
        den1 = (b * b - a * c) * (np.sqrt((a - c) ** 2 + 4 * b * b) - (a + c))
        den2 = (b * b - a * c) * (-np.sqrt((a - c) ** 2 + 4 * b * b) - (a + c))
        major = np.sqrt(num / den1)
        minor = np.sqrt(num / den2)

        # Angle to the horizontal
        if b == 0:
            if a < c:
                phi = 0
            else:
                phi = np.pi / 2
        else:
            if a < c:
                phi = np.arctan(2 * b / (a - c)) / 2
            else:
                phi = np.pi / 2 + np.arctan(2 * b / (a - c)) / 2

        return [x0, y0, 2 * major, 2 * minor, phi]

calc_parameters staticmethod

calc_parameters(coeffs)

Calculate ellipse center coordinates, semi-axes lengths, and the counterclockwise angle of rotation from the x-axis to the ellipse major axis. Visit http://mathworld.wolfram.com/Ellipse.html for how to estimate ellipse parameters.

:param coeffs: list of fitting coefficients :return: center: 1D ndarray, semi-axes: 1D ndarray, angle: float

Source code in deeplabcut/core/trackingutils.py
@staticmethod
@jit(nopython=True)
def calc_parameters(coeffs):
    """
    Calculate ellipse center coordinates, semi-axes lengths, and
    the counterclockwise angle of rotation from the x-axis to the ellipse major axis.
    Visit http://mathworld.wolfram.com/Ellipse.html
    for how to estimate ellipse parameters.

    :param coeffs: list of fitting coefficients
    :return: center: 1D ndarray, semi-axes: 1D ndarray, angle: float
    """
    # The general quadratic curve has the form:
    # ax^2 + 2bxy + cy^2 + 2dx + 2fy + g = 0
    a, b, c, d, f, g = coeffs
    b *= 0.5
    d *= 0.5
    f *= 0.5

    # Ellipse center coordinates
    x0 = (c * d - b * f) / (b * b - a * c)
    y0 = (a * f - b * d) / (b * b - a * c)

    # Semi-axes lengths
    num = 2 * (a * f * f + c * d * d + g * b * b - 2 * b * d * f - a * c * g)
    den1 = (b * b - a * c) * (np.sqrt((a - c) ** 2 + 4 * b * b) - (a + c))
    den2 = (b * b - a * c) * (-np.sqrt((a - c) ** 2 + 4 * b * b) - (a + c))
    major = np.sqrt(num / den1)
    minor = np.sqrt(num / den2)

    # Angle to the horizontal
    if b == 0:
        if a < c:
            phi = 0
        else:
            phi = np.pi / 2
    else:
        if a < c:
            phi = np.arctan(2 * b / (a - c)) / 2
        else:
            phi = np.pi / 2 + np.arctan(2 * b / (a - c)) / 2

    return [x0, y0, 2 * major, 2 * minor, phi]

SORTBox

Bases: SORTBase

Methods:

Name Description
match_detections_to_trackers

Assigns detections to tracked object (both represented as bounding boxes)

Source code in deeplabcut/core/trackingutils.py
class SORTBox(SORTBase):
    def __init__(self, max_age, min_hits, iou_threshold):
        self.max_age = max_age
        self.min_hits = min_hits
        self.iou_threshold = iou_threshold
        BoxTracker.n_trackers = 0
        super().__init__()

    def track(self, dets):
        self.n_frames += 1

        trackers = np.zeros((len(self.trackers), 5))
        for i in range(len(trackers)):
            trackers[i, :4] = self.trackers[i].predict()
        empty = np.isnan(trackers).any(axis=1)
        trackers = trackers[~empty]
        for ind in np.flatnonzero(empty)[::-1]:
            self.trackers.pop(ind)

        matched, unmatched_dets, unmatched_trks = self.match_detections_to_trackers(dets, trackers, self.iou_threshold)

        # update matched trackers with assigned detections
        animalindex = []
        for t, trk in enumerate(self.trackers):
            if t not in unmatched_trks:
                d = matched[np.where(matched[:, 1] == t)[0], 0]
                animalindex.append(d[0])
                trk.update(dets[d, :][0])  # update coordinates
            else:
                animalindex.append("nix")  # lost trk!

        # create and initialise new trackers for unmatched detections
        for i in unmatched_dets:
            trk = BoxTracker(dets[i, :])
            self.trackers.append(trk)
            animalindex.append(i)

        i = len(self.trackers)
        ret = []
        for trk in reversed(self.trackers):
            d = trk.state
            if (trk.time_since_update < 1) and (trk.hit_streak >= self.min_hits or self.n_frames <= self.min_hits):
                ret.append(
                    np.concatenate((d, [trk.id, int(animalindex[i - 1])])).reshape(1, -1)
                )  # for DLC we also return the original animalid
                # +1 as MOT benchmark requires positive >> this is removed for DLC!
            i -= 1
            # remove dead tracklet
            if trk.time_since_update > self.max_age:
                self.trackers.pop(i)

        if len(ret) > 0:
            return np.concatenate(ret)
        return np.empty((0, 5))

    @staticmethod
    def match_detections_to_trackers(detections, trackers, iou_threshold):
        """Assigns detections to tracked object (both represented as bounding boxes)

        Returns 3 lists of matches, unmatched_detections and unmatched_trackers
        """
        if not len(trackers):
            return (
                np.empty((0, 2), dtype=int),
                np.arange(len(detections)),
                np.empty((0, 5), dtype=int),
            )
        iou_matrix = np.zeros((len(detections), len(trackers)), dtype=np.float32)

        for d, det in enumerate(detections):
            for t, trk in enumerate(trackers):
                iou_matrix[d, t] = calc_iou(det, trk)
        row_indices, col_indices = linear_sum_assignment(-iou_matrix)

        unmatched_detections = []
        for d, _ in enumerate(detections):
            if d not in row_indices:
                unmatched_detections.append(d)
        unmatched_trackers = []
        for t, _ in enumerate(trackers):
            if t not in col_indices:
                unmatched_trackers.append(t)

        # filter out matched with low IOU
        matches = []
        for row, col in zip(row_indices, col_indices, strict=False):
            if iou_matrix[row, col] < iou_threshold:
                unmatched_detections.append(row)
                unmatched_trackers.append(col)
            else:
                matches.append([row, col])
        if not len(matches):
            matches = np.empty((0, 2), dtype=int)
        else:
            matches = np.stack(matches)
        return matches, np.array(unmatched_detections), np.array(unmatched_trackers)

match_detections_to_trackers staticmethod

match_detections_to_trackers(detections, trackers, iou_threshold)

Assigns detections to tracked object (both represented as bounding boxes)

Returns 3 lists of matches, unmatched_detections and unmatched_trackers

Source code in deeplabcut/core/trackingutils.py
@staticmethod
def match_detections_to_trackers(detections, trackers, iou_threshold):
    """Assigns detections to tracked object (both represented as bounding boxes)

    Returns 3 lists of matches, unmatched_detections and unmatched_trackers
    """
    if not len(trackers):
        return (
            np.empty((0, 2), dtype=int),
            np.arange(len(detections)),
            np.empty((0, 5), dtype=int),
        )
    iou_matrix = np.zeros((len(detections), len(trackers)), dtype=np.float32)

    for d, det in enumerate(detections):
        for t, trk in enumerate(trackers):
            iou_matrix[d, t] = calc_iou(det, trk)
    row_indices, col_indices = linear_sum_assignment(-iou_matrix)

    unmatched_detections = []
    for d, _ in enumerate(detections):
        if d not in row_indices:
            unmatched_detections.append(d)
    unmatched_trackers = []
    for t, _ in enumerate(trackers):
        if t not in col_indices:
            unmatched_trackers.append(t)

    # filter out matched with low IOU
    matches = []
    for row, col in zip(row_indices, col_indices, strict=False):
        if iou_matrix[row, col] < iou_threshold:
            unmatched_detections.append(row)
            unmatched_trackers.append(col)
        else:
            matches.append([row, col])
    if not len(matches):
        matches = np.empty((0, 2), dtype=int)
    else:
        matches = np.stack(matches)
    return matches, np.array(unmatched_detections), np.array(unmatched_trackers)

reconstruct_all_ellipses

reconstruct_all_ellipses(data, sd)

Reconstructs ellipses for multiple individuals based on their body part coordinates across multiple frames. Each ellipse is fitted to the coordinates using an EllipseFitter.

Parameters

data : pandas.DataFrame A multi-level DataFrame containing body part coordinates and likelihood values. The index represents frames, and the columns follow a multi-level structure: - Level 0: Scorer - Level 1: Individuals - Level 2: Body parts - Level 3: Coordinates ("x" and "y") and "likelihood". sd : float The standard deviation used by the EllipseFitter for fitting ellipses.

Returns

numpy.ndarray A 3D array of shape (A, F, 5), where: - A is the number of individuals (excluding "single" if present). - F is the number of frames. - Each row contains ellipse parameters [cx, cy, width, height, angle].

Notes

  • The method drops the "likelihood" column from the input DataFrame as it is not relevant for ellipse fitting.
  • If the "single" individual is present, it is excluded from the reconstruction process.
  • The EllipseFitter is used to fit ellipses to the body part coordinates for each individual in each frame.
  • NaN values are assigned when no valid ellipse can be fitted.
Source code in deeplabcut/core/trackingutils.py
def reconstruct_all_ellipses(data, sd):
    """Reconstructs ellipses for multiple individuals based on their body part
    coordinates across multiple frames. Each ellipse is fitted to the coordinates using
    an `EllipseFitter`.

    Parameters
    ----------
    data : pandas.DataFrame
        A multi-level DataFrame containing body part coordinates and likelihood values.
        The index represents frames, and the columns follow a multi-level structure:
        - Level 0: Scorer
        - Level 1: Individuals
        - Level 2: Body parts
        - Level 3: Coordinates ("x" and "y") and "likelihood".
    sd : float
        The standard deviation used by the `EllipseFitter` for fitting ellipses.

    Returns
    -------
    numpy.ndarray
        A 3D array of shape (A, F, 5), where:
        - A is the number of individuals (excluding "single" if present).
        - F is the number of frames.
        - Each row contains ellipse parameters [cx, cy, width, height, angle].

    Notes
    -----
    - The method drops the "likelihood" column from the input DataFrame as it is not
      relevant for ellipse fitting.
    - If the "single" individual is present, it is excluded from the reconstruction process.
    - The `EllipseFitter` is used to fit ellipses to the body part coordinates for each
      individual in each frame.
    - NaN values are assigned when no valid ellipse can be fitted.
    """
    xy = data.droplevel("scorer", axis=1).drop("likelihood", axis=1, level=-1)
    if "single" in xy:
        xy.drop("single", axis=1, level="individuals", inplace=True)
    animals = xy.columns.get_level_values("individuals").unique()
    nrows = xy.shape[0]
    ellipses = np.full((len(animals), nrows, 5), np.nan)
    fitter = EllipseFitter(sd)
    for n, animal in enumerate(animals):
        data = xy.xs(animal, axis=1, level="individuals").values.reshape((nrows, -1, 2))
        for i, coords in enumerate(tqdm(data)):
            el = fitter.fit(coords.astype(np.float64))
            if el is not None:
                ellipses[n, i] = el.parameters
    return ellipses