Detectron2のModel Zooで物体検出、インスタンスセグメンテーション、姿勢推定

概要

Detectron2のModel Zooにある訓練済みを使って、物体検出やインスタンスセグメンテーション、姿勢推定等を行う。 多くのモデルに対して一括で処理できるコードを作った。便利。

Detectron2

FacebookのAI研究グループ(FAIR)が開発している物体検出アルゴリズムを実装のためのソフトウェア。

環境

インストール

詳細は省略。ほぼ公式の通りやった。Windowsでやろうとしたら公式対応してないから大変そうな雰囲気。

メモ代わりに大雑把な手順を載せるが公式を読んでやったほうがいい。

  • OS、cuda等:dockerコンテナ利用
    • nvidia/cuda:10.1-cudnn7-devel
  • Python環境
    • Anaconda 2019.10でpython3.7環境構築
  • pytorch
    • conda install pytorch torchvision cudatoolkit=10.1 -c pytorch
  • detectron2

      git clone https://github.com/facebookresearch/detectron2.git
      cd detectron2
      python setup.py build develop
    

Detectron2 Model Zoo

学習済みのモデルが置いてある。 以下のようなものがあるが、詳細はウェブサイト参照。

  • タスク
    • COCO Object Detection
    • COCO Instance Segmentation
    • COCO Person Keypoint Detection
    • COCO Panoptic Segmentation
    • LVIS Instance Segmentation
    • Cityscapes
    • Pascal VOC
  • アルゴリズム
    • Faster R-CNN、RetinaNet、Mask R-CNN、Cascade R-CNN等
    • backbone: FPN、ResNetベースのもの

確認用コード

Model Zooにある学習済みモデルをロードするモジュールと、予測して可視化するためのモジュールがあったのでそれを使った。

以下のコードを実行するとModel Zooにあるモデルの42個について予測を行える。

from detectron2.data.detection_utils import read_image
from detectron2.model_zoo.model_zoo import ModelZooUrls
from detectron2.config import get_cfg
from demo.predictor import VisualizationDemo

img_path = 'input.jpg'
img = read_image(img_path, format="BGR")

for i, config_path in enumerate(ModelZooUrls.CONFIG_PATH_TO_URL_SUFFIX.keys()):
    # rpnとfast_rcnnは可視化対応していないので飛ばす
    if 'rpn' in config_path or 'fast_rcnn' in config_path:
        continue
    # config設定
    cfg = get_cfg()
    cfg.merge_from_file(f'configs/{config_path}')
    cfg.MODEL.WEIGHTS = ModelZooUrls.get(config_path)
    score_thresh = 0.5
    cfg.MODEL.RETINANET.SCORE_THRESH_TEST = score_thresh
    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = score_thresh
    cfg.MODEL.PANOPTIC_FPN.COMBINE.INSTANCES_CONFIDENCE_THRESH = score_thresh
    cfg.freeze()
    # 検出&可視化
    demo = VisualizationDemo(cfg)
    predictions, visualized_output = demo.run_on_image(img)
    # ファイル出力
    dataset_name, algorithm = config_path.split("/")
    algorithm = algorithm.split('.')[0]
    visualized_output.save(f'out/{i:02d}-{dataset_name}-{algorithm}.jpg')

結果

f:id:whoopsidaisies:20191218181902j:plainf:id:whoopsidaisies:20191218181904j:plainf:id:whoopsidaisies:20191218181907j:plainf:id:whoopsidaisies:20191218181911j:plain
f:id:whoopsidaisies:20191218181914j:plainf:id:whoopsidaisies:20191218181916j:plainf:id:whoopsidaisies:20191218181921j:plainf:id:whoopsidaisies:20191218181926j:plain
f:id:whoopsidaisies:20191218181930j:plainf:id:whoopsidaisies:20191218181932j:plainf:id:whoopsidaisies:20191218181935j:plainf:id:whoopsidaisies:20191218181937j:plain
f:id:whoopsidaisies:20191218181940j:plainf:id:whoopsidaisies:20191218181943j:plainf:id:whoopsidaisies:20191218181945j:plainf:id:whoopsidaisies:20191218181947j:plain
f:id:whoopsidaisies:20191218181950j:plainf:id:whoopsidaisies:20191218181953j:plainf:id:whoopsidaisies:20191218181956j:plainf:id:whoopsidaisies:20191218181959j:plain
f:id:whoopsidaisies:20191218182001j:plainf:id:whoopsidaisies:20191218182004j:plainf:id:whoopsidaisies:20191218182007j:plainf:id:whoopsidaisies:20191218182010j:plain
f:id:whoopsidaisies:20191218182013j:plainf:id:whoopsidaisies:20191218182016j:plainf:id:whoopsidaisies:20191218182020j:plainf:id:whoopsidaisies:20191218182023j:plain
f:id:whoopsidaisies:20191218182026j:plainf:id:whoopsidaisies:20191218182029j:plainf:id:whoopsidaisies:20191218182031j:plainf:id:whoopsidaisies:20191218182035j:plain
f:id:whoopsidaisies:20191218182039j:plainf:id:whoopsidaisies:20191218182041j:plainf:id:whoopsidaisies:20191218182044j:plainf:id:whoopsidaisies:20191218182047j:plain
f:id:whoopsidaisies:20191218182051j:plainf:id:whoopsidaisies:20191218182055j:plainf:id:whoopsidaisies:20191218182059j:plainf:id:whoopsidaisies:20191218182101j:plain
f:id:whoopsidaisies:20191218182104j:plainf:id:whoopsidaisies:20191218182107j:plainf:id:whoopsidaisies:20191218181911j:plainf:id:whoopsidaisies:20191218181907j:plain