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

# 目的

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

# 背景

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

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

https://github.com/ultralytics/yolov5/releases/tag/v6.1

対応形式

- 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`で用意されている学習済みモデルを指定しました。<br>自身で学習済みのモデルを指定する場合、`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によりインストール方法も異なってくるため、本投稿ではインストール方法については割愛とさせていただきます。

https://nodejs.org/ja/

```
cd tfjs-yolov5-example
npm install
```

### 2.3.モデル配置

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

### 2.4.Webアプリ起動

Webアプリを起動します。

```
npm start
```

こんな感じのページが立ち上がります。
点線の枠内に画像ファイルをドラッグドロップします。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/38724/1d6325e5-53f3-8dff-1455-e3b3b1cba904.png)

:::note warn
しばらく待っても点線の枠内に`Loading model...`が表示されたままの場合、モデルが読み込まれていない場合があります。`web_model`フォルダの位置が正しいか確認してください。
:::
検出が終わるとこのような感じで表示されます。うまくいってますね！！
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/38724/7b7d4d86-d560-bc90-7204-4af1bd1acdb3.png)

# ソースコード

https://github.com/MamedenQ/Yolov5ModelExportExample

ノートブックしかないので、下記リンクを叩くでも良いです。
<a href="https://colab.research.google.com/github/MamedenQ/Yolov5ModelExportExample/blob/main/YOLOv5%E3%83%A2%E3%83%87%E3%83%AB%E3%82%92%E5%88%A5%E3%81%AA%E5%BD%A2%E5%BC%8F%E3%81%AE%E3%83%A2%E3%83%87%E3%83%AB%E3%81%AB%E5%A4%89%E6%8F%9B.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a>

# 最後に

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

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

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