OpenCV + YOLOv3で物体検出を行う

2018年6月8日

どうも。帰ってきたOpenCVおじさんだよー。

そもそもYOLOv3って？

YOLO(You Look Only Onse)という物体検出のアルゴリズムで、画像を一度CNNに通すことで物体の種類が何かを検出してくれるもの、らしい。

使い方（ほぼ独自のため、正しいかとても怪しい……）

環境構築

事前にTensorflowなりKerasなりOpenCVなりをインストールしておく。Pyenvなどの方法は割愛。わからなければこちらをみて。

python -m venv ~/.venv/yolo
source ~/.venv/yolo/bin/activate
pip install --upgrade pip

pip install tensorflow
pip instsall keras
pip install opencv-python
pip install pillow

yolo.h5ファイルの作成

理解はあとにしておくことにして、yolo.weightsという学習データをyolo.h5というファイルに変換する必要がある。その方法は次の通り。

git clone https://github.com/xiaochus/YOLOv3
cd YOLOv3
python yad2k.py cfg/yolo.cfg yolov3.weights data/yolo.h5

とりあえずdemo.pyを実行してみる

demo.pyを実行すると、次のような結果が出た

python demo.py

ブログには書いていないのだけれど、実のところYOLOv2も使っていて、それよりも精度が上がっている模様。
ただ、動作がだいぶ重く感じる。

オリジナルでカメラからキャプチャーしてみた。

OpenCVでカメラから取り出した画像を物体検出してみた。ソースコードは以下の通り。

"""Demo for use yolo v3
"""
import os
import time
import argparse
import datetime

import cv2
import numpy as np

from PIL import Image
from keras.models import load_model
from model.yolo_model import YOLO

IMAGE_DIR = os.path.dirname(os.path.abspath(__file__)) + "/images/output/"

def process_image(img):
    """Resize, reduce and expand image.

    # Argument:
        img: original image.

    # Returns
        image: ndarray(64, 64, 3), processed image.
    """
    image = cv2.resize(img, (416, 416),
                       interpolation=cv2.INTER_CUBIC)
    image = np.array(image, dtype='float32')
    image /= 255.
    image = np.expand_dims(image, axis=0)

    return image


def get_classes(file):
    """Get classes name.

    # Argument:
        file: classes name for database.

    # Returns
        class_names: List, classes name.

    """
    with open(file) as f:
        class_names = f.readlines()
    class_names = [c.strip() for c in class_names]

    return class_names


def draw(image, boxes, scores, classes, class_names):
    """画像から検出されたオブジェクトを枠で囲う

    # Argument:
        image: original image.
        boxes: ndarray, boxes of objects.
        classes: ndarray, classes of objects.
        scores: ndarray, scores of objects.
        all_classes: all classes name.
    """
    for box, score, cl in zip(boxes, scores, classes):
        x, y, w, h = box

        top = max(0, np.floor(x + 0.5).astype(int))
        left = max(0, np.floor(y + 0.5).astype(int))
        right = min(image.shape[1], np.floor(x + w + 0.5).astype(int))
        bottom = min(image.shape[0], np.floor(y + h + 0.5).astype(int))

        cv2.rectangle(image, (top, left), (right, bottom), (255, 0, 0), 2)
        cv2.putText(image, '{0} {1:.2f}'.format(class_names[cl], score),
                    (top, left - 6),
                    cv2.FONT_HERSHEY_SIMPLEX,
                    0.6, (0, 0, 255), 1,
                    cv2.LINE_AA)

        print('class: {0}, score: {1:.2f}'.format(class_names[cl], score))
        print('box coordinate x,y,w,h: {0}'.format(box))

    print()

def _main(args):

    classes_path = os.path.expanduser(args.classes_path)

    with open(classes_path) as f:
        class_names = f.readlines()
    class_names = [c.strip() for c in class_names]

    yolo = YOLO(0.6, 0.5)
    cap = cv2.VideoCapture(0)

    while True:
        _, image = cap.read()

        # 画像を変換
        pimage = process_image(image)

        # 物体検出
        boxes, classes, scores = yolo.predict(pimage, image.shape)

        # 物体があったら枠で囲う
        if boxes is not None:
            draw(image, boxes, scores, classes, class_names)

        # 画像を保存
        image = Image.fromarray(image)
        file_name = datetime.datetime.now().strftime("%Y%m%d%H%i%s") + ".jpg"
        image.save(IMAGE_DIR + file_name)

if __name__ == '__main__':

    _main(parser.parse_args())