TensorFlow.js をサクッと試す

# はじめに
JavaScriptからTensorFlowを使える TensorFlow.js というものがあります。
https://www.tensorflow.org/js
https://github.com/tensorflow/tfjs

公式サンプルも色々あるのですが、シンプルにTensorFlow (Python) で学習したモデルを読み込んで、それを使ってブラウザ側で推論する部分を作ってみました。

# 検証環境

## 学習
- Google Colaboratory
    - TensorFlow 2.6.0
    - TensorFlow.js 3.8.0

## 推論
- Firefox 91.0
    - TensorFlow.js 3.8.0

# モデル学習

これは普通にやれば良いのですが、保存するときにTensorFlow.jsから読み込める形式にする必要があります。というわけで最初にpipパッケージをセットアップしておきます。

```
!pip install tensorflowjs
```

MNISTの手書き数字分類をするとして、データを準備。

```python
import numpy as np
import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense
from tensorflow.keras.datasets import mnist

# データセットの準備
(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train = x_train[:, :, :, np.newaxis].astype("float32") / 255.0  # (60000, 28, 28, 1)
x_test = x_test[:, :, :, np.newaxis].astype("float32") / 255.0    # (10000, 28, 28, 1)
```

続いて適当なモデルを作って学習させます。サンプルなので10エポックくらいでよいでしょう。GPUインスタンスを使えばすぐ終わります。

```python
# モデルの準備
model = Sequential([
  Conv2D(32, kernel_size=(3, 3), activation="relu", padding="same", input_shape=(28, 28, 1)),
  MaxPooling2D(pool_size=(2, 2)),
  Conv2D(64, kernel_size=(3, 3), activation="relu", padding="same"),
  MaxPooling2D(pool_size=(2, 2)),
  Flatten(),
  Dense(128, activation="relu"),
  Dense(10, activation="softmax")
])
model.compile("Adam", loss="sparse_categorical_crossentropy", metrics="sparse_categorical_accuracy")
# 学習
model.fit(x_train, y_train, batch_size=256, epochs=10, validation_data=(x_test, y_test))
```

モデルをTensorFlow.jsで使える形式に保存します。

```python
import tensorflowjs as tfjs
tfjs.converters.save_keras_model(model, "./tfjs_model")
```

# モデルのダウンロード

Google Drive経由でダウンロードします。まずはGoogle Driveをマウントします。

「ドライブをマウント」を押し、確認ダイアログで「Googleドライブに接続」を押すと
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/204084/b2093c81-81c9-6265-ee96-fcc7b059dc6e.png)

このように ``drive/MyDrive`` というディレクトリが利用できるようになります。ここにファイルをコピーすると、Google Drive経由でモデルをダウンロードできるようになります。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/204084/fcd058e8-188f-ac05-1e44-c50b6157f9ce.png)

以下のコマンドでモデルをコピーしたら、あとはブラウザで[Google Drive](https://drive.google.com/)にアクセスし、モデルをダウンロードしましょう。

```
!cp -r ./tfjs_model ./drive/MyDrive
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/204084/7df786c3-ac9d-8d85-6e54-b2c4fff207b6.png)

参考: [ColaboratoryでのGoogle Driveへのマウントが簡単になっていたお話 - Qiita](https://qiita.com/kado_u/items/45b76f9a6f920bf0f786)

# 推論

今度はクライアントのHTMLとJavaScriptを書いていきます。``tfjs.html`` という名前のファイルを以下の内容で作成します。
※ブラウザのセキュリティ対策のため、ローカル環境 (``file:///``) では動きません。XAMPPでもレンタルサーバーでも何でもよいので、 ``http://`` でアクセスできる場所を準備しましょう。

**（後述: 実はもう少し簡単に書けます）**

```tfjs.html
<!DOCTYPE html>
<html lang="ja">
  <head>
    <meta charset="utf-8">
  </head>

  <body>
    <div><img src="sample_5.png" id="input_image"></div>
    <div>Prediction: <span id="pred"></span></div>

    <!-- Load TensorFlow.js -->
    <script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs/dist/tf.min.js"></script>

    <!-- Place your code in the script tag below. You can also use an external .js file -->
    <script>
      // Notice there is no 'import' statement. 'tf' is available on the index-page
      // because of the script tag above.

      // https://blog.codecamp.jp/programming-javascript-Imagedata
      function createImageData(img) {
        var cv = document.createElement('canvas');
        cv.width = img.naturalWidth;
        cv.height = img.naturalHeight;
        var ct = cv.getContext('2d');
        ct.drawImage(img, 0, 0);
        var data = ct.getImageData(0, 0, cv.width, cv.height);
        return data;
      }

      async function run(){
        const model = await tf.loadLayersModel('./tfjs_model/model.json');
        const image_data = createImageData(document.getElementById('input_image'));
        const x = tf.tensor1d(new Float32Array(image_data.data)).gather(tf.range(0, 28*28*4, 4, 'int32')).div(tf.scalar(255)).reshape([1, 28, 28, 1]);
        const y = model.predict(x); // [1, 10]
        document.getElementById('pred').textContent = y.reshape([-1]).argMax().arraySync();
      }
      run();
    </script>
  </body>
</html>
```

以下の各ページを参考にしました。

- [tensorflow/tfjs: A WebGL accelerated JavaScript library for training and deploying ML models.](https://github.com/tensorflow/tfjs) の "Getting Started" のコード
- [TensorFlow.js API](https://js.tensorflow.org/api/latest/)
- [【javaScript学習】ImageDataへの画像読み込みと画像処理 | CodeCampus](https://blog.codecamp.jp/programming-javascript-Imagedata): 画像ファイルの読み込み方法

HTMLと同じ場所に、先程ダウンロードしたモデル（``tfjs_model`` ディレクトリ）を置いておきます。
また、ペイントなどのソフトで適当に手書きした数字（背景色を黒、文字を白にしてください）を保存し、同じ場所に置きます。ここでは **5** を書いて ``sample_5.png`` という名前で保存しました。
![sample_5.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/204084/587fd621-e5b6-2b5c-0418-40d82d5a95cf.png)

## 推論コード

``run()`` の部分を抜き出してみます。

```javascript
      async function run(){
        const model = await tf.loadLayersModel('./tfjs_model/model.json');
        const image_data = createImageData(document.getElementById('input_image'));
        const x = tf.tensor1d(new Float32Array(image_data.data)).gather(tf.range(0, 28*28*4, 4, 'int32')).div(tf.scalar(255)).reshape([1, 28, 28, 1]);
        const y = model.predict(x); // [1, 10]
        document.getElementById('pred').textContent = y.reshape([-1]).argMax().arraySync();
      }
```

1. ``tf.loadLayersModel`` を用いて学習済みモデルを読み込む
2. 画像を ``ImageData`` （中にピクセルデータが入っている）に変換する
3. TensorFlowの入力形式に変換する
4. 推論する
5. 結果を表示する

という流れです。特に3の部分にかなり引っかかるポイントがあります。

### TypedArray → Tensor

``image_data.data`` が ``Uint8ClampedArray`` 型なので、リファレンスを見るとそのまま ``Tensor`` に変換できそうですが、できません。実は ``Float32Array`` に一度変換しておく必要があります。

### ``image_data.data`` のデータ構造とスライシング

``RGBARGBA...`` のように1ピクセルが4バイトで表されています。今回はグレースケール入力ですから、Rの値だけを取り出すことにします。それが以下の ``gather`` の部分です。

```javascript
tf.tensor1d(...).gather(tf.range(0, 28*28*4, 4, 'int32'))
```

要するにNumPyやTensorFlowでいうところの ``...[0::4]`` をやりたいのですが、どうもTensorFlow.jsではこのような書き方はできないようなので、上のように取り出したいインデックスの配列を作って ``gather`` で取り出します。``tf.range()`` の第3引数に ``int32`` を指定しないと、float32の配列になってしまってうまくいきません。

### ブロードキャスト

モデルは入力ピクセル値として0～1のfloat32型を取るようになっているので、各ピクセル値を255で割るのですが、これがまた一工夫必要です。割り算には ``div()`` メソッドを使い、除数に ``tf.scalar()`` を使用しています。

```javascript
tf.tensor1d(...).gather(...).div(tf.scalar(255))
```

NumPyやTensorFlowのノリで以下のようにしたくなるところですが、うまくいきません。

```javascript
tf.tensor1d(...).gather(...) / 255      // NaN
tf.tensor1d(...).gather(...).div(255)   // 先頭の要素を255で割った結果1つしか返ってこない
```

参考: [ディープラーニングのお勉強～その13。TensorFlow.jsでMNISTリアルタイム推論してみる | mgo-tec電子工作](https://www.mgo-tec.com/blog-entry-tensorflow-js-m5cam01.html)

## ブラウザで動作確認

``tfjs.html`` にブラウザからアクセスします。数秒間待って「Prediction: 5」と出れば成功です。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/204084/3f84c742-c3d6-5a8e-a8dd-6c1b83d712d7.png)

# 推論（もう少し簡単に）

実はCanvas APIを介さずに ``<img>`` タグの画像を直接Tensorに変換できるようです。

```tfjs_2.html
<!DOCTYPE html>
<html lang="ja">
  <head>
    <meta charset="utf-8">
  </head>

  <body>
    <div><img src="sample_5.png" id="input_image"></div>
    <div>Prediction: <span id="pred"></span></div>

    <!-- Load TensorFlow.js -->
    <script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs/dist/tf.min.js"></script>

    <!-- Place your code in the script tag below. You can also use an external .js file -->
    <script>
      // Notice there is no 'import' statement. 'tf' is available on the index-page
      // because of the script tag above.

      async function run(){
        const model = await tf.loadLayersModel('./tfjs_model/model.json');
        const x = tf.browser.fromPixels(document.getElementById('input_image'), 1).reshape([1, 28, 28, 1]).div(tf.scalar(255));
        const y = model.predict(x); // [1, 10]
        document.getElementById('pred').textContent = y.reshape([-1]).argMax().arraySync();
      }
      run();
    </script>
  </body>
</html>
```

先程の画像読み込み部分のコードが、随分スッキリしました。

```javascript
tf.browser.fromPixels(document.getElementById('input_image'), 1)
```

これだけで、``<img>`` タグの画像が ``[height, width, 1]`` のTensorとして取得できます。デフォルトではチャンネル数が3 (RGB) で取得されますが、引数に ``1`` を指定することで最初の1チャンネル (R) だけを得ることができます。
あとは ``reshape`` でサンプルの次元を付け足して、255で割るだけ。簡単ですね。