【Python, darkflow】Webカメラの画像から人や猫を検出して切り抜いてみる

Darkflowというオープンソースの物体検出のライブラリを使って
Webカメラを使って人や猫を検出してその部分を切り抜くプログラムを紹介します。

人だけでも良かったのですが、私が猫が好きなので検出してみます。

ちなみにバックエンドはTensorFlowを使います。

バージョン等

  • Python 3.6.0
  • TensorFlow 1.7.0

2019年1月11日現在TensorFlowがバージョン2.0だとdarkflowは動かないので注意してください。

Darkflowとは(簡単に)

私も画像処理は専門ではないので、簡単に説明しますが、
Darkflowは物体検出物体認識をリアルタイムで行うことができるオープンソースのプログラムです。
以下のgithubのREADMEを見ると何ができるのかわかると思います。

https://github.com/thtrieu/darkflow

準備

TensorFlow 1系のインストール

先程もお話したようにTensorFlowは1系をインストールする必要があります。pipでバージョン指定してインストールします。

pip install tensorflow==1.7.0

darkflowのインストール

darkflowをインストールします。
まずgithubからdarkflowをcloneし、そのディレクトリに移動します。

git clone https://github.com/thtrieu/darkflow.git
cd darkflow

そして以下のコマンドでインストールします。

python setup.py build_ext --inplace

もし上記のコマンドでインストールできなかった時には次の2つの方法も試してください。

pip install -e .

or

pip install .

ディレクトリの作成と必要なデータの取得

ディレクトリ構成

まず、今回作るプログラムのディレクトリを作成しておきます。
ディレクトリ名をとりあえずfind_person_and_cat(適当に決めてください)にしておきます。

mkdir find_person_and_cat
cd find_person_and_cat

先に最終的なディレクトリ(find_person_and_cat)構成を示しておきます。

├── main.py # 実行用のプログラム
├── cfg     # 設定ファイル
│   ├── coco.names
│   └── yolo.cfg
└── weights  # 学習済みのモデル
    └── yolo.weights

上記のような構成になるようにもろもろ設定していきます。

設定ファイルの取得

設定ファイルや学習済みモデルのデータを保存しておく用のディレクトリを用意しておきます。

mkdir cfg

先程cloneしたdarkflowからdarkflow/cfg/yolo.cfgを今作成したcfgディレクトリに格納しておきます。(コピーなんかをしてください)

次にクラス名(分類する対象の名前)の一覧が書いてあるデータもダウンロードします。

wget https://raw.githubusercontent.com/pjreddie/darknet/master/data/coco.names

ルートディレクトリに戻っておきます。

cd ..

学習済みモデルの重みデータの取得

find_person_and_cat下にweightsというディレクトリを作っておきます。

mkdir weights
cd weights

そしてwgetコマンドで重みデータをダウンロードします。

wget https://pjreddie.com/media/files/yolov3.weights

ルートディレクトリに戻っておきます。

実装:人と猫を検出・認識して切り抜いてみる

では、Webカメラから画像データを取得して人と猫を見つけてその部分を切り抜いてみます。

先程のdarkflowの力を借りれば簡単にできます。

サンプルコード(main.py)を以下に示します。

# find_person_and_cat/main.py

import cv2
from darkflow.net.build import TFNet

# darkflowの設定
options = {"model": "./cfg/yolo.cfg",
           "load": "./weights/yolo.weights",
           "threshold": 0.1}

# darkflowの生成
tfnet = TFNet(options)

# カメラの起動
cap = cv2.VideoCapture(0)

class_names = ["cat", "person"]


def clip_img(img, result, img_path):
    """
        画像を切り抜く
    """
    x = result['topleft']['x']
    y = result['topleft']['y']
    brx = result['bottomright']['x']
    bry = result['bottomright']['y']
    h = abs(y-bry)
    w = abs(x-brx)

    clipped = img[y:y+h, x:x+w]
    cv2.imwrite(img_path, clipped)

    return clipped


def main():
    while(True):
        ret, img = cap.read()
        results = tfnet.return_predict(img)

        # 元画像の保存
        img_path = './img.jpg'
        # img_path = './{}.jpg'.format(datetime.datetime.now().strftime('%Y%m%d%H%M%S%f'))
        cv2.imshow("img", img)
        cv2.imwrite(img_path, img)
        for result in results:
            label = result['label']
            confidence = result['confidence']

            # 猫や人が含まれていて、確信度は0.6以上であれば切り抜く
            if label in class_names and confidence > 0.6:
                clipped_img_path = './img_clipped.jpg'
                # clipped_img_path = './{}_clipped.jpg'.format(datetime.datetime.now().strftime('%Y%m%d%H%M%S%f'))
                # 画像をクリップする
                clip_img(img, result, clipped_img_path)
                print("cliped")

        # qを押したら終了する
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()


if __name__ == '__main__':
    main()

ソースコードの詳しい説明はまた後日追記する予定です。
とりあえず指定したディレクトリ構成にし上記のコードを実行すると
Webカメラでひたすら撮影しながら、人か猫がいればその部分を切り抜いていきます。
(ひたすら画像を上書きしていきます。)

参考文献

https://yoheitaonishi.com/yolo9000/#i-3 https://arxiv.org/pdf/1612.08242.pdf https://pjreddie.com/darknet/yolo/
タイトルとURLをコピーしました