はじめに
前回の記事で紹介したcolabスクリプトでは物体検知の方のデモしか実行してくれません。
今回はGoogleドライブに環境構築をして、骨格点検出をしてみます！
環境構築
環境構築は前回の記事で紹介したcolabスクリプトとほとんど変わりません。
ノートブックの設定は「GPU」にしてください。設定方法についてはこちら。
Googleドライブのマウント
Googleドライブ上に環境を構築するので必須です。
from google.colab import drive
drive.mount('/content/drive')
CenterNetをCloneする
CenterNetレポジトリをクローンします。
ここではGoogleドライブ直下にあるColab Notebooksというディレクトリを指定しています。
%cd /content/drive/My Drive/Colab Notebooks/ !git clone -q --depth 1 'https://github.com/xingyizhou/CenterNet.git'
Modelをダウンロードする
CenterNetの作者が提供しているModelをダウンロードします。
ちなみに2つダウンロードしていますが、骨格点検出で利用するのは下のmulti_pose_dla_3x.pthの方です。
CenterNetのGitページに飛ぶとほかにもいろいろとモデルを提供しているのでよかったら観覧してください。
%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'
ここまでのコードについては一度実行した後に、Colabのランタイムを再起動、リセットしても再度実行する必要はありません。
依存パッケージ等のインストール
これ以降の部分についてはColabのランタイムが再起動、リセットされた場合にはもう一度実行する必要があります。
%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
DCNv2の修理及びビルド
DCNv2という変わったCNNについてはディレクトリ毎にビルドしなおす必要があるらしいです。そしてCenterNetにもともとあるものは壊れているみたいです。
%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を通す
CenterNetのパッケージを読み込むためにパスを通しておきます。
import sys
from os.path import join
sys.path.insert(0, join('/content/drive/My Drive/Colab Notebooks/CenterNet', 'src/lib'))
sys.path.append(join('/content/drive/My Drive/Colab Notebooks/CenterNet', 'src'))
sys.path.append(join('/content/drive/My Drive/Colab Notebooks/CenterNet', 'src/lib/models/networks/DCNv2'))
骨格点検出
ここまでこれば後は簡単です！
次のコードを実行することでretsには辞書型のデータが入ります（1が人のキーとなっています。他は調べてません）。
%cd /content/drive/My Drive/Colab Notebooks/CenterNet/src
from detectors.detector_factory import detector_factory
from opts import opts
MODEL_PATH = '../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)
img = '/content/drive/My Drive/Colab Notebooks/CenterNet/images/sample.jpg'
rets = detector.run(img)['results']
キーの要素は次のような二重配列になっています。
[[x1,y1,x2,y2,score,x3,y3,...,xn,yn][x1,y1,x2,y2,score,x3,y3,...,xn,yn][x1,y1,x2,y2,score,x3,y3,...,xn,yn][x1,y1,x2,y2,score,x3,y3,...,xn,yn]]
ここでx1,y1,x2,y2は物体を囲う最小の矩形、scoreは類似度です。
scoreの後には骨格点の座標が連なっています。
ということで、OpenCVをつかって簡単に骨格点を出力します。
import cv2
from google.colab.patches import cv2_imshow
img = cv2.imread(img, cv2.IMREAD_UNCHANGED)
ret = rets.get(1)
for rret in ret:
if rret[4] > 0.7:
print(rret)
cv2.rectangle(img, (int(rret[0]), int(rret[1])), (int(rret[2]), int(rret[3])), (0, 255, 0, 255), 2)
num = int((len(rret)-5)/2+1)
for i in range(num):
cv2.circle(img, (int(rret[2*(i-1)+5]), int(rret[2*(i-1)+6])), 5, (255, 255, 255), thickness=-1)
cv2_imshow(img)
実行してみた
ということでいつも通りぱくたそさんから適当にとってきた画像を入力画像として与えてみました。
曲げた足の足首まできれいに骨格点が取れました！
最後に
こちらで動画に対して骨格点検出を行っています。
CenterNetで次やることは自前データの学習かな？
Colabで簡単にできそうだったら、紹介するかもしれません。
コメント