はじめに
前回の記事で画像に映る人の骨格点を検出しました。
今回は、それを動画でもできるようにしたものです。コード等の説明は前回とほとんど同じなので、省略します。
Colabでの環境構築
環境構築については前回も書いているのでコンパクトに書きます。
ノートブックの設定は「GPU」です。
何をしているのか気になる方は前回の記事を参照してください。
ドライブのマウント
from google.colab import drive
drive.mount('/content/drive')
1度だけ実行
%cd /content/drive/My Drive/Colab Notebooks/ !git clone -q --depth 1 'https://github.com/xingyizhou/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'
ランタイムのリセット毎に実行
%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
%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
%cd /content/drive/My Drive/Colab Notebooks/CenterNet/src
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'))
ソースコード
ここからが動画で骨格点を検出するために必要なコードです。
骨格点検出の指定
前回と同様に、モデルパス、タスクを指定します。
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)
骨格点検出を行う
動画ファイルをOpenCVで読み込んで、フレームごとに骨格点検出を行うだけです。
ちなみに骨格点は次の順番だと思われます。
1:頭,2:顎,3:胸骨,4:右肩,5:左肩,6:右肘,7:左肘,8:右腰,9:左腰,10:右手首,11:左手首,12:右膝,13:左膝,14:右外果,15:左外果,16:右足尖,17:左足尖
import cv2
file_path = '/content/drive/My Drive/Colab Notebooks/CenterNet/images/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: #取り出せた場合
rets = detector.run(img)['results'] #骨格点検出
ret = rets.get(1) #人のキーを指定
for rret in ret: #人数分繰り返し
if rret[4] > 0.5: #閾値を0.5でセット
#人を矩形で囲う
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)
out.write(img)
else: #フレームが取り出せない（終点）
break
out.release()
実行結果
フリー動画素材サイトから取ってきた動画を使って実行してみました。
元動画と出力された動画を隣り合わせにしたものを置いておきます。
人が重なると骨格点についてはかなりチラつくように感じますね。
最後に
CenterNet2があるとかなんとか。いまいち調べてもよくわからず。
コメント