【CenterTrack】posetrackingを試す【Colab】

はじめに

前回の続きのようなものです。前回表示した矩形情報に加えて骨格点情報を表示するようにしています。

環境構築部分は前回と全く同じなので省略し、実装部分も前回とほぼ同様なのでコードだけ載せておきます。

ソースコード

ポーズトラッキング用のモデルを既にダウンロードしてある場合は大丈夫ですが、まだダウンロードしていない場合には「/content/drive/My Drive/Colab Notebooks/CenterTrack/models」に追加してください。

%cd /content/drive/My Drive/Colab Notebooks/CenterTrack/src
from detector import Detector
from opts import opts

MODEL_PATH = '../models/coco_pose_tracking.pth' #モデルのパス
TASK = 'tracking,multi_pose' #2dposeトラッキング
opt = opts().init('{} --load_model {}'.format(TASK, MODEL_PATH).split(' '))
detector = Detector(opt)

骨格点情報は全部で17点で「hps」を指定することで取り出せます。

CenterNetに比べてキーとかの指定で取り出せるようになっていて嬉しいところです。

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: #人
        #矩形,ID,描画
        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

        #骨格点描画
        for i in range(17):
          cv2.circle(img, (int(rret.get('hps')[2*i]), int(rret.get('hps')[2*i+1])), 2, (255, 255, 255), thickness=-1)
    out.write(img) #書き込み
  else: #フレームが取り出せない（終点）
    break
out.release()

実行してみた

前回と同じ動画を使っています。

最後に

面白い使い方ないかな。