Running DSST tracker on video with ground truth

dsst_trial1.py

import os
import glob

import dlib
import numpy as np

import pyTools.system.videoExplorer as VE
import pyTools.videoProc.annotation as A

# HELPER FUNCTIONS
def calculateOverlap(rect_a, rect_b):
    x11 = rect_a[0]
    x12 = rect_a[2]
    y11 = rect_a[1]
    y12 = rect_a[3]
    x21 = rect_b.left()
    x22 = rect_b.right()
    y21 = rect_b.top()
    y22 = rect_b.bottom()

    x_overlap = np.max([0, np.min([x12,x22]) - np.max([x11,x21])])
    y_overlap = np.max([0, np.min([y12,y22]) - np.max([y11,y21])])
    overlap_area = x_overlap * y_overlap

    rect_a_area =  np.abs(x11 - x12) * np.abs(y11 - y12)

    if overlap_area < rect_a_area:
        overlap = overlap_area / float(rect_a_area)
    else:
        overlap = rect_a_area / float(overlap_area)

    return overlap

def rectangle2array(rect):
    return np.array([rect.left(), rect.top(), rect.right(), rect.bottom()])

def rectangle2list(rect):
    return [rect.left(), rect.top(), rect.right(), rect.bottom()]


# GENERAL SETUP

# Path to the video frames
video_file = '/Volumes/Seagate_Backup_Plus_Drive 1/datasets/volleyball/graz-arbesback_5/graz-arbesbach_5_full.mp4'
anno_file = 'graz-arbesbach_5_small.csv'
lbls = ['player10', 'player14', 'player18', 'player2', 'player4', 'player9', 'player3']


vE = VE.videoExplorer()
vE.setVideoStream(video_file, frameMode='RGB')

# setup annotations
anno = A.Annotation()
anno.loadFromFile(anno_file)

filteredAnnos = []
for lbl in lbls:
    # get one annotation for each label
    filterTuple = A.AnnotationFilter(vials=None, annotators=['Peter'], behaviours=[lbl])
    filteredAnnos += [anno.filterFrameList(filterTuple, exactMatch=False)]

# allotcate Annotation object for result
resAnno = A.Annotation()


# Create the correlation tracker - the object needs to be initialized
# before it can be used
trackers = []
lost_tracks = []
succ_tracks = []
inits = []
for lbl in lbls:
    trackers += [dlib.correlation_tracker()]

    # win = dlib.image_window()
    # We will track the frames as we load them off of disk

    lost_tracks += [{'frames': [],
                     'meta': {}}]
    succ_tracks += [{'frames': [],
                     'meta': {}}]

    inits += [False]

# TRACKING
for k, frame in enumerate(vE):
    print("Processing Frame {}".format(k))

    for l in range(len(lbls)):
        filteredAnno = filteredAnnos[l]
        tracker = trackers[l]
        lost_track = lost_tracks[l]
        succ_track = succ_tracks[l]

        try:
            # get bounding-box of current frame number
            gt_bb = A.getPropertyFromFrameAnno(filteredAnno.getFrame(k), 'boundingBox')[0].astype(np.long)
        except TypeError:
            continue

        if not inits[l]:
            # initialize tracker
            tracker.start_track(frame, dlib.rectangle(*list(gt_bb)))
            inits[l] = True

        tracker.update(frame)
        pos = tracker.get_position()

        overlap = calculateOverlap(gt_bb, pos)

        if overlap < 0.25:
            # tracker lost target
            tracker.start_track(frame, dlib.rectangle(*list(gt_bb)))
            lost_track['frames'] += [k]
            lost_track['meta'][k] = {"confidence": 1,
                                  "boundingBox": rectangle2list(pos)}
        else:
            # tracking successful
            succ_track['frames'] += [k]
            succ_track['meta'][k] = {"confidence": 1,
                                  "boundingBox": rectangle2list(pos)}



# SAVE RESULT IN ANNOTATIONS
resAnno = A.Annotation(frameNo=anno.getLength())
for l, lbl in enumerate(lbls):
    filteredAnno = filteredAnnos[l]
    tracker = trackers[l]
    lost_track = lost_tracks[l]
    succ_track = succ_tracks[l]
    if lost_track['frames']:
        resAnno.addAnnotation(vial=None,
                              frames=lost_track['frames'],
                              annotator='DSST_TRACKER',
                              behaviour='{l} - lost track'.format(l=lbl),
                              metadata=lost_track['meta'])
    if succ_track['frames']:
        resAnno.addAnnotation(vial=None,
                              frames=succ_track['frames'],
                              annotator='DSST_TRACKER',
                              behaviour='{l} - successful track'.format(l=lbl),
                              metadata=succ_track['meta'])

resAnno.saveToFile('/Volumes/Seagate_Backup_Plus_Drive 1/datasets/volleyball/graz-arbesback_5/dsst.csv')