MMDetection を用いて画像から猫(だけ)を検出検出してみます。
環境等
- Windows 11
- Python 3.10.5
- mm-detection: 2.51.1
MMDetection のインストールなど
pip では以下のようにもろもろインストールできます。
pip install mmdet
pip install opencv-python
pip install torch torchvision torchaudio
pip install -U openmim
mim install mmcv-full
mim download mmdet --config yolov3_mobilenetv2_320_300e_coco --dest .実装
では、書いていきます。
まずは必要なモジュールを import します。
import cv2
import numpy as np
from mmdet.apis import inference_detector, init_detector設定ファイルや判定したい画像の path を指定します。
CONFIG_FILE = r"cat_det/yolov3_mobilenetv2_320_300e_coco.py"
CHECKPOINT_FILE = r"cat_det/yolov3_mobilenetv2_320_300e_coco_20210719_215349-d18dff72.pth"
CAT_IDX = 15 # 猫は15番目
THRESHOLD = 0.2
img_path = r"cat_det/cat2.jpg"今回 MMDetection で使用するモデルは COCO データセットで事前に学習してあるモデルで、
モデルの出力では、15 番目が猫となっているので、CAT_IDX=15 としています。
また、モデルから各ラベルの確率が取得できるので、THRESHOLD=0.2 としているので、20%以上の確率であれば、それは猫と判断する時に用いるので先に用意しています。
次に、モデルを初期化と猫を検出します。
model = init_detector(CONFIG_FILE, CHECKPOINT_FILE, device="cpu") # or device='cuda:0'
result = inference_detector(model, img_path)猫以外のスコアはいらないので、初期化します。
for i, _ in enumerate(result):
if i != CAT_IDX:
result[i] = np.empty((0, 5), dtype=np.float32)最後に猫がいれば、検出してセグメンテーションした結果を保存するコードを書いていきます。
if len(result[CAT_IDX]) != 0:
score = max(result[CAT_IDX][:, 4]) # bbox(0,1,2,3),score(4)という配列なので、最後の4番目を取得。その最大値をスコアとしているのでそれも取得
if score > THRESHOLD:
model.show_result(img_path, result, out_file="result.jpg")
result = cv2.imread("result.jpg")
cv2.imshow("frame", result)コメントに書いてありますが、score は配列の最後の要素に格納してあるので、
上のコードで取得します。
あとは、MMDetection に内蔵されいるセグメンテーション結果を出力する
show_result を使って結果を保存しています。
ここまでのコードを事項するとresult.jpgとして以下のような出力がされているはずです。
終わりに
今回はMMDetectionを使って簡単に猫を検出するプログラムを書いてみました。
MMDetectionは便利でありがたいです。
今後もMMDetectionを使っていろいろ書いていきたいです。
コメント