iou_tracker.py

# ---------------------------------------------------------
# IOU Tracker
# Copyright (c) 2017 TU Berlin, Communication Systems Group
# Licensed under The MIT License [see LICENSE for details]
# Written by Erik Bochinski
# ---------------------------------------------------------

from time import time
from util import load_mot, iou

class Tracker():
    def __init__(self, sigma_l=0, sigma_h=0.5, sigma_iou=0.5, t_max=5, verbose=False):
        self.sigma_l = sigma_l
        self.sigma_h = sigma_h
        self.sigma_iou = sigma_iou
        self.t_max = t_max
        self.frame = 0
        self.id_count = 0
        self.tracks_active = {}
        self.verbose = verbose

    #Clear the old tracks
    def clean_old_tracks(self):
        target_frame = self.frame - self.t_max
        if self.tracks_active.has_key(target_frame):
            if self.verbose:
                print("[LOG]: Tracks Deleted: {}".format(self.tracks_active[target_frame].keys()))
            del(self.tracks_active[target_frame])

    #Retrieve tracks in an correct matching order
    def retrieve_tracks(self):
        tracks = []
        frames = range(self.frame, self.frame - self.t_max, -1)
        for frame in frames: 
            if frame in self.tracks_active:
                if self.verbose:
                    print("[LOG]: Frame {} Tracks Retrieved: {}."format(frame,self.tracks_active[frame].keys()))
                tracks += self.tracks_active[frame].items()
        return tracks

    def track(self, detections):
        self.frame += 1
        self.tracks_active[self.frame] = {}
        #Clear the tracks in old frame
        self.clean_old_tracks()

        dets = [det for det in detections if det['score'] >= self.sigma_l]
        
        for id_, track in self.retrieve_tracks():
            if len(dets) > 0:
                # get det with highest iou
                best_match = max(dets, key=lambda x: iou(track['bbox'], x['bbox']))
                if iou(track['bbox'], best_match['bbox']) >= self.sigma_iou:
                    self.tracks_active[self.frame][id_] = best_match
                    if self.verbose:
                        print("[LOG]: Tracke {} updated with bbox: {}".format(id_ , best_match['bbox']))
                    # remove from best matching detection from detections
                    del dets[dets.index(best_match)]

        #Create new tracks
        for det in dets:
            self.id_count += 1
            self.tracks_active[self.frame][self.id_count] = det
            if self.verbose:
                print("[LOG]: Tracke {} added with bbox: {}".format(self.id_count, best_match['bbox']))


        #Return the current tracks 
        if self.verbose:
            print("[LOG]: Trackes {} returned".format(self.tracks_active[self.frame].keys()))
        return self.tracks_active[self.frame]  


  
def track_iou(detections, sigma_l, sigma_h, sigma_iou, t_min):
    """
    Simple IOU based tracker.
    See "High-Speed Tracking-by-Detection Without Using Image Information by E. Bochinski, V. Eiselein, T. Sikora" for
    more information.

    Args:
         detections (list): list of detections per frame, usually generated by util.load_mot
         sigma_l (float): low detection threshold.
         sigma_h (float): high detection threshold.
         sigma_iou (float): IOU threshold.
         t_min (float): minimum track length in frames.

    Returns:
        list: list of tracks.
    """

    tracks_active = []
    tracks_finished = []

    for frame_num, detections_frame in enumerate(detections, start=1):
        # apply low threshold to detections
        dets = [det for det in detections_frame if det['score'] >= sigma_l]

        updated_tracks = []
        for track in tracks_active:
            if len(dets) > 0:
                # get det with highest iou
                best_match = max(dets, key=lambda x: iou(track['bboxes'][-1], x['bbox']))
                if iou(track['bboxes'][-1], best_match['bbox']) >= sigma_iou:
                    track['bboxes'].append(best_match['bbox'])
                    track['max_score'] = max(track['max_score'], best_match['score'])

                    updated_tracks.append(track)

                    # remove from best matching detection from detections
                    del dets[dets.index(best_match)]

            # if track was not updated
            if len(updated_tracks) == 0 or track is not updated_tracks[-1]:
                # finish track when the conditions are met
                if track['max_score'] >= sigma_h and len(track['bboxes']) >= t_min:
                    tracks_finished.append(track)

        # create new tracks
        new_tracks = [{'bboxes': [det['bbox']], 'max_score': det['score'], 'start_frame': frame_num} for det in dets]
        tracks_active = updated_tracks + new_tracks

    # finish all remaining active tracks
    tracks_finished += [track for track in tracks_active
                        if track['max_score'] >= sigma_h and len(track['bboxes']) >= t_min]

    return tracks_finished


def track_iou_matlab_wrapper(detections, sigma_l, sigma_h, sigma_iou, t_min):
    """
    Matlab wrapper of the iou tracker for the detrac evaluation toolkit.

    Args:
         detections (numpy.array): numpy array of detections, usually supplied by run_tracker.m
         sigma_l (float): low detection threshold.
         sigma_h (float): high detection threshold.
         sigma_iou (float): IOU threshold.
         t_min (float): minimum track length in frames.

    Returns:
        float: speed in frames per second.
        list: list of tracks.
    """

    detections = detections.reshape((7, -1)).transpose()
    dets = load_mot(detections)
    start = time()
    tracks = track_iou(dets, sigma_l, sigma_h, sigma_iou, t_min)
    end = time()

    id_ = 1
    out = []
    for track in tracks:
        for i, bbox in enumerate(track['bboxes']):
            out += [float(bbox[0]), float(bbox[1]), float(bbox[2] - bbox[0]), float(bbox[3] - bbox[1]),
                    float(track['start_frame'] + i), float(id_)]
        id_ += 1

    num_frames = len(dets)
    speed = num_frames / (end - start)

    return speed, out