1.はじめに

Detectron2とは、Facebook AIが開発した、PyTorchベースの物体検出のライブラリです。今回はDetectron2を用いた自作データの学習と題して、犬のキーポイント検出を行っていこうと思います。

作業環境としてはGoogle Colabを利用します。

2.キーポイント検出とは

キーポイント検出とは、画像に写っている人などの対象物の特徴点を検出することです。
よく使われるのが、人体の手足等のポイントを検出することで、その人の姿勢推定を行う例があります。

参考画像

人のキーポイント検出を行う学習モデルは多数ありますが、新たな物体のキーポイント検出をする際は、自作データで学習を行う必要があります。

今回は"犬"のキーポイント検出を行っていこうと思います。

3.方法

3-1. 学習データ作成

アノテーションにはCOCO Annotatorを用います。
詳細はこちらのページを参考にしました。

犬の頭、首、右前足、左前足、お尻、右後足、左後足の7点をキーポイントとして設定すると、こんな感じになります。

3-2. 学習

まずdetectron2をインストールします。
公式チュートリアルを参考にしましょう。


!python -m pip install pyyaml==5.1
import sys, os, distutils.core
!git clone 'https://github.com/facebookresearch/detectron2'
dist = distutils.core.run_setup("./detectron2/setup.py")
!python -m pip install {' '.join([f"'{x}'" for x in dist.install_requires])}
sys.path.insert(0, os.path.abspath('./detectron2'))

下記コードが実行ができれば、正常にインストールされています。

import torch, detectron2
!nvcc --version
TORCH_VERSION = ".".join(torch.__version__.split(".")[:2])
CUDA_VERSION = torch.__version__.split("+")[-1]
print("torch: ", TORCH_VERSION, "; cuda: ", CUDA_VERSION)
print("detectron2:", detectron2.__version__)

結果

nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2021 NVIDIA Corporation
Built on Sun_Feb_14_21:12:58_PST_2021
Cuda compilation tools, release 11.2, V11.2.152
Build cuda_11.2.r11.2/compiler.29618528_0
torch:  1.12 ; cuda:  cu113
detectron2: 0.6

importするライブラリは下記の通りです。

# Some basic setup:
# Setup detectron2 logger
import detectron2
from detectron2.utils.logger import setup_logger
setup_logger()

# import some common libraries
import numpy as np
import os, json, cv2, random
from google.colab.patches import cv2_imshow

# import some common detectron2 utilities
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, DatasetCatalog

データセットを登録します。
下記コードでdog-projectという名前でデータが登録され、今後データを読み出す際はこのdog-projectという名前を利用します。

from detectron2.data.datasets import register_coco_instances
# register_coco_instances(name, metadata, json_file, image_root) # parameter
register_coco_instances("dog-project", {}, "/content/drive/MyDrive/data/dog/test.json", "/content/drive/MyDrive/data/dog")

データセットの詳細を設定します。
Keypoint検出に特有な項目が下記3点です。
①keypoint_names
設定したkeypointの名称を順番にリスト形式で渡します。
list[str]

②keypoint_flip_map
画像を水平方向反転した際に、ラベルを入れ替える項目を設定します。
データオーギュメンテーションの際に用いられます。
ex.) 人：右手と左手、右足と左足
list[tuple[str]]

③keypoint_connection_rules
各キーポイントの対応と表示する際の色を(R,G,B)で渡します。
list[tuple(str, str, (r, g, b))]

*参照

詳細設定の際は、
MetadataCatalog.get("dog-project(*register_coco_instancesで設定した名前)")
を使用しましょう。

from detectron2.data import MetadataCatalog

keypoint_names = ['head', 'neck', 'front-right', 'front-left', 'hip', 'hind-right', 'hind-left']
keypoint_flip_map = [('front-right','front-left'),('hind-right','hind-left')]
keypoint_connection_rules = [('head','neck',(128,0,0)),('neck','front-right',(0,128,0)),
                             ('neck','front-left',(0,0,128)),('neck','hip',(255,0,0)),('hip','hind-right',(0,255,0)),
                             ('hip','hind-left',(0,0,255))]

MetadataCatalog.get("dog-project").thing_classes = ["Dog"]
MetadataCatalog.get("dog-project").thing_dataset_id_to_contiguous_id = {1:0}
MetadataCatalog.get("dog-project").keypoint_names = keypoint_names
MetadataCatalog.get("dog-project").keypoint_flip_map = keypoint_flip_map
MetadataCatalog.get("dog-project").keypoint_connection_rules = keypoint_connection_rules

読み込んだデータを確認します。

dog_metadata = MetadataCatalog.get("dog-project")
dataset_dicts = DatasetCatalog.get("dog-project")

for d in random.sample(dataset_dicts, 3):
  img = cv2.imread(d["file_name"])
  visualizer = Visualizer(img[:, :, ::-1], metadata=dog_metadata, scale=1.0)
  vis = visualizer.draw_dataset_dict(d)
  cv2_imshow(vis.get_image()[:, :, ::-1])

こんな形で画像が表示されればOKです。

それでは学習を行っていきます。

注意点は下記2点です。
①全体のクラス数設定
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
cfg.MODEL.RETINANET.NUM_CLASSES = 1

ここの部分は自分が設定したインスタンスのクラス数(ここでは犬のみなので1)を設定しましょう。

②keypointの数
cfg.MODEL.ROI_KEYPOINT_HEAD.NUM_KEYPOINTS = 7
cfg.TEST.KEYPOINT_OKS_SIGMAS = np.ones((7,1),dtype=float).tolist()

ここの部分は自分が設定したkeypointの数(ここでは"head", "neck", "fornt-right", "front-left", "hip", "hind-right", "hind-left"の7)を設定しましょう。

参照

OKSについての説明はこちらを参照

from detectron2.engine import DefaultTrainer

cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-Keypoints/keypoint_rcnn_R_50_FPN_1x.yaml"))
cfg.DATASETS.TRAIN = ("dog-project",)
cfg.DATASETS.TEST = ()
cfg.DATALOADER.NUM_WORKERS = 2
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-Keypoints/keypoint_rcnn_R_50_FPN_1x.yaml")  # Let training initialize from model zoo
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = 0.00025  # pick a good LR
cfg.SOLVER.MAX_ITER = 2000    # 300 iterations seems good enough for this toy dataset; you will need to train longer for a practical dataset
cfg.SOLVER.STEPS = []        # do not decay learning rate
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128   # faster, and good enough for this toy dataset (default: 512)
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
cfg.MODEL.RETINANET.NUM_CLASSES = 1
cfg.MODEL.ROI_KEYPOINT_HEAD.NUM_KEYPOINTS = 7
cfg.TEST.KEYPOINT_OKS_SIGMAS = np.ones((7,1),dtype=float).tolist()
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = DefaultTrainer(cfg) 
trainer.resume_or_load(resume=False)
trainer.train()

学習が実行されます。

3-3. 推論

学習されたモデルは、"./output/model_final.pth"に保存されます。
推論用のモデルはそこを参照しましょう。

cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "/content/drive/MyDrive/output/model_final.pth")
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7
cfg.MODEL.DEVICE = "cpu"
predictor = DefaultPredictor(cfg)

学習データにない画像でテストしてみます。

imgPath = "/content/drive/MyDrive/data/dog/test.png"
im = cv2.imread(imgPath)
    
outputs = predictor(im)
v = Visualizer(im[:, :, ::-1],
               metadata=dog_metadata, 
               scale=1.0
)
v = v.draw_instance_predictions(outputs["instances"].to("cpu"))
cv2_imshow(v.get_image()[:, :, ::-1])

実行結果がこちら。

右前足(濃い緑の線)がおかしなところにいっていますが、おおよそ良いのではないでしょうか。今回は学習データ15枚なので、もっと増やせば精度あがると思います。

4. まとめ

今回はDetectron2を使った自作データでのキーポイント検出を試してみました。非常に手軽にできるため、色々なことに応用できそうですね。

5. 参考

参考YOU TUBE動画
https://www.youtube.com/watch?v=xX9fPmclEHY