__init__.py

from __future__ import absolute_import

import concurrent.futures
import cv2 as cv
import gc
import numpy as np
import time
import torch
from typing import Union

from ..utils.metrics import iou


class Tracker(object):

    def __init__(self, name, is_deterministic=False):
        self.name = name
        self.is_deterministic = is_deterministic
        if self.is_using_cuda:
            print('Detect the CUDA devide')
            self._timer_start = torch.cuda.Event(enable_timing=True)
            self._timer_stop = torch.cuda.Event(enable_timing=True)
        self._timestamp = None

    def init(self, image, box):
        raise NotImplementedError()

    def update(self, image):
        raise NotImplementedError()

    @property
    def is_using_cuda(self):
        self.cuda_num = torch.cuda.device_count()
        if self.cuda_num == 0:
            return False
        else:
            return True

    def _start_timing(self) -> Union[float, None]:
        if self.is_using_cuda:
            self._timer_start.record()
            timestamp = None
        else:
            timestamp = time.time()
            self._timestamp = timestamp
        return timestamp

    def _stop_timing(self) -> float:
        if self.is_using_cuda:
            self._timer_stop.record()
            torch.cuda.synchronize()
            # cuda event record return duration in milliseconds.
            duration = self._timer_start.elapsed_time(
                self._timer_stop
            )
            duration /= 1000.0
        else:
            duration = time.time() - self._timestamp
        return duration

    def track(self, seq_name, img_files, anno, restart_flag, visualize, seq_result_dir, save_img, method):
        frame_num = len(img_files)
        box = anno[0, :]  # the information of the first frame
        boxes = np.zeros((frame_num, 4))  # save the tracking result
        boxes[0] = box
        times = np.zeros(frame_num)  # save time

        fail_count = 0  # fail_count records the failures in R-OPE mechanism

        init_positions = []  # save the restart locations
        if visualize:
            display_name = 'Display: ' + seq_name
            cv.namedWindow(display_name, cv.WINDOW_NORMAL | cv.WINDOW_KEEPRATIO)
            cv.resizeWindow(display_name, 960, 720)

        # TODO: improve the imread speed
        with concurrent.futures.ProcessPoolExecutor() as executor:
            executor.map(cv.imread, img_files)

        for f, img_file in enumerate(img_files):

            image = cv.imread(img_file)
            height = image.shape[0]
            width = image.shape[1]
            img_resolution = (width, height)

            # start_time = time.time() 
            self._start_timing()
            if f == 0:
                self.init(image, box)
                times[f] = self._stop_timing()
            if fail_count >= 10 and method == 'restart' and f in restart_flag:
                # the tracker will be restarted when the cumulative number of failures reaches 10
                print('init again in %s' % f)
                init_positions.append(f)
                self.init(image, anno[f, :])
                fail_count = 0
            else:
                frame_box = self.update(image)
                frame_box = np.rint(frame_box)
                times[f] = self._stop_timing()

                current_gt = anno[f, :].reshape((1, 4))
                frame_box = np.array(frame_box)
                track_result = frame_box.reshape((1, 4))
                bound = img_resolution
                seq_iou = iou(current_gt, track_result, bound=bound)

                # check failures
                if method == 'restart' and (anno[f, :] != np.array([0, 0, 0, 0])).all():
                    if seq_iou < 0.5:
                        # failure occures in present frame
                        fail_count += 1
                    else:
                        # re-locate the target
                        fail_count = 0

                boxes[f, :] = frame_box

                if method == 'restart':
                    print(seq_name, self.name, ' Tracking %d/%d' % (f, frame_num - 1), 'time:%.2f' % times[f],
                          'fail count:', fail_count, frame_box)
                else:
                    print(seq_name, self.name, ' Tracking %d/%d' % (f, frame_num - 1), 'time:%.2f' % times[f],
                          frame_box)

                if save_img or visualize:
                    frame_disp = image.copy()
                    state = [int(s) for s in frame_box]
                    state[0] = 0 if state[0] < 0 else state[0]
                    state[1] = 0 if state[1] < 0 else state[1]
                    state[2] = width - state[0] if state[0] + state[2] > width else state[2]
                    state[3] = height - state[1] if state[1] + state[3] > height else state[3]
                    font_face = cv.FONT_HERSHEY_SIMPLEX
                    cv.putText(frame_disp, 'No.%06d' % (f), (50, 100), font_face, 0.8, (0, 255, 0), 2)
                    if (anno[f, :] != np.array([0, 0, 0, 0])).all():
                        cv.putText(frame_disp, 'seq iou: %2f' % (seq_iou), (50, 130), font_face, 0.8, (0, 255, 0), 2)

                    cv.rectangle(frame_disp, (state[0], state[1]), (state[2] + state[0], state[3] + state[1]),
                                 (0, 255, 0), 5)
                    gt = [int(s) for s in anno[f, :]]
                    cv.rectangle(frame_disp, (gt[0], gt[1]), (gt[2] + gt[0], gt[3] + gt[1]), (0, 0, 255), 5)

                    if visualize:
                        cv.imshow(display_name, frame_disp)
                    if save_img:
                        save_path = "{}/{:>06d}.jpg".format(seq_result_dir, f)
                        cv.imwrite(save_path, frame_disp)

                    del frame_disp
                    gc.collect()

                key = cv.waitKey(1)
                if key == ord('q'):
                    break

            del image
            gc.collect()

        if visualize:
            cv.destroyAllWindows()

        if method == None:
            return boxes, times
        elif method == 'restart':
            return boxes, times, init_positions