Zenn
🍒

【Colabで動く】Detectron2のFaster R-CNNで自作データセットの物体検出

2022/10/12に公開
Python
ディープラーニング
機械学習
PyTorch
データ分析
tech

この記事は何か

Detectron2（Facebook AI Researchが開発したPyTorchベースの物体検出ライブラリ）でFaster R-CNNのPre-Trainedモデルを使って自作データセットの物体検出をしてみます。実行環境はGoogle Colaboratoryを利用します。Colab + Detectron2 + Faster R-CNN + 自作データセットの組み合わせの記事はほとんど見受けられなかったので、備忘録がてらこの記事を書いています。

Detectron2のインストール

Detectron2の前提となるライブラリを入れていきます。Colabが使っているCUDAのバージョンを確認しましょう。

!nvcc --version

2022年9月時点だとCUDAのバージョンは11.1でした。PyTorchの公式サイトでCUDAのバージョンに対応するPyTorchやtorchvisionのバージョンを確認します。CUDAのバージョンが11.1のときは、以下のコマンドでPyTorchなどをインストールすれば良いようです。

# CUDA 11.1
pip install torch==1.10.1+cu111 torchvision==0.11.2+cu111 torchaudio==0.10.1 -f https://download.pytorch.org/whl/torch_stable.html

Colabで以下のコードを実行すると、Detectron2がインストールされます。

!pip uninstall --yes torchtext
!pip install torch==1.10.1+cu111 torchvision==0.11.2+cu111 torchaudio==0.10.1 -f https://download.pytorch.org/whl/torch_stable.html
import cv2
!python -m pip install detectron2 -f https://dl.fbaipublicfiles.com/detectron2/wheels/cu111/torch1.10/index.html

データセットを用意する

COCO FormatでJSONを出力できる適当なアノテーションツールを使ってデータを作成していきます。単純作業で辛いですが頑張りましょう。COCO Format以外でも自分でパースすれば使用可能ですが、COCO Formatだと専用のデータロード関数が用意されているので後々少し楽になります。筆者はFastLabelというアノテーションツールを使っています。

オリジナルの画像とアノテーションファイルをGoogle ドライブにアップロードします。アップロード後、以下のコードでColabからドライブにアクセスできるようにします。

from google.colab import drive


drive.mount("/content/drive")

register_coco_instances関数でデータセットをDetectron2に登録します。アノテーションファイル（annotations.json）はtrain, val, testそれぞれ別のものになります。

from detectron2.data.datasets import register_coco_instances


register_coco_instances("my_dataset_train", {}, "/path/to/annotations.json", "/path/to/your/train/images")
register_coco_instances("my_dataset_val", {}, "/path/to/annotations.json", "/path/to/your/validation/images")
register_coco_instances("my_dataset_test", {}, "/path/to/annotations.json", "/path/to/your/test/images")

筆者はここでトラップにひっかかったのですが、アノテーションツールで生成したアノテーションファイルをそのまま登録するとFiltered out {} instances without valid segmentation.と言われてアノテーションが読み込まれませんでした。おそらくですが、Detectron2は主にセグメンテーションのために使われるライブラリのため、物体検出用のアノテーション（矩形の左上と右下の座標だけ与えられている）では不十分で、矩形の頂点の座標をすべて与えてやる必要があるようです。ソースコードを読むと、頂点が2つ以下のアノテーションに関しては（2次元の図形を成さないため）無視するような仕様になっているようです。
結構調べたのですが正攻法での解決策が見つからなかったので、アノテーションファイルを編集し、以下画像でもともと赤い座標しか与えられていなかったところを、青い座標も与えるようにしてやりました。

具体的には、COCO Formatのアノテーションファイルを次の通り変更します。

変更前
"segmentation": [
    [
        x1,
        y1,
        x2,
        y2
    ]
変更後
"segmentation": [
    [
        x1,
        y1,
        x2,
	y1,
	x2,
	y2,
	x1,
        y2
    ]

正直こんなことやるなら生のPyTorch書いた方が早かったかなと思いながら上記の処理を適当なPythonスクリプトを用意してやってしまいました。
ここまでできたら、アノテーション済みの画像を確認してみましょう。

from google.colab.patches import cv2_imshow
from detectron2.data import MetadataCatalog, DatasetCatalog


my_dataset_train_metadata = MetadataCatalog.get("my_dataset_train")
dataset_dicts = DatasetCatalog.get("my_dataset_train")


import random
from detectron2.utils.visualizer import Visualizer


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

セグメンテーション仕様のアノテーションを与えてやったので、以下のような見た目になるはずです（塗りつぶされてますよね）

モデルの学習を行う

まずは、学習用の各種設定を行います。

from detectron2.config import get_cfg
from detectron2 import model_zoo


cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x.yaml"))
cfg.DATASETS.TRAIN = ("my_dataset_train",)
cfg.DATASETS.TEST = ("my_dataset_val",)
cfg.DATALOADER.NUM_WORKERS = 2
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x.yaml")
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = 0.00025
cfg.SOLVER.MAX_ITER = 1500
cfg.SOLVER.STEPS = []
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
cfg.OUTPUT_DIR = "/path/to/your/output/folder"

今回はFaster R-CNNの学習を行いたいので、事前学習済みのFaster R-CNN（faster_rcnn_X_101_32x8d_FPN_3x.yaml）をWEIGHTSとして指定しています。用意したデータセットの物体のクラスは1種類（ナス）のみですので、cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1となっています。cfg.OUTPUT_DIRには、学習済みのモデルの重みを出力したいフォルダを指定してください。
学習用のクラスを定義します。多分コピペで良いはず。

from detectron2.engine import DefaultTrainer
from detectron2.evaluation import COCOEvaluator


class CocoTrainer(DefaultTrainer):

    @classmethod
    def build_evaluator(cls, cfg, dataset_name, output_folder=None):
  
        if output_folder is None:
            os.makedirs("coco_eval", exist_ok=True)
            output_folder = "coco_eval"
  
        return COCOEvaluator(dataset_name, cfg, False, output_folder)

以下のコードで学習を実行します。

import os


os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = CocoTrainer(cfg) 
trainer.resume_or_load(resume=False)
trainer.train()

学習済みのモデルで推論を行う

訓練後の重みはmodel_final.pthに保存されています。重みをロードして、テストデータで推論を行ってみましょう。

from detectron2.engine import DefaultPredictor


cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")
cfg.DATASETS.TEST = ("my_dataset_test", )
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7
predictor = DefaultPredictor(cfg)
test_metadata = MetadataCatalog.get("my_dataset_test")


from detectron2.utils.visualizer import ColorMode
import glob


for imageName in glob.glob("/path/to/your/test/images/*.jpg"):
    im = cv2.imread(imageName)
    outputs = predictor(im)
    v = Visualizer(im[:, :, ::-1],
                  metadata=test_metadata, 
                  scale=0.8
                   )
    out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
    cv2_imshow(out.get_image()[:, :, ::-1])

それっぽい結果になりました！

まとめ

Faster R-CNNを試したいときはDetectron2ではなくて生のPyTorchを書いた方が早いかもしれない…

Discussion