【CenterNet】動画に映る人の骨格点を検出する【Colab】

スポンサーリンク
mov-centernetCenterNet
スポンサーリンク

はじめに

前回の記事で画像に映る人の骨格点を検出しました。

今回は、それを動画でもできるようにしたものです。コード等の説明は前回とほとんど同じなので、省略します。

Colabでの環境構築

環境構築については前回も書いているのでコンパクトに書きます。

何をしているのか気になる方は前回の記事を参照してください。

ドライブのマウント

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で読み込んで、フレームごとに骨格点検出を行うだけです。

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があるとかなんとか。いまいち調べてもよくわからず。

コメント

タイトルとURLをコピーしました