本記事では、Global Tracking Transformersと呼ばれる機械学習手法を用いて、動画の物体追跡を行う方法をご紹介します。

Global Tracking Transformers

概要

Global Tracking Transformers(GTR)は、動画中の複数のオブジェクトを追跡するTransformerベースの物体追跡技術(global multi-object tracking)です。
GTRは、動画から生成したフレーム画像からオブジェクトの特徴量をエンコードし、trajectory queriesを使用してそれぞれのオブジェクトを追跡します。 また、ペアワイズなどの組み合わせの関連付けを必要とせず、物体検出と物体追跡をEnd to Endにトレーニング可能にしています。

MOT17ベンチマークや、検出クラス数の多いTAOデータセットなどで従来技術に匹敵、あるいは上回るパフォーマンスを示しています。

詳細はこちらの論文をご参照ください。

本記事では、上記手法を用いて、任意の動画の物体追跡を実施する方法をご紹介します。

デモ(Colaboratory)

それでは、実際に動かしながら物体追跡を行っていきます。
ソースコードは本記事にも記載していますが、下記のGitHubでも取得可能です。
GitHub - Colaboratory demo

また、下記から直接Google Colaboratoryで開くこともできます。

なお、このデモはPythonで実装しています。
Pythonの実装に不安がある方、Pythonを使った機械学習について詳しく勉強したい方は、以下の書籍やオンライン講座などがおすすめです。

環境セットアップ

それではセットアップしていきます。 Colaboratoryを開いたら下記を設定しGPUを使用するようにしてください。

GTRの動作に必要なdetectron2をインストールします。

%cd /content

!git clone https://github.com/facebookresearch/detectron2.git
%cd /content/detectron2
!pip install -e .

論文発表元のGitHubからGlobal Tracking Transformersのコード取得します。

%cd /content

!git clone https://github.com/xingyizhou/GTR.git --recurse-submodules

ライブラリをインストールします。

%cd /content/GTR
!pip install -r requirements.txt
!pip install --upgrade gdown

以上で環境セットアップは完了です。

学習済みモデルのセットアップ

続いて、論文発表元より提供されている学習済みモデルをセットアップします。

import os

%cd /content/GTR
!mkdir pretrained

pretrained_ckpt = 'pretrained/GTR_TAO_DR2101.pth'

# https://drive.google.com/file/d/1TqkLpFZvOMY5HTTaAWz25RxtLHdzQ-CD/view
if not os.path.exists(pretrained_ckpt):
  !gdown --id 1TqkLpFZvOMY5HTTaAWz25RxtLHdzQ-CD \
          -O {pretrained_ckpt}

gdownを使用して、/content/GTR/pretrainedにモデルをダウンロードしています。

テスト動画のセットアップ

任意の動画をGoogle Colaboratoryにアップロードします。

本記事では、こちらのPexels様の動画を使用させて頂きます。

%cd /content/GTR
!rm -rf upload
!mkdir -p upload
%cd upload

from google.colab import files

uploaded = files.upload()
uploaded = list(uploaded.keys())
file_name = uploaded[0]

video_path = os.path.join("/content/GTR/upload", file_name)
print("upload file here:", video_path)

output_video_path = os.path.join('/content/GTR/output', 'result_'+file_name)

物体追跡

それでは物体追跡を実行します。

%%time

%cd /content/GTR

!python demo.py \
  --config-file configs/GTR_TAO_DR2101.yaml \
  --video-input {video_path} \
  --output {output_video_path} \
  --opts MODEL.WEIGHTS {pretrained_ckpt}

オブジェクトの追跡結果を追記した動画がoutputで指定したディレクトリに出力されます。
約8秒ほどの動画で3分ほど処理時間を要しました。検出クラスの多さからか、やや遅めといった印象です。

出力された動画を表示します。

from moviepy.editor import *
from moviepy.video.fx.resize import resize

clip = VideoFileClip(output_video_path)
clip = resize(clip, height=420)
clip.ipython_display()

出力結果は以下の通りです。
人物の追跡もさることながら、多くのオブジェクトが追跡されています。

まとめ

本記事では、Global Tracking Transformersを用いた物体追跡を行いました。
追跡可能なオブジェクトは非常に多いですが、処理速度が気になるところです。

また本記事では、機械学習を動かすことにフォーカスしてご紹介しました。
もう少し学術的に体系立てて学びたいという方には以下の書籍などがお勧めです。ぜひご一読下さい。

また動かせるだけから理解して応用できるエンジニアの足掛かりに下記のUdemyなどもお勧めです。

参考文献

1.  論文 - Global Tracking Transformers

2. GitHub - xingyizhou/GTR