PyTorchでFaster R-CNN 物体検出モデルを作成する

PyTorchでFaster R-CNNの物体検出モデルの作成する手順を紹介します。任意のNotebookとGPUの環境があれば実行可能と思いますが、以降ではWatson StudioのNotebookを使用します。

Watson StudioではGPU版のNotebookは2020年3月現在ベータ版での提供となります。使用する場合は以下リンク先からリクエストすることで使用可能となります。
https://medium.com/ibm-watson/gpu-powered-notebook-is-coming-to-watson-studio-56876f60c056

教師データの作成

以下の様なWatsonのロゴマークを物体検出の対象とします。左側が現行のWatsonロゴ(label: watson)、右側がその他のWatsonロゴ(label: watson_others)になります。

画像のアノテーションツールを使ってロゴを学習させていきます。画像のアノテーションツールはさまざまものがありますが、私はlabelImgを使用しました。導入手順はリンク先の他、色々なところで紹介されているので割愛します。学習のための画像をWeb等で収集した上で、labelImg で watson および watson_others をアノテーションしていきます。アノテーション結果は Pascal VOC 形式で出力します。

学習で使用した画像とアノテーション結果をGitに置いているので、以降のNotebookではこちらをダウンロードして使用可能です。

学習

Watson Studioでプロジェクトを作成して、以下リンク先から Notebook を作成します。

Notebookを実行して行くと、Faster R-CNNによるモデルが作成されます。

テスト

Notebookの最後のセルでいくつかテスト結果を載せていますが、うまく物体検出ができているようです。

教師データを入れ替えれば、同様の手順で物体検出モデルを比較的少ないコードで実装可能ですので、お試しください。