Detectron2のModel Zooで物体検出、インスタンスセグメンテーション、姿勢推定
概要
Detectron2のModel Zooにある訓練済みを使って、物体検出やインスタンスセグメンテーション、姿勢推定等を行う。 多くのモデルに対して一括で処理できるコードを作った。便利。
Detectron2
FacebookのAI研究グループ(FAIR)が開発している物体検出アルゴリズムを実装のためのソフトウェア。
環境
- ubuntu 18.04
- GeForce GTX TITAN X
インストール
詳細は省略。ほぼ公式の通りやった。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
- アルゴリズム
確認用コード
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')