【Darknet】リアルタイムオブジェクト認識 YOLOをTensorflowで試す

  • 31
    いいね
  • 0
    コメント

この記事は Retty Advent Calendar 7日目です。

昨日は、のりぴーさん(@noripi )のJavaのプロダクトをKotlinに移行してみた話でした。

はじめに

DeepLearningのライブラリにはTensorflowやChainer、Caffeなど有名なものが複数ありますが、ここで扱うのはTensorflowと、意外と知られていないDarknetというFrameworkです。

Darknetを使うことでYOLO(You Only Look Once)というリアルタイムオブジェクト認識やDeepDreamのような画像加工、AlphaGoのような囲碁を試すことができます。
さて、今回はそんなDarknetのYOLO(オブジェクト認識)をTensorflowで試してみようというお話です。
※現在YOLOはv2が公開されていますが、今回Tensorflowで扱うものはv1です。
ありがたいことにYOLOで学習されたモデルをTensorflow用に変換するライブラリが公開されているのでこれを使います。

Darknetのインストール

Darknetの良いところは何といってもインストールがとても簡単かつ既に実用的で使いやすいモデルが用意されていることです。
公式サイトにも書いていますがインストールは以下の3行をターミナルで実行するだけです。

git clone https://github.com/pjreddie/darknet
cd darknet
make

簡単!

そしてDarknetはGPUにも対応していますがデフォルトではCPUのみとなっています、GPUに対応させるには以下のようにMakefileの1行目をGPU=0からGPU=1に書き換えます。

Makefile
GPU=1

ちなみに僕の実行環境ではGPUは使えないためCPU版で行なっていきます。

リアルタイムオブジェクト認識YOLO

モデルダウンロード

Darknetのインストールが完了したら早速YOLOを使ってみましょう。
まずは学習済みのモデル(.weights)ファイル(258 MB)をダウンロードします。
またはターミナルから以下を実行してもokです。
wget http://pjreddie.com/media/files/yolo.weights

実行

./darknet detect cfg/yolo.cfg yolo.weights data/dog.jpg
モデルをダウンロードし、上記の実行コマンドを入力するとpredictions.pngという画像が生成されオブジェクト認識されていることがわかります。
predictions.png

TensorflowでYOLO

Darknetのチュートリアルが終わったので本題です。
冒頭でも述べましたがありがたいことにYOLOで学習されたモデルをTensorflow用に変換するライブラリが公開されています。
早速これもcloneしちゃいましょう。
git clone https://github.com/gliese581gg/YOLO_tensorflow.git

モデルダウンロード

こちらも同様、Tensorflow用に既に変換されたYOLOのモデル(.ckpt)が公開されていますのでダウンロードします。(358 MB)
gitリポジトリにはYOLO_smallとYOLO_tinyの2つがありますが、YOLO_tinyは精度は落ちるけど速度が速いようです。(今回はYOLO_smallを使用)

実行

.ckptファイルをダウンロードしたら、YOLO_tensorflow/weights/に入れます。そして次のコマンドを実行すると結果が表示されます。
python YOLO_small_tf.py -fromfile test/person.jpg

↓のようにすることで結果画像を保存することもできます。
python YOLO_small_tf.py -fromfile test/person.jpg -tofile_img predictions.png
predictions.png

predictions.png

YOLO_smallモデルでは犬が認識されていませんね。

自分のデータで学習

ディープラーニングの楽しい部分はなんといっても自分で好きなものを学習させることができるところです。
YOLOで認識させたいものの学習も、またありがたいことに学習用のチュートリアルが作られています。
学習の方法についてはここでは詳しく述べませんが、BBox-Label-Toolというものを使い教師データを作成しトレーニング用のコマンドを実行するという形です。
このチュートリアルではstop標識とyield標識を学習させています。

学習されたモデルをTensorfrowで使えるものに変換

上記の学習で出来上がるモデルは(.weights)ファイルであり、このままではTensorflowでは使えません。変換にはTensorflow用に変換するライブラリを使います。
cloneするとYOLO_weight_extractor/YOLO_weight_extractor.tar.gzというファイルがあるので解凍します。
これを使ってstop標識とyield標識を学習させたモデル(.weights)ファイルを.ckptファイルに変換してみます。

ネットワーク構成の変更

モデルを変換する際、学習に使ったネットワーク構成と同じネットワークを使用しなければなりません。
YOLOではネットワーク構成がcfg/フォルダの中に入っています。上の学習の例ではこちらのネットワーク構成ファイルが使用されているのでダウンロードしてcfg/の中に入れます。

モデルのダウンロード

モデルについても学習後に完成したものが必要なのでダウンロード(※1.09 GBあります)
これはYOLO_weight_extractor/に置いておきます。

ソースコードの変更

デフォルトのソースではクラス数が20個になっていますのでstop標識とyield標識の2つに変更するため以下のように修正。

src/yolo.c(12行目)
char *voc_names[] = {"stopsign", "yieldsign"};
image voc_labels[2];
src/yolo.c(352行目)
draw_detections(im, l.side*l.side*l.n, thresh, boxes, probs, voc_names, voc_labels, 2);
src/yolo.c(412行目)
for(i = 0; i < 2; ++i){
        char buff[256];
        sprintf(buff, "data/labels/%s.png", voc_names[i]);
        voc_labels[i] = load_image_color(buff, 0, 0);
    }

修正が完了したらターミナルでmakeコマンドを実行します。

変換モデル作成コードの変更

変換するためのコードもデフォルトのままではクラス数が違いますので変更します。

YOLO_full_builder.py(47行目)
self.fc_31 = self.fc_layer(31,self.fc_29,1452,flat=False,linear=True)

デフォルトでは1470となっている部分を1452に修正するのですがこの数はcfgファイルの225行目に設定されているsideという変数、226行目のnum、クラス数(classes)の3つで決まる数です。計算式は(5*num + classes)*(side*side)となります。
今回扱うネットワークモデルではside=11、num=2、classes=2となっているので(5*2+2)*(11*11)=1452となります。

またネットワーク構成が保存されるディレクトリがデフォルトでは作者のものになっているので修正します。

YOLO_full_builder.py(6行目)
weights_dir = 'cjy/'

ラベルの作成

学習されたラベルデータの登録も必要となります。
data/labels/の中にmake_labels.pyというコードがあるのでこのコードを開くと

make_labels.py
l = ["person","bicycle","car", .....]

という部分がありますのでこの配列に学習したラベルを追加します、今回は"stopsign""yieldsign"の2つを追加します。
追加したらpython make_labels.pyコマンドで実行します。
するとlabelsフォルダの中にstopsign.pngとyieldsign.pngが出来上がります。

モデル変換

やっと準備が整いました。コード類の修正が終わったらようやくモデルの変換作業です。
変換はとても簡単です。まずはターミナルにて以下のコマンドを実行
./darknet yolo test cfg/yolo_2class_box11.cfg yolo_2class_box11_3000.weights
実行ししばらく待つとEnter Image Path:という文が表示されます、この文が表示されれば完了なのでctrl+cで終了して構いません。
完了するとcjy/フォルダにネットワーク構成用のファイルが複数出来上がります。
次にpython YOLO_full_builder.pyコマンドを実行すると.ckptファイルが出来上がります。

Tensorflow用に変換したモデルでYOLOを試す

いよいよ最後です。
用意されたモデルではデフォルトのままtensorflowで試すことができますが新しく作ったモデルを試すために修正する部分があります。
またYOLO_full_builder.pyで作成されたモデルを試すコードが用意されてないようなのでYOLO_small_tf.pyを修正します。
修正部分は以下

YOLO_small_tf.py(15行目)
weights_file = 'weights/YOLO_full.ckpt'
YOLO_small_tf.py(22行目)
classes = ["stopsign","yieldsign"]
YOLO_small_tf.py(74行目)
self.fc_29 = self.fc_layer(29,self.conv_28,4096,flat=True,linear=False)
YOLO_small_tf.py(75行目)
self.fc_32 = self.fc_layer(30,self.fc_29,1452,flat=False,linear=True)
YOLO_small_tf.py(77行目)
# self.fc_32 = self.fc_layer(32,self.fc_30,1470,flat=False,linear=True) コメントアウト
YOLO_small_tf.py(155行目)
def interpret_output(self,output):
    side = 11
    num = 2
    class_num = 2
    print "class_num:",class_num,"side:",side
    probs = np.zeros((side,side,num,class_num))
    class_probs = np.reshape(output[0:side*side*class_num],(side,side,class_num))
    scales = np.reshape(output[side*side*class_num:side*side*class_num+side*side*num],(side,side,num))
    boxes = np.reshape(output[side*side*class_num+side*side*num:],(side,side,num,4))
    offset = np.transpose(np.reshape(np.array([np.arange(side)]*side*num),(num,side,side)),(1,2,0))

    boxes[:,:,:,0] += offset
    boxes[:,:,:,1] += np.transpose(offset,(1,0,2))
    boxes[:,:,:,0:2] = boxes[:,:,:,0:2] / (side*1.0)
    boxes[:,:,:,2] = np.multiply(boxes[:,:,:,2],boxes[:,:,:,2])
    boxes[:,:,:,3] = np.multiply(boxes[:,:,:,3],boxes[:,:,:,3])

    boxes[:,:,:,0] *= self.w_img
    boxes[:,:,:,1] *= self.h_img
    boxes[:,:,:,2] *= self.w_img
    boxes[:,:,:,3] *= self.h_img

    for i in range(2):
        for j in range(class_num):
            probs[:,:,i,j] = np.multiply(class_probs[:,:,j],scales[:,:,i])

お疲れ様です。これで全ての準備が整いました。さっき出来上がったYOLO_full.ckptファイルをYOLO_tensorflow/weights/に入れて
コマンドを実行
python YOLO_small_tf.py -fromfile test/stopsign.JPEG -tofile_img predictions.png
stopsign.JPEGは自分で用意します。
predictions3.png

まとめ

いかがでしたでしょうか?
Tensorflow用に変換するのは大変ですがわざわざ変換せずとも簡単に試すことができるので是非試して遊んでみて下さい!
また今回深くは触れなかった学習についても試してみると面白そうです!

この投稿は Retty Inc. Advent Calendar 20167日目の記事です。