OpenCV Video Module GOTURN Tracker Sample

goturn.py

import os
import sys
import cv2 as cv
import numpy as np

def main():
    # キャプチャを開く
    directory = os.path.dirname(__file__)
    #capture = cv.VideoCapture(os.path.join(directory, "test.mp4")) # 画像ファイル
    capture = cv.VideoCapture(0) # カメラ
    if not capture.isOpened():
        exit()
    
    # トラッカーを作成する
    # https://github.com/opencv/opencv_extra/tree/c4219d5eb3105ed8e634278fad312a1a8d2c182d/testdata/tracking
    params = cv.TrackerGOTURN_Params()
    params.modelBin = os.path.join(directory, "goturn.caffemodel")
    params.modelTxt = os.path.join(directory, "goturn.prototxt")
    tracker = cv.TrackerGOTURN_create(params)

    # オブジェクトを選択する
    result, image = capture.read()
    if result is False:
        sys.exit()
    box = cv.selectROI(image, False)

    # トラッカーを初期化する
    tracker.init(image, box)

    while True:
        # フレームをキャプチャして画像を読み込む
        result, image = capture.read()
        if result is False:
            cv.waitKey(0)
            break

        # トラッカーを更新する
        result, box = tracker.update(image)

        # 追跡したバウンディングボックスを描画する
        if result is True:
            color = (0, 0, 255)
            thickness = 2
            cv.rectangle(image, box, color, thickness, cv.LINE_AA)
        else:
            # TODO: Re-Iniit Tracker
            break

        # 画像を表示する
        cv.imshow("tracking", image)
        key = cv.waitKey(10)
        if key == ord('q'):
            break
    
    cv.destroyAllWindows()

if __name__=='__main__':
    main()