Class SORTBox

deeplabcut/core/trackingutils.py:613–708

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)