(第5回)Python + OpenCV で遊んでみる(YOLOを用いた物体検出編)

はじめに

前回まではOpenCVに同梱されているカスケード型の検出器を用いて、静止画および動画を使って顔検出を行いました。
今回は、YOLOと呼ばれる物体検出法を用いた物体検出を行ってみたいと思います。

YOLOとは

YOLOとは「You Only Look Once」（一目見るだけで）の頭文字をとった略語で、一目見ただけで物体を検出できるという特徴があります。

前回までの検出法は、Sliding Window Approachな手法で、画像内を検出範囲の大きさを変えたり、動かしたりして複数回の検証を行い、パターンに一致する部分を検出する手法でした。

YOLOは、画像を一度だけCNN(Convolutional Neural Network:畳み込みニューラルネットワーク)に通すだけで、画像内の複数の物体を検出および認識を同時に行うアルゴリズムです。
1度だけCNNを通すだけなので、高速化が図られており、リアルタイムでの検出・認識にもよく用いられています。

現在は2018年に発表されたYOLOv3(Version 3)があり、YOLOv2(Version 2)ではCNNが19層モデルであったのがYOLOv3では53層モデルになっており、認識精度が大幅にアップしています。

YOLO: Real-Time Object Detection

https://pjreddie.com/darknet/yolo/

You only look once (YOLO) is a state-of-the-art, real-time object detection system.

pjreddie.com

必要なライブラリ

今回はYOLOv3で検出・認識し、OpenCVを用いて画像処理します。

YOLOv3:PyTorch用のYOLOv3を用います。

git clone https://github.com/ayooshkathuria/pytorch-yolo-v3.git

PyTorch:Pythonの機械学習ライブラリ
Pandas:データ解析支援ライブラリ
Cython:PythonをC/C++に変換するライブラリ（高速化を図るため）

conda install pandas
conda install pytorch torchvision -c pytorch
pip install cython

conda install opencv (OpenCVを導入していない場合のみ)
pip install matplotlib (matplotlibを導入していない場合のみ)

ライブラリではないですが、重みファイルを以下からダウンロードしておきます。

cd pytorch-yolo-v3
wget https://pjreddie.com/media/files/yolov3.weights

YOLOを用いた物体検出

今回使用するプログラムはPyTorch用YOLOv3に同梱されているvideo_demo.pyの内容を改変したものです。

from __future__ import division
import time
import torch
import torch.nn as nn
from torch.autograd import Variable
import numpy as np
import cv2
from util import *
from darknet import Darknet
from preprocess import prep_image, inp_to_image
import pandas as pd
import random
import argparse
import pickle as pkl

def prep_image(img, inp_dim):
    # CNNに通すために画像を加工する
    orig_im = img
    dim = orig_im.shape[1], orig_im.shape[0]
    img = cv2.resize(orig_im, (inp_dim, inp_dim))
    img_ = img[:,:,::-1].transpose((2,0,1)).copy()
    img_ = torch.from_numpy(img_).float().div(255.0).unsqueeze(0)
    return img_, orig_im, dim

def write(x, img):
    # 画像に結果を描画
    c1 = tuple(x[1:3].int())
    c2 = tuple(x[3:5].int())
    cls = int(x[-1])
    label = "{0}".format(classes[cls])
    color = random.choice(colors)
    cv2.rectangle(img, c1, c2,color, 1)
    t_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_PLAIN, 1 , 1)[0]
    c2 = c1[0] + t_size[0] + 3, c1[1] + t_size[1] + 4
    cv2.rectangle(img, c1, c2,color, -1)
    cv2.putText(img, label, (c1[0], c1[1] + t_size[1] + 4), cv2.FONT_HERSHEY_PLAIN, 1, [225,255,255], 1);
    return img

def arg_parse():
    # モジュールの引数を作成
    parser = argparse.ArgumentParser(description='YOLO v3 Cam Demo') # ArgumentParserで引数を設定する
    parser.add_argument("--confidence", dest = "confidence", help = "Object Confidence to filter predictions", default = 0.25)
    # confidenceは信頼性
    parser.add_argument("--nms_thresh", dest = "nms_thresh", help = "NMS Threshhold", default = 0.4)
    # nms_threshは閾値
   
    parser.add_argument("--reso", dest = 'reso', help =
                        "Input resolution of the network. Increase to increase accuracy. Decrease to increase speed",
                        default = "160", type = str)
                        # resoはCNNの入力解像度で、増加させると精度が上がるが、速度が低下する。
    return parser.parse_args() # 引数を解析し、返す

if __name__ == '__main__':
    cfgfile = "cfg/yolov3.cfg" # 設定ファイル
    weightsfile = "yolov3.weights" # 重みファイル
    num_classes = 80 # クラスの数

    args = arg_parse() # 引数を取得
    confidence = float(args.confidence) # 信頼性の設定値を取得
    nms_thesh = float(args.nms_thresh) # 閾値を取得
    start = 0
    CUDA = torch.cuda.is_available() # CUDAが使用可能かどうか

    num_classes = 80 # クラスの数
    bbox_attrs = 5 + num_classes

    model = Darknet(cfgfile) #modelの作成
    model.load_weights(weightsfile) # modelに重みを読み込む

    model.net_info["height"] = args.reso
    inp_dim = int(model.net_info["height"])

    assert inp_dim % 32 == 0
    assert inp_dim > 32

    if CUDA:
        model.cuda() #CUDAが使用可能であればcudaを起動

    model.eval()

    cap = cv2.VideoCapture(0) #カメラを指定

    assert cap.isOpened(), 'Cannot capture source' #カメラが起動できたか確認

    frames = 0
    start = time.time()
    while cap.isOpened(): #カメラが起動している間

        ret, frame = cap.read() #キャプチャ画像を取得
        if ret:
            # 解析準備としてキャプチャ画像を加工
            img, orig_im, dim = prep_image(frame, inp_dim)

            if CUDA:
                im_dim = im_dim.cuda()
                img = img.cuda()

            output = model(Variable(img), CUDA)
            output = write_results(output, confidence, num_classes, nms = True, nms_conf = nms_thesh)

            # FPSの表示
            if type(output) == int:
                frames += 1
                print("FPS of the video is {:5.2f}".format( frames / (time.time() - start)))
                cv2.imshow("frame", orig_im)

                # qキーを押すとFPS表示の終了
                key = cv2.waitKey(1)
                if key & 0xFF == ord('q'):
                    break
                continue

            output[:,1:5] = torch.clamp(output[:,1:5], 0.0, float(inp_dim))/inp_dim
            output[:,[1,3]] *= frame.shape[1]
            output[:,[2,4]] *= frame.shape[0]

            classes = load_classes('data/coco.names') # 識別クラスのリスト
            colors = pkl.load(open("pallete", "rb"))

            list(map(lambda x: write(x, orig_im), output))

            cv2.imshow("frame", orig_im)
            key = cv2.waitKey(1)
            # qキーを押すと動画表示の終了
            if key & 0xFF == ord('q'):
                break
            frames += 1
            print("FPS of the video is {:5.2f}".format( frames / (time.time() - start)))

        else:
            break

このプログラムをAnaconda Promptを起動して、実行します。

python 上記のプログラムを入力したpyファイル.py

ブログにアップロードするため、画像が粗く小さくなってしまいましたが、ペットボトルがbottle（たまにcup）、テーブルがdining table、椅子がchairと検出・認識しています。
また手を差し出すと、personとして手が人の一部だとも認識しています。

私のWindowsタブレットではFPSが2程度であり、多少カクカクした動きになりました。
今回の処理はCPUのみで動作しており、GPUを搭載したマシン等を使用すればもっとスムーズな動画で処理が行えるかもしれませんね。

おわりに

今回は、YOLOv3を利用して物体検出を行いました。
リアルタイムでも使える早さで物体検出が出来ているのは驚きでした。
また、物体認識の精度が高く、物体を認識している様は、映画「ターミネーター」を思い出させますね。

次回は、画像から文字認識をさせるOCRを行ってみたいと思います。

Season 1