[GTR] 機械学習で動画の物体追跡 [Python]

2022年3月31日木曜日

Artificial Intelligence

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

アイキャッチ

Global Tracking Transformers

概要

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

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

GTR_概要図

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

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

デモ(Colaboratory)

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

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

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

環境セットアップ

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

「ランタイムのタイプを変更」→「ハードウェアアクセラレータ」を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

AIで副業ならココから!

まずは無料会員登録

プロフィール

メーカーで研究開発を行う現役エンジニア
組み込み機器開発や機会学習モデル開発に従事しています

本ブログでは最新AI技術を中心にソースコード付きでご紹介します


Twitter

カテゴリ

このブログを検索

ブログ アーカイブ

TeDokology