YOLO形式のアノテーションファイルをTFRecordに変換する

はじめに

YOLO形式のアノテーションファイルを、TFRecordに変換する必要があったので、その方法をメモします。Googleで検索すると以下のGitHubが検索でヒットするのですが、使ってみようとするとライブラリのインストールなど色々と詰まったので情報を共有します。
ドキュメントに不備があるほか、ツールの設定方法が洗練されていない印象を抱いたので、個人で使う用途に限定したほうがよいかもしれません。
https://github.com/mwindowshz/YoloToTfRecords

事前準備

以下のバージョンで変換を確認しました。TensorFlowはバージョン2.7とするのが良さそうです(執筆時の最新バージョン。2.5で実行したところエラーが発生)。バージョンの不整合を起こす場合、TensorFlow公式が推奨するとおりにvenvを利用するか、Dockerで環境を作るのが良いかもしれません。

• Ubuntu v.20.04
• Python v.3.9.7
• TensorFlow v.2.7
• Protocol Buffers v.3.14.0

PythonとTensorFlowは以下を参考にインストールします。
【Pythonのインストール】
https://www.python.jp/install/install.html

$ sudo apt update
$ sudo apt install build-essential libbz2-dev libdb-dev \
  libreadline-dev libffi-dev libgdbm-dev liblzma-dev \
  libncursesw5-dev libsqlite3-dev libssl-dev \
  zlib1g-dev uuid-dev tk-dev
$ wget "https://www.python.org/ftp/python/3.9.7/Python-3.9.7.tar.xz"
$ tar xJf Python-3.9.7.tar.xz
$ cd Python-3.9.7
$ ./configure
$ make
$ sudo make install

TensorFlowのインストール
https://www.tensorflow.org/install/pip?hl=ja

$ pip3 install tensorflow==2.7.0

YoloToTfRecordsをgitでチェックアウトします。

$ git clone https://github.com/mwindowshz/YoloToTfRecords.git ~/YoloToTfRecords

このスクリプトを実行するには、OpenCVとPascal Voc Writerというライブラリをインストールする必要があります。次のコマンドを実行することで、インストールすることができます。
※Pythonのバージョンに応じてパスを変更

pip3 install opencv-python pascal-voc-writer
export PYTHONPATH=$PYTHONPATH:/usr/local/lib/python3.9/site-packages

また、併せて次のライブラリをインストールしてください。

pip3 install jinja2 pandas pillow tf_slim scipy tensorflow.io

続けて、Protocol Buffersをインストールします。
※引用： http://google.github.io/proto-lens/installing-protoc.html

PROTOC_ZIP=protoc-3.14.0-linux-x86_64.zip
curl -OL https://github.com/protocolbuffers/protobuf/releases/download/v3.14.0/$PROTOC_ZIP
sudo unzip -o $PROTOC_ZIP -d /usr/local bin/protoc
sudo unzip -o $PROTOC_ZIP -d /usr/local 'include/*'
rm -f $PROTOC_ZIP

TensorFlow Model GardenをgitでチェックアウトしてObject Detection APIをインストールします。
※参考： https://qiita.com/wqwq/items/1a5c4e8671203d190fca

git clone https://github.com/tensorflow/models ~/models
cd ~/models/research/
protoc object_detection/protos/*.proto --python_out=.
export PYTHONPATH=$PYTHONPATH:`pwd`:`pwd`/slim

変換

0.すべての画像(jpg)とアノテーションデータ(txt)を一つのディレクトリにまとめておきます
※ここでは、/home/hoge/yoloにファイルを置くものとします

1. YOLO形式からPascal VOC形式のXMLファイルを生成する
変換ツールのpascal_voc_writer/__init__.pyを編集します。
137行目のclassesに、アノテーションのラベルの名前をリストで定義します。
続けて、変換対象のjpgファイルをフルパスで列挙したファイルを作成し、140行目のtrainFilesにそのファイルのフルパスを指定します。

trainFiles = "/home/hoge/all.txt" # for example

変換対象のjpgファイルのフルパスを列挙したファイルは、例えば下のコマンドを打てば作成できます。

find /home/hoge/yolo -name "*.jpg" > /home/hoge/all.txt

また、100行目のコードを確認します。コメントを逆にとると、'background'がない場合は加算すると結果が変わってしまうので、'+1'を削除する必要があります。

classID = int(row[0])+1 #we add 1 because in ssd 0 is backgroud

そして、pyファイルを(Run Codeと言っているので)実行します。

python3 __init__.py

実行が完了すると、/home/hoge/yoloにXMLファイルが生成されます。

2. CSVファイルを生成する
変換ツールのGenerateTFRecords/xml_to_csv.pyを編集します。55行目のsource_file_listに、1.で作成した画像のリストファイル(all.txt)をフルパスで指定してください。また、56行目のdest_csv_fileとして同じディレクトリを指定してCSVを出力するようにします。

source_file_list = "/home/hoge/all.txt" # for example
dest_csv_file = "/home/hoge/all.csv" # for example

そじて、pyファイルを実行します。

python3 xml_to_csv.py

3. TFRecordファイルを出力する
変換ツールのGenerateTFRecords/generate_tfrecord.pyを編集します。TensorFlow 1系のコードなのか、2系で実行しようとするとエラーが発生してしまいます。以下に示すように修正してください。

# 23行目
flags = tf.compat.v1.flags
# 43行目
with tf.io.gfile.GFile(os.path.join(path, '{}'.format(group.filename)), 'rb') as fid:
# 84行目
writer = tf.compat.v1.python_io.TFRecordWriter(FLAGS.output_path)
# 98行目
tf.compat.v1.app.run()

READMEに書いていないのですが、31行目のclass_text_to_int関数のdictの内容を、タグ付け内容に応じて書き換える必要があります。

dict = {'signal_blue':0, 'signal_yellow':1, 'signal_red':2} # for example

上で書いた修正ができたら、下のコマンドを実行してください。csv_inputは2で生成したCSVファイルへのパスです。

python3 generate_tfrecord.py --csv_input=/home/hoge/all.csv --output_path=/home/hoge/all.tfrecord

変換に成功すると以下のメッセージが出力されます。

Successfully created the TFRecords: /home/hoge/all.tfrecord

おわりに

控えめに言って、難しい。。今回紹介した修正をしても、今後は使えなくなってしまうかもしれません。あとは、正しく出力されているか、CSVファイルやTFRecordの出力結果を確認したほうが良さそう。今回は、学習で使用しようとしているTensorflowのバージョンが2.5である一方、Object Detection APIをインストールするのに最新のTensorflow(ver. 2.7)が必要となったため(エラーが発生)、最新のTensorFlow用の環境を別に用意する必要があり面倒でした。