はじめに
最近はもっぱら人物追跡系を触っているのですが、今回はSORTの検出器をCenterNetに変えたものをColabで試してみました。
SORTは検出器から得たデータに追跡IDを割り当てるためのアルゴリズムで、複数の技術を組み合わせたものです。
論文では検出器にFasterCNNを利用していたと思いますが、今回はCenterNetを使ってみました。
環境構築
まずはノートブックの設定を「GPU」に変更します。（わからない場合はこちらを参照してください。）
次にGoogleDriveのマウントをします。
from google.colab import drive
drive.mount('/content/drive')
以後、CenterNetとSORTの環境構築を順に行います。
CenterNet環境構築
CenterNetについては以前の記事でも触れているので、コードのみで済ませます。
Gitリポジトリのクローン
%cd /content/drive/My Drive/Colab Notebooks/ !git clone -q --depth 1 'https://github.com/xingyizhou/CenterNet.git'
Modelのダウンロード
%cd /content/drive/My Drive/Colab Notebooks/CenterNet/models !pip install -q gdown !gdown 'https://drive.google.com/uc?id=1pl_-ael8wERdUREEnaIfqOV_VF2bEVRT' !gdown 'https://drive.google.com/uc?id=1PO1Ax_GDtjiemEmDVD7oPWwqQkUu28PI'
依存パッケージのインストール
CenterNetの環境構築についてはこれ以降、ランタイムのリセット毎に実行してください。
%cd /content/drive/My Drive/Colab Notebooks/CenterNet !pip install -U torch==1.4 torchvision==0.5 -f https://download.pytorch.org/whl/cu101/torch_stable.html !pip install -q -r requirements.txt
DCNの修復
%cd /content/drive/My Drive/Colab Notebooks/CenterNet/src/lib/models/networks !rm -rf DCNv2 !git clone https://github.com/CharlesShang/DCNv2.git %cd /content/drive/My Drive/Colab Notebooks/CenterNet/src/lib/models/networks/DCNv2 !chmod u+x make.sh !./make.sh
Pathを通す
import sys
sys.path.insert(0, '/content/drive/My Drive/Colab Notebooks/CenterNet/src/lib')
sys.path.append('/content/drive/My Drive/Colab Notebooks/CenterNet/src')
sys.path.append('/content/drive/My Drive/Colab Notebooks/CenterNet/src/lib/models/networks/DCNv2')
以上でCenterNetの準備は完了です！
SORT環境構築
SORTの環境構築は簡単です。
Gitリポジトリのクローン
著者のGitリポジトリからダウンロードします。
%cd /content/drive/My Drive/Colab Notebooks/ !git clone -q --depth 1 'https://github.com/abewley/sort.git'
依存パッケージのインストール
SORTの環境構築ではこれ以降をランタイムのリセット毎に実行してください。
%cd /content/drive/My Drive/Colab Notebooks/sort/ !pip install -r requirements.txt
Pathを通す
sys.path.append('/content/drive/My Drive/Colab Notebooks/sort/')'
以上でSORTの準備は完了です！
実際に動かしてみる！
では実際に動かしてみます。
モデルのパスなどのCenterNetの設定
SORTを動かすには矩形情報と類似度が必要なので、それらが取得できるモデルであればどれでも動くと思います。
%cd /content/drive/My Drive/Colab Notebooks/
from detectors.detector_factory import detector_factory
from opts import opts
MODEL_PATH = '/content/drive/My Drive/Colab Notebooks/CenterNet/models/multi_pose_dla_3x.pth' #モデルのパス
TASK = 'multi_pose' #骨格点検出
opt = opts().init('{} --load_model {}'.format(TASK, MODEL_PATH).split(' '))
detector = detector_factory[opt.task](opt)
動画の読み込みと結果の出力
リポジトリのsort.pyを参考にしています。
7行目の動画ファイルと17行目の出力ファイルの部分を適宜変更してください。
私の場合は勝手に類似度で間引きをしていますが、これが不要な場合は5行目の「tsd = 0.4」を「tsd = 0.0」に変更するか、適宜コードを変更してください。
import cv2
import numpy as np
from sort import *
tsd = 0.4 #類似度の閾値（勝手に設定）
video_path = '/content/drive/My Drive/Colab Notebooks/video/demo.mp4' #動画ファイルのパス
cap = cv2.VideoCapture(video_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('/content/drive/My Drive/Colab Notebooks/out_demo.mp4', fourcc, frame_rate, (width, height))
mot_tracker = Sort() #SORTインスタンス生成
print('start')
#ループ処理
while True:
gg, img = cap.read() #フレーム取り出し
if gg: #フレームが取り出せた場合
ret = detector.run(img)['results'] #CenterNet実行
ret_sort = np.empty((0, 5)) #Sortに渡すデータ
for rret in ret.get(1): #get(1)で人物
if rret[4] > tsd: #類似度以上
ret_sort = np.append(ret_sort, np.array([rret[0:5]]), axis=0) #Sortに渡す
track_bbs_ids = mot_tracker.update(ret_sort) #Sort実行
for ret_track in track_bbs_ids:
cv2.rectangle(img, (int(ret_track[0]), int(ret_track[1])), (int(ret_track[2]), int(ret_track[3])), (0, 255, 0, 255), 2) #矩形描画
cv2.putText(img, str(ret_track[4]), (int(ret_track[0])+30, int(ret_track[1])+60), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 0, 255), thickness=2) #トラッキングID
out.write(img) #動画書き込み
else:
break
out.release() #動画完了
print('finish')
出力結果
出力結果をYouTubeに上げました。
最後に
SORTが思っていた以上に使いやすいので、いいなーと思いましたが精度としてはやはりdeepSORTの方がいいのかもしれないですね。
そのうちdeepSORTについても記事書きます。
コメント