こんにちは。投資エンジニアの三年坊主(@SannenBouzu)です。
今回は、ディープラーニングで物体検出する際に、複数の画像から一度に物体を検出する方法(batch inference)を紹介します。
- predictorを呼ぶ回数を減らせる
- トータルの処理時間を減らせる(未確認ですがおそらく)
- 今回紹介するフレームワーク(Detectron2)だと実装が少し面倒
画像からの物体検出
- 入力:画像
- 出力:画像の中で「何が」「どこに」あるか
のような「物体検出」を行う機械学習モデルを考えます。
すでに学習済みのモデルファイルを読み込んで、画像からの物体検出を行います。
ライブラリ・フレームワーク【PyTorch/Detectron2】
PyTorchベースの、Detectron2というフレームワークを使います。
公式のColab Notebookから、セットアップを抜き出して・・・
1 2 3 4 5 6 7 8 9 10 11 12 13 | import detectron2 from detectron2.utils.logger import setup_logger setup_logger() # import some common libraries import numpy as np import os, json, cv2, random # import some common detectron2 utilities from detectron2.checkpoint import DetectionCheckpointer from detectron2.engine import DefaultPredictor from detectron2.modeling import build_model from detectron2.config import get_cfg |
Detectron2の使い方自体は、公式のColab Notebookを参照してください。
2枚の画像を読み込むところまで、コードを書きます。
1 2 3 4 5 6 7 8 9 10 | cfg = get_cfg() ... cfg.MODEL.WEIGHTS = './some_trained_model.pth' ... image_file_1 = "image1.jpg" image_file_2 = "image2.jpg" im1 = cv2.imread(image_file_1) im2 = cv2.imread(image_file_2) |
一枚の画像からの物体検出
公式のColab Notebookの通り、DefaultPredictorを生成し、画像を渡すと物体検出の結果が返ってきます。
1 2 3 4 5 6 7 8 9 | predictor = DefaultPredictor(cfg) # inference: image 1 outputs1 = predictor(im1) instances1 = outputs1["instances"].to("cpu") # inference: image 2 outputs2 = predictor(im2) instances2 = outputs2["instances"].to("cpu") |
複数の画像からの物体検出
問題はこちら。
DefaultPredictorは一度に一枚の画像しかinferenceできません。
なので、複数の画像をまとめて(predictorではなく)modelを直接呼ぶ必要があります。
@darkAlert as said above, you need to call a model, not predictor, with batch of inputs. As the docs (https://detectron2.readthedocs.io/modules/engine.html#detectron2.engine.defaults.DefaultPredictor) said the DefaultPredictor does not support batch of inputs.
引用:Do you support batch inference? · Issue #282 · facebookresearch/detectron2 · GitHub
上記issueのポイントは以下の通り。
- DetectionCheckpointer(model).load(file_path) で重みをロードする
- numpy 画像の次元は (height, width, channel) の順番だが、PyTorch(のテンソル)では (channel, height, width)(下記 transpose(2, 0, 1) で順番を入れ替えている)
- model.eval() または model.train(False) でモデルをinference modeにする
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 | model = build_model(cfg) # returns a torch.nn.Module model.eval() # inference mode checkpointer = DetectionCheckpointer(model) checkpointer.load('./some_trained_model.pkl') def ndarray_as_tensor(im: np.ndarray): return torch.as_tensor(im.astype("float32").transpose(2, 0, 1)) inputs = [{"image": ndarray_as_tensor(im1)}, {"image": ndarray_as_tensor(im2)}] # inputs is ready # inference: image 1 & 2 together outputs = model(inputs) instances1 = outputs[0]["instances"].to("cpu") instances2 = outputs[1]["instances"].to("cpu") |
大まかな流れは以上です。
DefaultPredictorクラスのdocstringにも書いてあるように、直接modelを呼ぶ方法と比べて、DefaultPredictorでは以下が追加されています。
cfg.MODEL.WEIGHTS
からチェックポイントをロードする- 常にBGR画像を入力として受け取り、
cfg.INPUT.FORMAT
に応じて画像変換を適用する cfg.INPUT.{MIN,MAX}_SIZE_TEST
に応じたリサイズを適用するバッチ(複数画像)の代わりに、一度に一つの入力画像を受け取り一つの出力を生成する複数画像を入力したい場合は無関係
DefaultPredictorの挙動を厳密に再現するなら、このあたりもDefaultPredictorの実装を真似するとよさそうです。
関連記事:こちらも読まれています
【2019年版】Pythonインストール・Mac編【長く安全に使える環境構築】
Pythonを快適に使いこなすMac環境【現役エンジニアおすすめはPro 13インチ】
【Keras/NumPy】データ拡張(Data Augmentation)における画像操作