detectron2の使い方

yoshidaです。今回はdetectron2の使い方について、COCOフォーマットのデータがある際のインスタンスセグメンテーションについてのコードを、一行ずつ見ながらその解説をしていこうと思います。

準備

drive.mount('/content/drive')
from google.colab import drive

いつものdriveをマウントするコード。

!pip install torch==1.10.0+cu113 torchvision==0.11.1+cu113 torchaudio===0.10.0+cu113 -f https://download.pytorch.org/whl/cu113/torch_stable.html

Pytorchを、CUDA 11.1, torch==1.10.0 でダウンロード。コードの詳細は下のPytorch公式より。
https://pytorch.org/get-started/previous-versions/

!pip install pyyaml==5.1

pyyamlのインストール。これがあることでyamlファイルが扱えるようになる。

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

%cd /content/drive/MyDrive/Detectron2
!pip install detectron2 -f https://dl.fbaipublicfiles.com/detectron2/wheels/$CUDA_VERSION/torch$TORCH_VERSION/index.html

torchをimport, torchとcudaのバージョンを取ってきて、それをあらかじめ決まっている型に当てはめることにより、
https://detectron2.readthedocs.io/en/latest/tutorials/install.html
の、”Install Pre-Built Detectron2 (Linux only)”に従ってDetectron2をインストールする。

# !git clone https://github.com/facebookresearch/detectron2.git

一応githubからdetectron2をクローンしているが、これをしなくとも動いたため不要であると考えられる。何かしら設定を変えなければならないときは、下記のget_cfgから設定すればよいと考えられる。

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 glob import glob

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

detectron2内外の各種必要なライブラリを取ってくる。

訓練

from detectron2.data.datasets import register_coco_instances
register_coco_instances("train", {}, "/content/drive/MyDrive/Detectron2/json/train.json", "/content/drive/MyDrive/Detectron2/images")

register_coco_instances("val", {}, "/content/drive/MyDrive/Detectron2/json/val.json", "/content/drive/MyDrive/Detectron2/images")

train_metadata = MetadataCatalog.get("train")
val_metadata = MetadataCatalog.get("val")

register_coco_instancesで、インスタンスレベル（物体検出やセグメンテーションなど）のデータセットが既にCOCO形式のjsonファイルであれば、データセットとその関連するメタデータを簡単に登録することが可能。
register_coco_instanceの第1引数がデータセット名になっていて、それをMetadataCatalog.getでメタデータをを取ってこれる。

print(train_metadata)
dataset_dicts = DatasetCatalog.get("train")

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

print(val_metadata)
dataset_dicts = DatasetCatalog.get("val")

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

train, valそれぞれにおいて、先程同様にDatasetCatalog.getでデータセットを取ってきて、そのなかから一つを取る。それがjpgで取ってくるため、cv2.imreadで画像をndarray形式に変換し、Visualizerを初期化し、visualizer.draw_dataset_dictでアノテーションを画像に反映し、それをcv2_imshowで表示する。cv2.imreadはBGRとなるようにうけとるため、img[:, :, ::-1]でRGB形式に戻す。

from detectron2.engine import DefaultTrainer

cfg = get_cfg()
# cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_101_FPN_3x.yaml"))
cfg.DATASETS.TRAIN = ("train",)
cfg.DATASETS.TEST = ()
cfg.DATALOADER.NUM_WORKERS = 2
# cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")  # Let training initialize from model zoo
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_101_FPN_3x.yaml")
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = 0.00025  # pick a good LR
cfg.SOLVER.MAX_ITER = 300    # 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 = 2  # 2 classes
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = DefaultTrainer(cfg) 
trainer.resume_or_load(resume=False)
trainer.train()

デフォルトの学習ロジックを持つ、DefaultTrainerをimportする。
get_cfg()で一旦デフォルトのconfigを取ってくる。
trainer = DefaultTrainer(cfg) : configからtrainerを作成
trainer.resume_or_load(resume=False) : 最後のチェックポイント、またはMODEL.WEIGHTSからロードする
trainer.train() : 学習を実行

configについては以下の表のようになっている。

%load_ext tensorboard
%tensorboard --logdir output

tensorboardで学習の様子を確認。

推論

cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_101_FPN_3x.yaml"))
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 2 
cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "/content/drive/MyDrive/Detectron2/output/model_final.pth")
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.8
cfg.MODEL.DEVICE = "cpu"
predictor = DefaultPredictor(cfg)

学習同様にconfigを整備する。
DefaultPredictor : configからpredictorを作る

dataset_dicts = DatasetCatalog.get("val")
for imagePath in dataset_dicts: 
    im = cv2.imread(imagePath["file_name"])
    
    outputs = predictor(im)
    v = Visualizer(im[:, :, ::-1],
                   metadata=val_metadata, 
                   scale=1.0
    )
    v = v.draw_instance_predictions(outputs["instances"].to("cpu"))
    cv2_imshow(v.get_image()[:, :, ::-1])

検証用データセットから画像を取ってきて、predictorにかけて表示する。
ここでoutputs = predictor(im[:, :, ::-1])としなくてよい理由はよくわかっていない。

from detectron2.evaluation import COCOEvaluator, inference_on_dataset
from detectron2.data import build_detection_test_loader
evaluator = COCOEvaluator("val", output_dir="./output")
val_loader = build_detection_test_loader(cfg, "val")
print(inference_on_dataset(predictor.model, val_loader, evaluator))

COCOEvaluatorでCOCOの指標で評価するevaluatorを作り、val_loaderでデータをロードし、inference_on_datasetにかけて評価する。

考察及びまとめ

インスタンスセグメンテーションをする際にdetectron2でよく使われているのが、2017年のMask R-CNNということであるが、もしかしてdetectron2のように簡単に使えるライブラリを使うのではなく、YOLOか何かの最新技術を使った方がよいのかもしれない。
普段、ライブラリの仕様に困った時にQiitaやStack Overflowを多用していたが、英語に臆さずちゃんと、一次情報である公式のリファレンスを見ることがいかに大事かを学んだ。

参考資料

Detectron2で自作データを学習し物体検出をしてみた【インスタンスセグメンテーション】

zenn.dev

detectron2で物体検出。COCOフォーマットのデータ準備から評価まで - Qiita

0. 初めに2020年末ぐらいにkaggleで開催された胸部レントゲン画像コンペに参加しました。胸のレントゲン画像にある病変部位を見つけるという、物体検知タスクのコンペでした。このコンペで銅メダ…

qiita.com