OpenCVによる物体追跡 - nobUnagaの日記

OpenCVには物体追跡アルゴリズムがcontrib-libとして提供されている。 Web上にたくさんサンプルが落ちているので参考にやってみる。

インストール

追跡アルゴリズムはcontrib提供なので別途インストールが必要。
VideoCaptureを使うにはOpenCVのコンパイル時にVideo-I/Oがアクティブでなければならない，などの注意があるらしいが今回は問題が発生しなかったので詳細は不明。

$ pip3 install opencv-python
$ pip3 install opencv-contlib-python

サンプルとメモ

's'キーを押すとROI(Region of interest)選択モードになる。これにはopenCVが提供するselectROI関数を利用する(便利な関数だなぁ)。fromCenterをFalseにしないとマウスの位置が矩形の中心になるので，左上角から矩形を描きたい場合にはFalseにしておく。
Webの例だとROIの選択は一回のみで，再度's'キーで止めてROIを選択しても適切に動作しない場合が多い。そこを修正しようとしたが，ROIを上書きして，Trackerをinitするだけでは適切にリセットされないようで，'s'キーを押すたびにTrackerをコンストラクトしている。実用的ではないけど実験には十分。
imshow()はwaitKeyで指定される秒数しか描画されない，描画を継続したい場合にはwaitKey(0)を指定する(引数無しでも良さそうだが，手元の環境ではうまく行かなかった)。また，waitKeyが無くimshow単体では適切に描画さえもされないので注意する。
destroyAllWindowsにウィンドウを個別に渡している例も見られたが，destroyAllWindowsは引数が無くても，全てのウィンドウを破棄してくれるみたいだ。
使用するアルゴリズムの基本的な選択基準は，精度最優先=CSRT, 速度最優先=MOSSE。KCFはBoostingとMILの改良版なので，Boosting, MILは基本的には使わない。起動の予測が容易な場合にはMedianFlow。TLDは他のアルゴリズムと異なる特性を持つようだが，現在のOpenCVの実装にはバグがあるのか，適切に動かない(らしい)。試した印象ではCSRTと他で大きく精度が違う印象。

import argparse
import imutils
import cv2


def main(args):
    # Generate tracker
    TRACKERS = {
        "csrt": cv2.TrackerCSRT_create,
        "kcf": cv2.TrackerKCF_create,
        "boosting": cv2.TrackerBoosting_create,
        "mil": cv2.TrackerMIL_create,
        "tld": cv2.TrackerTLD_create,
        "median": cv2.TrackerMedianFlow_create,
        "mosse": cv2.TrackerMOSSE_create
    }
    tracker = TRACKERS[args.tracker]()

    # Load video stream
    cap = cv2.VideoCapture(args.video)
    assert cap.isOpened(), "Given video can't be opened."

    window = 'tracking test'
    target = None
    while True:
        is_read, frame = cap.read()
        if is_read:
            frame = imutils.resize(frame, width=500)  # ndarray, (H,W,C)
            H, W, _ = frame.shape
            if target is not None:
                is_success, bbox = tracker.update(frame)
                if is_success:
                    x, y, w, h = [int(d) for d in bbox]
                    cv2.rectangle(frame, (x, y), (x+w, y+h), (0, 255, 0), 2)
                else:
                    print('Tracking is failed.')
            cv2.imshow(window, frame)
            k = cv2.waitKey(10) & 0xFF
            if k == ord('q'):
                break
            elif k == ord('s'):  # select target
                # After selection, ENTER or SPACE. Re-select by ESC
                del(tracker)
                tracker = TRACKERS[args.tracker]()
                target = cv2.selectROI(window, frame, fromCenter=False,
                                       showCrosshair=True)
                tracker.init(frame, target)
        else:
            break

    cap.release()
    cv2.destroyAllWindows()


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-v', '--video', type=str, help='video file path')
    parser.add_argument('-t', '--tracker', type=str, default='csrt',
                        help='tracker type')
    args = parser.parse_args()

    main(args)