
import os
import numpy as np
import time

import cv2
import operator

from obj_tracker import Sort

scale = 2
total_time = 0.0
total_frames = 0

pre_det = []
tracks = np.empty((0, 5))

# create sort tracker
sort = Sort(5, 0, 0.2)

# do track and render on image
def track_and_render(detections, img):
    global total_frames
    global total_time
    global pre_det
    global tracks

    total_frames += 1
    start_time = time.time()

    # call the tracker algo
    if not operator.eq(pre_det, detections):
        tracks = sort.update(np.array(detections))
        pre_det = detections
    else:
        # if skip frame in detection process
        tracks = sort.update()
    
    # get the fps, calculate avg speed in 100 frames
    cycle_time = time.time() - start_time
    total_time += cycle_time
    if total_frames % 100 == 0:
        print("Total Tracking took: %.3f seconds for %d frames or %.1f FPS" % (total_time, total_frames, total_frames / total_time))
        total_time = 0.0
        total_frames = 0
    
    # render on frame image
    for track in tracks: 
        # [left, top, right, bottom, id]
        pt1 = (int(track[0]), int(track[1]))
        pt2 = (int(track[2]), int(track[3]))
        cv2.rectangle(img, pt1, pt2, (0, 255, 0), 2)

        cv2.putText(img, str(int(track[4])),
                    (pt1[0], pt1[1] - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.7,
                    (0, 255, 255), 1) 

    # return frame image to caller
    return img


# main
if __name__ == '__main__':
    
    # detection data, generated by object-detector algo model, including frame image
    path = 'mot_benchmark\\MY_TEST\\www3'

    # read test data from disk in a loop one frame by frame and send it to sort tracker
    for i in range(400000):
        # read frame image
        bg = cv2.imread(path + '\\' + str(i % 800 + 1) + '.jpg')
        # rescale
        bg = cv2.resize(bg, (int(1280 / scale), int(720 / scale)))
        
        # read detection data in current frame
        dets = []
        with open(path + '\\' + str(i % 800 + 1) + '.txt') as txt:
            lines = txt.readlines()
            for line in lines:
                items = line.split(' ')
                left = float(items[1])
                top = float(items[2])
                right = float(items[3])
                bottom = float(items[4]) 
                
                dets.append((left / scale, top / scale, right / scale, bottom / scale, int(items[0])))
        
        # send to the tracker and draw the result on background image
        if len(dets) > 0:
            bg = track_and_render(dets, bg)
            cv2.imshow('tracker', bg)
            cv2.waitKey(40)
    
    cv2.destroyAllWindows()
