はじめに
CenterTrackの使い方が分かるサイトが全くなかったので、環境構築から残しておきます。
CenterNetについてはこちらのカテゴリで触れてきましたが、今回はCenterNetのモデルベースの物体追跡です。
ちなみにColab由来のエラーなのかはわかりませんが、プログラムの1か所を修正しないといけませんでした（2020/10/08現在）。
この修正と同じものがpullリクエストにあったので、近いうちに修正が入ると思います。
環境構築
Colabratoryでの環境構築です（動作確認：2020/10/08）。
Googleドライブのマウント
Googleドライブ上に環境を作ることをお勧めします。
from google.colab import drive
drive.mount('/content/drive')
CenterTrackのcloneとmodelのダウンロード（1度だけ実行）
CenterTrackのレポジトリをクローンします。
ここではGoogleDrive直下のColab Notebooksを指定しています。
CenterTrackフォルダの中にmodelsを作成し、各種モデルをダウンロードしています（このコードは必須ではありません、必要なものを選んでダウンロードすることをお勧めします）。
GoogleDrive上にデータが保存されるはずなのでこの作業は1度だけ行えば大丈夫です。
%cd /content/drive/My Drive/Colab Notebooks/ !git clone -q --depth 1 'https://github.com/xingyizhou/CenterTrack.git' %cd /content/drive/My Drive/Colab Notebooks/CenterTrack/ %mkdir models %cd /content/drive/My Drive/Colab Notebooks/CenterTrack/models !pip install -q gdown !gdown 'https://drive.google.com/uc?id=1WXBlzHsxHQTELvusJSgEWw_wydC6u7XB' !gdown 'https://drive.google.com/uc?id=1e8zR1m1QMJne-Tjp-2iY_o81hn2CiQRt' !gdown 'https://drive.google.com/uc?id=1tJCEJmdtYIh8VuN8CClGNws3YO7QGd40' !gdown 'https://drive.google.com/uc?id=1H0YvFYCOIZ06EzAkC2NxECNQGXxK27hH'
プログラムの修正
Colabで次の作業に進む前に、ダウンロードしたgitリポジトリ内のプログラムの修正をします。
前述のとおり、現在は修正を自分で行わないとエラーが出る状態です。
この作業を後から行うと結局以下の作業がすべてやり直しなので、念のため確認しておくことを推奨します。
opt = self.parse()
「CenterTrack/src/lib/opts.py」の397行目が上記のようになっている場合には次のように修正を加えてください。
opt = self.parse(args)
ランタイムのリセット毎に実行が必要
各種パッケージのインストールと、DCNv2のmake実行、src/libのpathを通しています。
この作業はランタイムがリセットされるたびに行わなくてはいけません。
%cd /content/drive/My Drive/Colab Notebooks/CenterTrack/
!pip install -U torch==1.4 torchvision==0.5 -f https://download.pytorch.org/whl/cu101/torch_stable.html
!pip install -U pytest==3.8 motmetrics==1.1.3 folium==0.2.1
!pip install -q -r requirements.txt
%cd /content/drive/My Drive/Colab Notebooks/CenterTrack/src/lib/model/networks
!rm -rf DCNv2
!git clone https://github.com/CharlesShang/DCNv2.git
%cd /content/drive/My Drive/Colab Notebooks/CenterTrack/src/lib/model/networks/DCNv2
!chmod u+x make.sh
!./make.sh
%cd /content/drive/My Drive/Colab Notebooks/CenterTrack/src
import sys
from os.path import join
sys.path.insert(0, join('/content/drive/My Drive/Colab Notebooks/CenterTrack', 'src/lib'))
sys.path.append(join('/content/drive/My Drive/Colab Notebooks/CenterTrack', 'src'))
sys.path.append(join('/content/drive/My Drive/Colab Notebooks/CenterTrack', 'src/lib/model/networks/DCNv2'))
これで環境構築は終了です！
CenterNetとほとんど同様でした。
人物追跡の実装
とりあえずモデルの準備です。モデルパスとタスクを指定します。
ここまでは人物追跡以外も同じはず。
%cd /content/drive/My Drive/Colab Notebooks/CenterTrack/src
from detector import Detector
from opts import opts
MODEL_PATH = '../models/coco_tracking.pth' #モデルのパス
TASK = 'tracking' #2Dトラッキング
opt = opts().init('{} --load_model {}'.format(TASK, MODEL_PATH).split(' '))
detector = Detector(opt)
人物追跡ソースコード
後はモデルに通すことで自動で結果が得られるのでそのデータをもとに矩形で囲ったりトラッキングIDを出力するだけです。
今回はOpenCVを使って実装しました。
import cv2
file_path = '/content/drive/My Drive/Colab Notebooks/CenterTrack/videos/test.mp4' #動画ファイルのパス
cap = cv2.VideoCapture(file_path) #動画の読み込み
#出力用設定
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) #幅
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) #高
frame_rate = int(cap.get(cv2.CAP_PROP_FPS)) #FPS
fourcc = cv2.VideoWriter_fourcc('m','p','4','v') #mp4出力
out = cv2.VideoWriter('out.mp4', fourcc, frame_rate, (width, height))
while True:
gg, img = cap.read() #フレームの取り出し
if gg: #取り出せた場合
ret = detector.run(img)['results'] #データの取得
#物体毎に処理
for rret in ret:
if rret.get('class') == 1: #人
#描画
cv2.rectangle(img, (int(rret.get('bbox')[0]), int(rret.get('bbox')[1])), (int(rret.get('bbox')[2]), int(rret.get('bbox')[3])), (0, 255, 0, 255), 2) #人物矩形
cv2.putText(img, str(rret.get('tracking_id')), (int(rret.get('bbox')[0])+30, int(rret.get('bbox')[1])+60), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 0, 255), thickness=2) #トラッキングID
out.write(img) #書き込み
else: #フレームが取り出せない（終点）
break
out.release()
動画を読み込んで、srcのフォルダに結果の動画を出力しています。
実行結果
こちらのサイトから適当に取ってきた動画で試してみた結果です。
最後に
ところどころトラッキングが間違っているところがありますが、元動画を白黒で用意したのが悪いかもです。徹夜して適当にやってたのでごめんなさい。そろそろ寝ます（朝9時）。
コメント