GluonCVでObject Tracking(SiamRPN models)

環境

Windows10 Pro
GPUなし
Python 3.8.2

バージョンの確認(pip freeze)

インストールが必要なのは「mxnet」と「gluoncv」と「opencv-python」のみ。

pip install mxnet
pip install gluoncv
pip install opencv-python

その他は勝手についてくる。

certifi==2020.4.5.1
chardet==3.0.4
cycler==0.10.0
gluoncv==0.7.0
graphviz==0.8.4
idna==2.6
kiwisolver==1.2.0
matplotlib==3.2.1
mxnet==1.6.0
numpy==1.18.4
opencv-python==4.2.0.34
Pillow==7.1.2
portalocker==1.7.0
pyparsing==2.4.7
python-dateutil==2.8.1
pywin32==227
requests==2.18.4
scipy==1.4.1
six==1.14.0
tqdm==4.46.0
urllib3==1.22

実行ファイル

import os
import mxnet as mx
from gluoncv import model_zoo
from gluoncv.model_zoo.siamrpn.siamrpn_tracker import SiamRPNTracker

import cv2 

# mp4データを読み込む
video_frames = []
video_path = 'test.mp4'
cap = cv2.VideoCapture(video_path)
while(True):
    ret, img = cap.read()
    if not ret:
        break
    video_frames.append(img)

# モデルを取得する
net = model_zoo.get_model('siamrpn_alexnet_v2_otb15', pretrained=True)
tracker = SiamRPNTracker(net)

#最初のポジション
#(左上X座標、左上Y座標、横の大きさ、縦の大きさ)
gt_bbox = [212, 228, 119, 202]

if not os.path.exists('./predictions'):
    os.makedirs('./predictions')

for ind, frame in enumerate(video_frames):
    if ind == 0:
        tracker.init(frame, gt_bbox, ctx=mx.cpu())
        pred_bbox = gt_bbox
    else:
        outputs = tracker.track(frame, ctx=mx.cpu())
        pred_bbox = outputs['bbox']
    pred_bbox = list(map(int, pred_bbox))
    cv2.rectangle(frame, (pred_bbox[0], pred_bbox[1]),
                      (pred_bbox[0]+pred_bbox[2], pred_bbox[1]+pred_bbox[3]),
                      (0, 255, 255), 3)
    cv2.imwrite(os.path.join('./predictions', '%04d.jpg'%(ind+1)), frame)

結果

f:id:touch-sp:20200513213938g:plain

ffmpeg -i predictions/%04d.jpg -f gif output.gif

補足

物体の最初の位置を調べるために「BBoX-Label-Tool」を使用した。
Python3系で使用するために以下のサイトを参考にさせて頂いた。
BBoX-Label-ToolをPython3で使用する - ロボット、電子工作、IoT、AIなどの開発記録

2020年10月25日追記

物体の最初の位置を調べるためにPythonスクリプトを書いた。
touch-sp.hatenablog.com