YOLOv3 – 学習済みモデルで画像から人や車を検出する方法

YOLOv3 – 学習済みモデルで画像から人や車を検出する方法

概要

YOLOv3 の MSCOCO の学習済みモデルで画像から人や車を検出する方法について紹介します。

Advertisement

環境

  • Ubuntu 18.04
  • Windows 10

準備

YOLOv3 の Pytorch 実装である nekobean/pytorch_yolov3 を使用します。

まず、レポジトリをクローンします。

git clone https://github.com/nekobean/pytorch_yolov3.git
cd pytorch_yolov3

依存ライブラリをインストールします。

pip install -r requirements.txt

重みをダウンロードします。

./weights/download_weights.sh

ダウンロードが完了すると、weights ディレクトリ以下に2つのファイルがあるはずです。

weights/
|-- download_weights.sh
|-- yolov3-tiny.weights
`-- yolov3.weights
  • yolov3.weights: MSCOCO で学習した YOLOv3 の重み
  • yolov3-tiny.weights: MSCOCO で学習した YOLOv3-tiny の重み

Windows の場合、bash が使えないので、手動で yolov3.weights 及び yolov3-tiny.weights をダウンロードして、weights/ ディレクトリに置いてください。

推論する

MSCOCO の学習済みモデル

公式サイトで配布されている yolov3.weights 及び yolov3-tiny.weights は、次の80クラスを含むおよそ10万枚の画像で構成されるデータセットで学習した重みです。 以下の物体を検出したい場合は、自分でデータセットを作成して学習を行わなくても、学習済みモデルを利用するとすぐに検出を行うことができます。

クラス ID クラス名 クラス ID クラス名 クラス ID クラス名 クラス ID クラス名
0 20 ゾウ 40 ワイングラス 60 ダイニングテーブル
1 自転車 21 クマ 41 カップ 61 トイレ
2 22 シマウマ 42 フォーク 62 テレビ
3 バイク 23 キリン 43 ナイフ 63 ノートパソコン
4 飛行機 24 リュックサック 44 スプーン 64 マウス
5 バス 25 45 ボウル 65 リモコン
6 電車 26 ハンドバッグ 46 バナナ 66 キーボード
7 トラック 27 ネクタイ 47 リンゴ 67 携帯電話
8 28 スーツケース 48 サンドウィッチ 68 電子レンジ
9 信号機 29 フリスビー 49 オレンジ 69 オーブン
10 消火栓 30 スキー 50 ブロッコリー 70 トースター
11 ストップサイン 31 スノーボード 51 キャロット 71 シンク
12 パーキングメーター 32 スポーツボール 52 ホットドッグ 72 冷蔵庫
13 ベンチ 33 カイト 53 ピザ 73
14 34 野球バット 54 ドーナツ 74 時計
15 35 野球グラブ 55 ケーキ 75 花瓶
16 36 スケートボード 56 椅子 76 ハサミ
17 37 サーフボード 57 ソファー 77 テディベア
18 38 テニスラケット 58 鉢植え 78 ヘアードライヤー
19 39 ボトル 59 ベッド 79 歯ブラシ
Advertisement

YOLOv3 と YOLOv3-tiny について

YOLOv3-tiny は YOLOv3 よりパラメータ数を減らした軽量版のモデルです。 検出精度は YOLOv3 に劣りますが、高速に動作します。計算リソースが潤沢でなく、YOLOv3 を動かすのが難しい場合は YOLOv3-tiny を試してみてください。

画像から検出する

--input に画像ファイルのパス、または画像があるディレクトリを指定します。--output には出力ディレクトリを指定します。

python detect_image.py \
    --input <画像ファイルまたは画像があるディレクトリ> \
    --output <出力ディレクトリ> \
    --weights weights/yolov3.weights \
    --config config/yolov3_coco.yaml

試しに data/herd_of_horses.png に対して、検出を行ってみます。

python detect_image.py --input data/herd_of_horses.png --output output --weights weights/yolov3.weights --config config/yolov3_coco.yaml

実行が完了すると、output/herd_of_horses.png に結果が出力されます。

YOLOv3-tiny を使う

YOLOv3-tiny を使用する場合、--weightsweights/yolov3-tiny.weights--configconfig/yolov3tiny_coco.yaml を指定してください。

python detect_image.py --input data/herd_of_horses.png --output output --weights weights/yolov3-tiny.weights --config config/yolov3tiny_coco.yaml

YOLOv3-tiny は高速ですが、検出精度は YOLOv3 と比較すると、下がります。

画像から人を検出する

YOLOv3 を利用して人検出を行うサンプルコードです。クラス「人」は MSCOCO の学習済みモデルに含まれているため、YOLOv3 の検出結果からクラス名が person となっている矩形だけ抽出すればよいです。

people.png

yolov3_pathgit clone した pytorch_yolov3 ディレクトリのパスを指定してください。

In [1]:
import sys
from pathlib import Path

import cv2
from IPython.display import Image, display

# git clone した pytorch_yolov3 ディレクトリのパスを指定してください
yolov3_path = Path("../pytorch_yolov3")

sys.path.append(str(yolov3_path))
from yolov3.detector import Detector

config_path = yolov3_path / "config/yolov3_coco.yaml"
weights_path = yolov3_path / "weights/yolov3.weights"


def imshow(img):
    """ndarray 配列をインラインで Notebook 上に表示する。
    """
    ret, encoded = cv2.imencode(".jpg", img)
    display(Image(encoded))


# 検出器を作成する。
detector = Detector(config_path, weights_path)

# 画像を読み込む。
img = cv2.imread("people.png")

# 検出する。
detections = detector.detect(img)

# 人の検出結果のみ抽出する。
people = list(filter(lambda x: x["class_name"] == "person", detections[0]))

# 検出結果を画像に描画する。
for bbox in people:
    cv2.rectangle(
        img,
        (int(bbox["x1"]), int(bbox["y1"])),
        (int(bbox["x2"]), int(bbox["y2"])),
        color=(0, 255, 0),
        thickness=2,
    )

imshow(img)
Darknet format weights file loaded. ../pytorch_yolov3/weights/yolov3.weights
Advertisement

画像から自動車を検出する

自動車も、「自転車」「車」「バイク」「バス」「トラック」が MSCOCO の学習済みモデルに含まれているので、検出結果からこれらのクラスでフィルタすれば、同様に検出できます。

cars.png

In [2]:
import sys
from pathlib import Path

import cv2
from IPython.display import Image, display

# git clone した pytorch_yolov3 ディレクトリのパスを指定してください
yolov3_path = Path("../pytorch_yolov3")

sys.path.append(str(yolov3_path))
from yolov3.detector import Detector

config_path = yolov3_path / "config/yolov3_coco.yaml"
weights_path = yolov3_path / "weights/yolov3.weights"


def imshow(img):
    """ndarray 配列をインラインで Notebook 上に表示する。
    """
    ret, encoded = cv2.imencode(".jpg", img)
    display(Image(encoded))


# 検出器を作成する。
detector = Detector(config_path, weights_path)

# 画像を読み込む。
img = cv2.imread("cars.png")

# 検出する。
detections = detector.detect(img)

# 人の検出結果のみ抽出する。
target = ["bicycle", "car", "motorcycle", "bus", "truck"]
cars = list(filter(lambda x: x["class_name"] in target, detections[0]))

# 検出結果を画像に描画する。
for bbox in cars:
    cv2.rectangle(
        img,
        (int(bbox["x1"]), int(bbox["y1"])),
        (int(bbox["x2"]), int(bbox["y2"])),
        color=(0, 255, 0),
        thickness=2,
    )

imshow(img)
Darknet format weights file loaded. ../pytorch_yolov3/weights/yolov3.weights