opencv/object_tracking_vittrack

1# This file is part of OpenCV Zoo project.2# It is subject to the license terms in the LICENSE file found in the same directory.3 4import numpy as np5import cv2 as cv6 7class VitTrack:8    def __init__(self, model_path, backend_id=0, target_id=0):9        self.model_path = model_path10        self.backend_id = backend_id11        self.target_id = target_id12 13        self.params = cv.TrackerVit_Params()14        self.params.net = self.model_path15        self.params.backend = self.backend_id16        self.params.target = self.target_id17 18        self.model = cv.TrackerVit_create(self.params)19 20    @property21    def name(self):22        return self.__class__.__name__23 24    def setBackendAndTarget(self, backend_id, target_id):25        self.backend_id = backend_id26        self.target_id = target_id27 28        self.params.backend = self.backend_id29        self.params.target = self.target_id30 31        self.model = cv.TrackerVit_create(self.params)32 33    def init(self, image, roi):34        self.model.init(image, roi)35 36    def infer(self, image):37        is_located, bbox = self.model.update(image)38        score = self.model.getTrackingScore()39        return is_located, bbox, score40