Detectron2で小銭を数える

はじめに

この記事ではDetectron2というFacebook AIが開発している物体検出ライブラリを利用して，自作のデータセットに対して物体検出をしてみます．実際にやってみるとDetectron2に関する日本語・英語の情報があまり見つからず少し苦労したので，私と同じようにDetectron2で物体検出をしようとしている人の助けになれば幸いです．

データセットとして日本硬貨４種類（1円，5円，10円，100円）が含まれた画像を用意し，PretrainedのFaster R-CNNを訓練しています．結果として次のように硬貨を検知できました．

※データセットに50円，500円硬貨が含まれていないのは，たまたま財布に入っていなかったためです．
※今気づきましたがhundredをスペルミスしてますね．恥ずかしいですが直すのが面倒なのでこのままで行きます．

僕は深層学習に関して全く詳しくないため，訓練時のパラメータや設定などは公式チュートリアルそのままか，あるいは適当に決めています．なのでその辺りに関しては説明しません（できません）．データセットの準備や訓練の手順を説明し，上の画像のようなそこそこ使えそうな結果を得ることを目指しています．

作成した硬貨のデータセットはGitHubに上げています．

訓練時のコードはGoogle Colabにあります．

公式チュートリアルでは，独自のデータセット（風船の画像）のInstance Segmentationを通して基本的な使い方を学ぶことができます．

データセット作成

画像を用意する

iPhoneで硬貨の画像を30枚程度撮りました．30枚という枚数は適当です．
各画像には硬貨が1~5枚写っています．iPhoneで撮った画像はサイズがとても大きく，訓練に時間がかかってしまいそうだったので大幅に縮小しました．

Annotationをつける

撮影した画像の硬貨に対してBounding Boxを付けていきます．

Bounding Boxのデータを保存する形式にはいろいろな種類がありますが，その中でもDetectron2はCOCOフォーマットという形式をサポートしています．もちろんそれ以外の形式でも使えますが，COCOフォーマットで作っておくとデータセットの読み込みがとても楽です．

そのためにCOCO Annotatorというツールを利用しました．
使い方については以下のリンクを参考にしました．初めて使うと少し戸惑う部分があるので公式のGetting Startedに沿って進めていくのがおすすめです．

• COCO Annotator Getting Started
• COCO Annotatorを使って骨格検知用のデータセットを作ってみる

Detectron2

ここから実際にDetectron2を使っていきます．コードはGoogle Colabにあります．

セットアップ

基本的に公式のチュートリアルそのままです．Detectron2をインストールしたあとに「ランタイムの再起動」を行う必要があるので注意してください．

データセット登録

GitHubから硬貨データセットをダウンロードし展開します．

!wget https://github.com/Be4rR/JapaneseCoins/blob/main/coins.zip?raw=true -q -O coins.zip
!unzip coins.zip > /dev/null

ダウンロードしたデータセットをregister_coco_instancesで登録します．

register_coco_instances("coins", {}, "./coins/coco-1612779490.2197058.json", "./coins")

nameは重複しなければなんでも良いみたいです．後でデータセットを指定するときに使います．
json_fileはAnnotationファイルのパスを指定します．
image_rootは画像を含むフォルダのパスを指定します．

データセットが正しく登録できているか確認します．
データセットの中からランダムに３枚選び，画像とAnnotationを表示します．

coins_metadata = MetadataCatalog.get("coins")
dataset_dicts = DatasetCatalog.get("coins")

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

hundredをスペルミスしてますがデータセットの登録は問題なさそうです．

ちなみにimg[:, :, ::-1]は画像のRGB形式⇔BGR形式を変換しています．
画像を読み込むときcv2.imreadはBGR形式で画像を読み込むのに対して，VisualizerはRGB形式で画像を受け取るため変換が必要になります．
参考：Python, OpenCVでBGRとRGBを変換するcvtColor

訓練

Faster RCNNをModel Zooからダウンロードして訓練（Fine Tune）します．

いろんなパラメータがありますが，とりあえず変えないといけないのは分類のクラス数cfg.MODEL.ROI_HEADS.NUM_CLASSESです．今回は１円，５円，１０円，１００円を分類するので4にします．

学習率など他にも大事なパラメータがありますが，決め方をよく知らないのでスルーします．

cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml"))
cfg.DATASETS.TRAIN = ("coins",)
cfg.DATASETS.TEST = ()
cfg.DATALOADER.NUM_WORKERS = 2
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml")
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = 0.0004
cfg.SOLVER.MAX_ITER = (
    500
)
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = (
    128
)
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 4

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

タスクごとに様々なモデルが使えます．どんなモデルがあるのかについては次のリンクを参考にしてください．

硬貨を検出させてみる

訓練したモデルを読み込みます．
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TESTは検出の閾値です．

cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.6
predictor = DefaultPredictor(cfg)

Annotationが付いていないIMG_4693.jpg～IMG_4695.jpgを使って硬貨を検出させてみます．

for num in [4693, 4694, 4695]:    
    im = cv2.imread(f"./coins/IMG_{num}.jpg")
    outputs = predictor(im)
    v = Visualizer(im[:, :, ::-1],
                   metadata=coins_metadata, 
                   scale=1.0
    )
    v = v.draw_instance_predictions(outputs["instances"].to("cpu"))
    cv2_imshow(v.get_image()[:, :, ::-1])

たまに重複して検出してしまったりしていますが，概ねうまく検出できています．

おわりに

荒削りですがデータセット作成から訓練，簡単なテストまでできました．
モデルの評価まではできませんでしたが，気になる方は公式のチュートリアルに説明があるので見てみてください．