DNN: Detectron2 Keypoint 検出モデル

Detectron2でキーポイント検出モデル(keypoint detection model)の推論を試したメモ。以下のDetectron2 Beginner’s Tutorialを和訳して説明を加えたもの。TutorialがGoogle Colabというクラウドサービスを使ったJupyterノートブックなのでローカルマシンで動くようにPythonスクリプトを少し変更している。

Detectron2 Beginner’s Tutorial

過去記事のMask RCNNとスクリプト自体はほとんど同じ。違うのはモデルの設定とウェイトファイルのみ。ほとんど同じスクリプトでいろいろなモデルを試せるのはDetectron2の良いところ。

以下のスクリプトを~/src/detectron2/myprog/keypoint_inference.pyとして保存する。

# kepoint_inference.py
from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog
import cv2

# 画像の読み込み
im = cv2.imread("./happy_mini.jpg")

# ネットワークの設定
cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-Keypoints/keypoint_rcnn_R_50_FPN_3x.yaml"))
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7  # set threshold for this model
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-Keypoints/keypoint_rcnn_R_50_FPN_3x.yaml")

# 推論
predictor = DefaultPredictor(cfg)
outputs   = predictor(im)

# 結果の表示
v = Visualizer(im[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2)
v = v.draw_instance_predictions(outputs["instances"].to("cpu"))
cv2.imshow('Results', v.get_image()[:, :, ::-1])
cv2.waitKey(0) 
cv2.destroyAllWindows()

サンプル画像のダウンロード