目的

YOLOv5で学習したモデルを別形式に変換する方法を身につける。
TensorFlow.js形式に変換しWebアプリ上で物体検出する。

背景

YOLOv5のソースコードを見てニヤニヤしていたところ、モデル形式を変換するソースコードを見つけた。

下記のリリースノートには11の形式に公式に対応しましたよとの記述があり、性能や精度の検証が目的とのこと。

対応形式

PyTorch
TorchScript
ONNX
OpenVINO
TensorRT
CoreML
TensorFlow SavedModel
TensorFlow GraphDef
TensorFlow Lite
TensorFlow Edge TPU
TensorFlow.js

今のところ性能や精度の測定には興味はないが、TensorFlow.jsで読み込める形式にもエクスポートができるため、Webアプリに組み込めたら夢が広がるなと思った。

実行環境

Google Colab
nodejs

内容

ソースコードについてはGitHubにアップしてあります。
GitHubへのリンクは本投稿後半にあるソースコードの章を参照願います。

1.モデルをTensorFlow.js形式に変換

この章はGoogle Colab、または、お使いのパソコンでの操作となります。

変換元のモデルはYOLOv5で用意してある学習済みモデルを利用します。

YOLOv5のインストール

!git clone https://github.com/ultralytics/yolov5

%cd yolov5
!pip install -r /content/yolov5/requirements.txt

モデルを変換

export.pyを実行しモデルを変換します。

!python export.py --weights yolov5s.pt --include tfjs

yolov5フォルダの直下にyolov5s_web_modelというフォルダが出来上がっているので、フォルダごとローカルへ保存しておいてください。

他にもyolov5s_saved_modelというフォルダが出来上がってますが、これはTensorFlow.js形式へ変換する際に別の形式を経由しているために出力された中間ファイルとなります。
TensorFlow.jsモデルへの変換のみ、内部の変換の流れを紹介すると、下記の流れで変換されているようです。
PyTorch＞TensorFlow SavedModel＞TensorFlow GraphDef＞TensorFlow.js
詳しく知りたい方はexport.pyの中身を覗いてみてください。

今回指定したオプションについての説明です。
あまり細かなことまで調べきれていないため、概要程度の説明となります。
別な投稿にて詳しい解説を予定しています。

オプション名	指定値	説明
--weights	yolov5s.pt	変換する元のモデルを指定します。今回は`YOLOv5`で用意されている学習済みモデルを指定しました。自身で学習済みのモデルを指定する場合、`runs/train/exp/weights/best.pt`のようにパスを指定します。
--include	tfjs	変換後形式を指定します。指定可能な値は次の表で解説します。

--includeに指定可能な値は下記です。
ご希望の変換後形式に対応した指定値を--includeにご指定願います。

変換後形式	指定値
TorchScript	torchscript
ONNX	onnx
OpenVINO	openvino
TensorRT	engine
CoreML	coreml
TensorFlow SavedModel	saved_model
TensorFlow GraphDef	pb
TensorFlow Lite	tflite
TensorFlow Edge TPU	edgetpu
TensorFlow.js	tfjs

2.Webアプリに組み込み物体検出

ここからはお使いのパソコンでの操作となります。

2.1.Webアプリ生成

export.pyの冒頭で紹介されているお試しWebアプリがあるので、それをクローンします。

git clone https://github.com/zldrobit/tfjs-yolov5-example.git

2.2.依存ライブラリインストール

クローンしたプロジェクトに移動し、依存ライブラリをインストールします。
インストールとアプリの実行にはnodejsが必要となります。お使いの環境に合わせてインストール願います。
OSによりインストール方法も異なってくるため、本投稿ではインストール方法については割愛とさせていただきます。

cd tfjs-yolov5-example
npm install

2.3.モデル配置

先程変換したモデルを配置します。
クローンしたWebアプリフォルダにあるpublicフォルダ直下に、先程変換したモデルのフォルダyolov5s_web_modelをフォルダごとコピーし、フォルダ名をweb_modelと変更してください。

2.4.Webアプリ起動

Webアプリを起動します。

npm start

こんな感じのページが立ち上がります。
点線の枠内に画像ファイルをドラッグドロップします。

しばらく待っても点線の枠内にLoading model...が表示されたままの場合、モデルが読み込まれていない場合があります。web_modelフォルダの位置が正しいか確認してください。

検出が終わるとこのような感じで表示されます。うまくいってますね！！

ソースコード

ノートブックしかないので、下記リンクを叩くでも良いです。

最後に

コマンド実行するのみで、コードを一切書かずにモデルの変換を行うことができて驚いた。
以前PytorchのモデルをTensorFlow.jsで読み込める形式に変換しようと調べていた時はエラー続きで結局断念していた記憶があったが、何も困ることなく進めることができた。

webアプリに組み込めたことで夢が広がりそうだなとも思ったが、アプリで動作させるのはちょっと重たい。
TensorFlow.js側の話になってしまうが、この辺がもう少し素早く動いてくれるようになれば良いと思った。やはり同居は難しいか。

ともあれ、形式を変換する方法は身についたのでよしとする。