Object Trackingを行うために物体の最初の位置を取得する

Object Tracking

2021年10月18日追記
OpenCVの「selectROI」を使うと簡潔に書ける。

import cv2
import sys

cap = cv2.VideoCapture(sys.argv[1])
ret, img = cap.read()
cap.release()

source_window = "draw_rectangle"
cv2.namedWindow(source_window)

rect = cv2.selectROI(source_window, img, False, False)

print('xmin:%d'%rect[0])
print('ymin:%d'%rect[1])
print('width:%d'%rect[2])
print('height:%d'%rect[3])

cv2.destroyAllWindows()




ここから元記事です。


以前GluonCVを使ったObject Trackingの記事を書いた。
touch-sp.hatenablog.com
Trackingしたい物体の最初の位置を知る必要がある。
前は「BBoX-Label-Tool」を使用したが、使用する前にはビデオファイルから画像を抽出する必要がある。
またPython3系で「BBoX-Label-Tool」を使用するためには少しスクリプトを書き変える必要がある。

今回書いたPythonスクリプト

import cv2
import sys

drawing = False
complete_region = False
ix,iy,width,height = -1,-1,0,0
 
# マウスコールバック関数
def my_mouse_callback(event,x,y,flags,param):
    global ix,iy,width,height,drawing,complete_region
 
    if event == cv2.EVENT_MOUSEMOVE:      # マウスが動いた時
        if(drawing == True):
            width = x - ix
            height = y - iy
 
    elif event == cv2.EVENT_LBUTTONDOWN:  # マウス左押された時
        drawing = True
 
        ix = x
        iy = y
        width = 0
        height = 0
 
    elif event == cv2.EVENT_LBUTTONUP:    # マウス左離された時
        drawing = False
        complete_region = True
 
        if(width < 0):
            ix += width
            width *= -1
        if(height < 0):
           iy += height
           height *= -1

cap = cv2.VideoCapture(sys.argv[1])
ret, img = cap.read()       #最初のフレームを取得
temp = img.copy()           # 画像コピー

source_window = "draw_rectangle"
cv2.namedWindow(source_window)
cv2.setMouseCallback(source_window, my_mouse_callback)

while(1):
    cv2.imshow(source_window,temp)

    if(drawing):            # 左クリック押されてたら
        temp = img.copy()   # 画像コピー
        cv2.rectangle(temp,(ix,iy),(ix + width, iy+ height),(0,255,0),2)  # 矩形を描画

    if(complete_region):    # 矩形の選択が終了したら
        complete_region = False

        print('xmin:%d'%ix)
        print('ymin:%d'%iy)
        print('width:%d'%width)
        print('height:%d'%height)

    # キー操作
    k = cv2.waitKey(1) & 0xFF
    if k == 27:             # esc押されたら終了
        break

cap.release()
cv2.destroyAllWindows()

あとで使用するためにこのスクリプトは「get_position.py」として保存しておく。

実際にObject Trackingを行う

まずはサンプルビデオをダウンロード。

from gluoncv import utils
video_path = 'https://raw.githubusercontent.com/dmlc/web-data/master/gluoncv/tracking/Coke.mp4'
utils.download(video_path)

これだけで「Coke.mp4」というファイルがダウンロードされる。
次に先ほど作った「get_position.py」を使う。

python get_position.py Coke.mp4


缶を選ぶとコンソール上に結果が出力される。

xmin:297
ymin:158
width:44
height:85

この結果を用いて以下を実行する。

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 = 'Coke.mp4'
cap = cv2.VideoCapture(video_path)
fps = int(cap.get(cv2.CAP_PROP_FPS))
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

fourcc = cv2.VideoWriter_fourcc('m','p','4','v')
out = cv2.VideoWriter('output.mp4',fourcc, fps, (width,height))

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 = [297, 158, 44, 85]

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)
    out.write(frame)

cap.release()
out.release()
cv2.destroyAllWindows()

「output.mp4」というファイルが作成される。
mp4をgifに変換したものをのせておく。

はまった点

mp4で保存したければ

fourcc = cv2.VideoWriter_fourcc('m','p','4','v')
out = cv2.VideoWriter('output.mp4',fourcc, fps, (width,height))

aviで保存したければ

fourcc = cv2.VideoWriter_fourcc('M', 'J', 'P', 'G')
out = cv2.VideoWriter('output.avi',fourcc, fps, (width,height))