YOLOでの学習の練習にMNISTの教師データを自前で作って学習させてみました。

必要なソフトのインストール

Kerasのインストール

KerasがPython環境に入っていない場合は入れてください。

pip install tensorflow
pip install keras

darknet YOLOのインストール

darknet YOLOのリポジトリを取得する

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

ビルドなど、うまくいかない場合はdarknet YOLOの公式サイト参照。

pjreddie.com

学習データの作成

github.com

学習データを作成するためのスクリプトを作り、GitHubにあげました。まずは、このリポジトリを取得。

git clone https://github.com/uchidama/MNIST-TrainDataForYOLO.git

とか。

[1] MNISTの画像とラベルの作成

python mnist_to_jpg_and_label.py

画像とラベルの入ったテキストファイルが7万枚ずつ作成されます。

[2] train.txt と test.txtの作成

python generate_train_txt_and_test_txt.py

訓練用の画像一覧 train.txt とテスト用の画像一覧 test.txtが作成されます。以上で学習データの作成は終了です。

YOLOでの学習

[3] リポジトリのcfgディレクトリ、dataディレクトリ内のファイルをdarknetにコピー

cp cfg/tiny-yolo-mnist.cfg <darknet_dir>/cfg
cp cfg/voc-mnist.data <darknet_dir>/cfg
cp data/voc-mnist.names <darknet_dir>/data

tiny-yolo-mnist.cfg の上の方にあるbatch, subdivisionsをコメントアウトして切り替えておく

[net]
# Training
batch=64
subdivisions=2
# Testing
#batch=1
#subdivisions=1

[4] <darknet_dir>/cfg/voc-mnist.dataを編集。2で作成したtrain.txt、test.txtにパスを通す

train  = <path-to-mnist-train>/train.txt
valid  = <path-to-mnist-test>/test.txt

[5] 学習に使うトレーニング済のウェイトをダウンロードする

cd <darknet_dir>
wget https://pjreddie.com/media/files/darknet19_448.conv.23

[6] 学習モデルの保存ディレクトリを作成

mkdir backup

[7] 学習する

./darknet detector train cfg/voc-mnist.data cfg/tiny-yolo-mnist.cfg darknet19_448.conv.23

backupディレクトリの中に学習したモデル(weights)が保存される。

学習したモデルで識別する

./darknet detector test <data file> <cfg file> <weights> <predict image>

コマンドとしては、こういう指定。

./darknet detector test cfg/voc-mnist.data cfg/tiny-yolo-mnist.cfg weights/tiny-yolo-mnist_500000.weights ~/MNIST-TrainDataForYOLO/JPEGImages/60015.jpg

なんで、こんな感じで識別できます。

学習済のモデル

学習すると時間かかるんで、学習済みの重みもGitHubに置いときました。 https://github.com/uchidama/MNIST-TrainDataForYOLO/blob/master/weights/tiny-yolo-mnist_500000.weights

参考