• ホーム
  • Chainer
  • ChainerCVのデモンストレーションプログラムを読んでみた（推論編）

ChainerCVのデモンストレーションプログラムを読んでみた（推論編）

ChainerCVの物体認識（SSD）のデモンストレーションのファイルを読んでみました。今回は推論の部分です。

目次

  1. 本編
  2. プログラムの中身
  3. まとめ

本編

ChainerCV の物体認識（SSD）のデモンストレーションプログラムのうち、 デモそのもののファイル（demo.py） を読んでみます。

デモの内容は、学習済みのSSD300またはSSD512を使って任意の画像の物体認識をするというものです。

プログラムの中身

コードの最初の部分は大量のインポート文です。Pythonですからね。

引数の解釈

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--model', choices=('ssd300', 'ssd512'), default='ssd300')
    parser.add_argument('--gpu', type=int, default=-1)
    parser.add_argument('--pretrained-model')
    parser.add_argument(
        '--dataset', choices=('voc',), default='voc')
    parser.add_argument('image')
    args = parser.parse_args()

argparseモジュールを使って、コマンドライン引数の内容をパースします。パースした結果をargsオブジェクトに格納します。

モデルの選択

if args.model == 'ssd300':
        cls = SSD300
    elif args.model == 'ssd512':
        cls = SSD512

引数の指定に基づいて、推論に使用するネットワークを選択してclsという名前のインスタンスにします。SSD300とSSD512というのは、ChainerCVにプリセットされているリンク（ネットワーク）のクラスです。

データセットの選択

if args.dataset == 'voc':
    if args.pretrained_model is None:
        args.pretrained_model = 'voc0712'
    label_names = voc_bbox_label_names

datasetコマンドライン引数がvocの場合で、pretrained_modelに何も指定されていない場合は、pretrained_modelをvoc0712に設定します。ここで自前の学習データを使うのか、それともChainerCVのvoc0712を使うのか選択するわけですね。

voc_bbox_label_namesはVOCのラベルのデータセットで、インポートすることでセットできるようになっているようです。

モデルのインスタンスを作る

model = cls(n_fg_class=len(label_names),
            pretrained_model=args.pretrained_model)

ここでモデルのインスタンスを作ります。

n_fg_classが推論するときのクラス数の指定です。クラス数と事前学習データをネットワークオブジェクトに指定して、モデルのインスタンスを作るわけです。

推論をする

img = utils.read_image(args.image, color=True)
bboxes, labels, scores = model.predict([img])
bbox, label, score = bboxes[0], labels[0], scores[0]

vis_bbox(
    img, bbox, label, score, label_names=label_names)
plt.show()

ChainerCVのutilsモジュールのread_imageを使って画像を読み込みます。ファイルから読み込んだ画像をndarrayに変換するようです。

モデルのオブジェクトのpredictメソッドに読み取った画像をndarray形式で渡すと、ボックスの座標とラベルとスコアのそれぞれのリストが返ってきます。

リストの中から1つ選んで、vis_boxという結果表示用のユーティリティに渡して、描画します。

これ、scoresの最大値をとるのではなくて、最初のをとるのですね。

まとめ

SSDを使って推論するのであれば、pretrained_modelに自前の学習データとラベルを指定しさえすれば、このまま使えそうな感じです。

ChainerCVのデモンストレーションプログラムを読んでみた   MNISTの画像をPNGに変換してみた

公開日

広告

Chainerカテゴリの投稿