TensorFlow 公式の Pix2Pix チュートリアルで作ったモデルを TensorFlow.js + p5.js で動かす


[affinelayer/pix2pix-tensorflow](https://github.com/affinelayer/pix2pix-tensorflow) の [export-checkpoint.py](https://github.com/affinelayer/pix2pix-tensorflow/blob/master/server/tools/export-checkpoint.py) を動かすのに一苦労していたところ、[公式の Pix2Pix チュートリアル]
(https://www.tensorflow.org/tutorials/generative/pix2pix) を見つけたので、それで学習したモデルを [TensorFlow.js](https://www.tensorflow.org/js/guide) で**とりあえず**動かした時のメモです。
![output-palette.gif](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/693216/3496586e-c704-2bcb-6cfb-877fbfd31d98.gif)

# 1. データ作成

オリジナルデータではなく [CMP Facade](http://cmp.felk.cvut.cz/~tylecr1/facade/) を使うのであれば不要です。

<details><summary>データ作成する話</summary><div>

## 1.1. データの置き方

対となる 2 枚の 256 x 256 の画像を準備して、1 枚(512 x 256 の画像)にします。CMP Facade で合計 606 画像なので、同じぐらいはあった方が良さそうです。右側の画像から左側の画像を推測するので、注意です。データの置き方は [CMP Facade](http://cmp.felk.cvut.cz/~tylecr1/facade/) を真似て、

- facades/
    - test/
        - 0001.jpg
        - 0002.jpg
        - ...
    - train/
        - 0001.jpg
        - 0002.jpg
        - ...
    - val/
        - 0001.jpg
        - 0002.jpg
        - ...

上記のように、facades ディレクトリの下に train/ val/ test/ ディレクトリがあり、それぞれの下に画像ファイルを置くようにします(楽なため)。[CMP Facade](http://cmp.felk.cvut.cz/~tylecr1/facade/) のファイル数が 106 : 400 : 100 なので、だいたい同じような比率で分けます。

## 1.2. データの作り方

ここは変換したい内容によって千差万別ですが、福笑いを試した時の手順です。
[例によって](https://qiita.com/tset-tset-tset/items/53bfbaa2c2ed0dc77947)、OpenVINO で動画から顔画像とランドマークを取ってきて、適当につなげています。

```bash
cd /cygdrive/c/Users/${USER}/Downloads/output/

## 1 fps の動画にしておく
ls -l input.mp4
ffmpeg -i input.mp4 -r 1 input_r1.mp4 -y

## OpenVINO で raw データを取得
echo 'source ${INTEL_OPENVINO_DIR}/bin/setupvars.sh

cd ${INTEL_CVSDK_DIR}/inference_engine/demos/
sed -i "s/*)/interactive_face_detection_demo)/g" CMakeLists.txt
./build_demos.sh

${INTEL_CVSDK_DIR}/deployment_tools/tools/model_downloader/downloader.py \
  --name face-detection-adas-0001,facial-landmarks-35-adas-0002 \
  --output_dir /content/model/ \
  --precisions FP32

echo `date`: start detection

/root/omz_demos_build/intel64/Release/interactive_face_detection_demo \
  -i /Downloads/output/input_r1.mp4 \
  -m /content/model/intel/face-detection-adas-0001/FP32/face-detection-adas-0001.xml \
  -m_lm /content/model/intel/facial-landmarks-35-adas-0002/FP32/facial-landmarks-35-adas-0002.xml \
  -no_show \
  -no_wait \
  -async \
  -r > /Downloads/output/raw_landmark.txt' | docker run -v /c/Users/${USER}/Downloads:/Downloads -u root -i --rm openvino/ubuntu18_dev:2020.4

## raw データから、ffmpeg + convert コマンドを作成
perl -ne '$i++ if m{^\[0,1\]}; printf("ffmpeg -loglevel error -ss ".($i-1)." -i input_r1.mp4 -vframes 1 -vf crop=$3:$4:$1:$2 %05d.jpg -y\n", ++$j), $s = $3 if m{\((\d+),(\d+)\)-\((\d+),(\d+)\) WILL BE RENDERED}; push @h, [$1, $2] if m{^(\S+), (\S+)}; if ($#h == 34) { if ($s) { printf(q(convert -size %dx%1$d canvas:white -stroke black -strokewidth 5 -draw "line %.0d,%.0d %.0d,%.0d" -draw "line %.0d,%.0d %.0d,%.0d" -draw "line %.0d,%.0d %.0d,%.0d" -draw "line %.0d,%.0d %12$.0d,%13$.0d" out.jpg; convert +append %05d.jpg out.jpg %16$05d_face.jpg) . "\n", $s, $h[0][0]*$s, $h[0][1]*$s, $h[1][0]*$s, $h[1][1]*$s, $h[2][0]*$s, $h[2][1]*$s, $h[3][0]*$s, $h[3][1]*$s, $h[8][0]*$s, $h[8][1]*$s, ($h[10][0]+$h[11][0])*$s/2, ($h[10][1]+$h[11][1])*$s/2, $h[9][0]*$s, $h[9][1]*$s, $j); } @h = (); $s = 0; } ' raw_landmark.txt > ffmpeg_convert.sh
head ffmpeg_convert.sh
  # ffmpeg -loglevel error -ss 0 -i input_r1.mp4 -vframes 1 -vf crop=333:333:505:359 00001.jpg -y
  # convert -size 333x333 canvas:white -stroke black -strokewidth 5 -draw "line 123,173 85,161" -draw "line 180,189 217,191" -draw "line 117,262 150,283" -draw "line 186,280 150,283" out.jpg; convert +append 00001.jpg out.jpg 00001_face.jpg
  # ffmpeg -loglevel error -ss 1 -i input_r1.mp4 -vframes 1 -vf crop=333:333:505:359 00002.jpg -y
  # convert -size 333x333 canvas:white -stroke black -strokewidth 5 -draw "line 123,174 85,162" -draw "line 179,189 216,191" -draw "line 117,262 150,284" -draw "line 186,280 150,284" out.jpg; convert +append 00002.jpg out.jpg 00002_face.jpg
  # ...
sh ffmpeg_convert.sh
mogrify -resize 512x256! *_face.jpg
```
![00004_face.jpg](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/693216/25cfe686-d71b-211b-ca0f-9fc82cadc00f.jpeg)

上記のような画像がたくさんできます。

</div></details>

# 2. TensorFlow で学習

TensorFlow 公式にある、[Pix2Pix  |  TensorFlow Core](https://www.tensorflow.org/tutorials/generative/pix2pix) チュートリアルにある[ノートブック](https://github.com/tensorflow/docs/blob/master/site/en/tutorials/generative/pix2pix.ipynb)を使って学習を行います。2020/10/07 現在、Google Colab 上の TensorFlow が 2.3.0 なので、2.3.0 で学習しています。

## 2.1. 学習開始まで

[ノートブック](https://github.com/tensorflow/docs/blob/master/site/en/tutorials/generative/pix2pix.ipynb) を上から順に実行していきます。Google Drive をマウントしておきます。

### 2.1.1. 「Load the dataset」でデータをロード

自分のデータを使う際は、ここを書き換えておきます。

```bash
～～
!cd /content/; unzip /content/drive/My\ Drive/my_facades.zip

#PATH = os.path.join(os.path.dirname(path_to_zip), 'facades/')
PATH = '/content/facades/'
```

### 2.1.2. 「Build the Generator」の [BatchNormalization](https://www.tensorflow.org/api_docs/python/tf/keras/layers/BatchNormalization)

実際に JS 側で実行されることになる Generator 作成「Build the Generator」ですが、ここでハマりました。Generator 作成時、下記関数内で [BatchNormalization](https://www.tensorflow.org/api_docs/python/tf/keras/layers/BatchNormalization) レイヤが追加されるのですが、[TensorFlow.js](https://www.tensorflow.org/js/guide) 2.4.0 でのここの動きが追い切れていません。
まずそのままでは、[tensorflowjs_converter](https://github.com/tensorflow/tfjs/tree/master/tfjs-converter) で未対応の input があるらしく変換したモデルを JS で predict した際に、結果が全 NaN になって動きませんでした。
`fused=False` を渡すと [tensorflowjs_converter](https://github.com/tensorflow/tfjs/tree/master/tfjs-converter) は通り、predict で数値が返るようにはなります。が、TensorFlow 2.3.0 の [BatchNormalization](https://www.tensorflow.org/api_docs/python/tf/keras/layers/BatchNormalization) と [TensorFlow.js](https://www.tensorflow.org/js/guide) 2.4.0 の [batchNormalization](https://js.tensorflow.org/api/2.4.0/#layers.batchNormalization) の相性なのか、Python のノートブックの時と全く違うノイズのような結果が返ってきます。
その原因はわからず…　将来的に [TensorFlow.js](https://www.tensorflow.org/js/guide) で修正されるのか、既にパッチがあるのかも不明ですが、**とりあえず**現行で動かすために不本意ながら該当箇所をコメントアウトして回避しています。(品質は落ちるので悩ましいですが…)

```python
def downsample(filters, size, apply_batchnorm=True):
  initializer = tf.random_normal_initializer(0., 0.02)

  result = tf.keras.Sequential()
  result.add(
      tf.keras.layers.Conv2D(filters, size, strides=2, padding='same',
                             kernel_initializer=initializer, use_bias=False))

#  if apply_batchnorm:
#    result.add(tf.keras.layers.BatchNormalization(fused=False))
#    result.add(tf.keras.layers.BatchNormalization())

  result.add(tf.keras.layers.LeakyReLU())

  return result
```

```python
def upsample(filters, size, apply_dropout=False):
  initializer = tf.random_normal_initializer(0., 0.02)

  result = tf.keras.Sequential()
  result.add(
    tf.keras.layers.Conv2DTranspose(filters, size, strides=2,
                                    padding='same',
                                    kernel_initializer=initializer,
                                    use_bias=False))

#  result.add(tf.keras.layers.BatchNormalization(fused=False))
#  result.add(tf.keras.layers.BatchNormalization())

  if apply_dropout:
      result.add(tf.keras.layers.Dropout(0.5))

  result.add(tf.keras.layers.ReLU())

  return result
```

## 2.2. Google Drive に保存

デフォルトでは、カレントディレクトリ `/content/` 下の `training_checkpoints` ディレクトリにチェックポイントが作られていきます。途中で Google Colab が終わってしまうことも考えて、`fit` 実行前に、Google Drive へのシンボリックリンクにしておいた方が良いです。
また、次回の時用に前回チェックポイントを読むコードもいれます。

```bash
!mkdir -p /content/drive/My\ Drive/training_checkpoints/
!ln -s /content/drive/My\ Drive/training_checkpoints/
checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))
```

`checkpoint.restore` はチェックポイントがあれば読み込みますが、なければ何もしないのでそのまま呼んでおいて大丈夫そうです。


## 2.3. 学習完了 ～ エクスポート

ノートブックでは学習後、何枚か画像をテストして終わりですが、学習モデルを [TensorFlow.js](https://www.tensorflow.org/js/guide) 用にエクスポートするコードを追加します。

```python
generator.save('generator_model')
```

`.h5` をつけないで保存すれば、[TF2.x のデフォルトの SavedModel フォーマット](https://www.tensorflow.org/tutorials/keras/save_and_load)で保存されます。

```bash
!pip install tensorflowjs
!tensorflowjs_converter generator_model generator_tfjs_graph_model
!cp -r generator_tfjs_graph_model ./training_checkpoints/
```

その SavedModel を [tensorflowjs_converter](https://github.com/tensorflow/tfjs/tree/master/tfjs-converter) で `generator_tfjs_graph_model` に[変換](https://www.tensorflow.org/js/guide/conversion)しておきます。前述の Generator 作成時の `fused=False` が無いと、ここで以下のような WARN が出ていました。

```
WARNING:tensorflow:Didn't find expected Conv2D or DepthwiseConv2dNative input to 'StatefulPartitionedCall/functional_1/sequential_10/batch_normalization_9/FusedBatchNormV3'
WARNING:tensorflow:Didn't find expected Conv2D or DepthwiseConv2dNative input to 'StatefulPartitionedCall/functional_1/sequential_11/batch_normalization_10/FusedBatchNormV3'
WARNING:tensorflow:Didn't find expected Conv2D or DepthwiseConv2dNative input to 'StatefulPartitionedCall/functional_1/sequential_12/batch_normalization_11/FusedBatchNormV3'
WARNING:tensorflow:Didn't find expected Conv2D or DepthwiseConv2dNative input to 'StatefulPartitionedCall/functional_1/sequential_13/batch_normalization_12/FusedBatchNormV3'
WARNING:tensorflow:Didn't find expected Conv2D or DepthwiseConv2dNative input to 'StatefulPartitionedCall/functional_1/sequential_14/batch_normalization_13/FusedBatchNormV3'
WARNING:tensorflow:Didn't find expected Conv2D or DepthwiseConv2dNative input to 'StatefulPartitionedCall/functional_1/sequential_15/batch_normalization_14/FusedBatchNormV3'
WARNING:tensorflow:Didn't find expected Conv2D or DepthwiseConv2dNative input to 'StatefulPartitionedCall/functional_1/sequential_16/batch_normalization_15/FusedBatchNormV3'
```

# 3. Tensorflow.js + p5.js でモデルを読み込んで表示

[TensorFlow 公式からもリンクがある](https://www.tensorflow.org/js/tutorials) [ml5.js](https://ml5js.org/) ですが、[ml5.js](https://ml5js.org/) で準備された Pre-trained Model を使うことに特化しているようです。2020/10/07 時点で最新の [ml5.js](https://ml5js.org/) は 0.5.0 ですが、使っている [TensorFlow.js](https://www.tensorflow.org/js/guide) が 1.2.2 と古いので、自分で学習したモデルを使うには直接 [TensorFlow.js](https://www.tensorflow.org/js/guide) を使った方が良さそうでした。(`tf.browser.*` にブラウザ用のラッパもあるので)

実装の内容としては、モデルに合わせた input に加工しているところの調査が面倒だったぐらいで、他は特段難しい部分はありませんでした。

```javascript
    resultTensor = tf.tidy(() => {
        // tf.browser.fromPixels() returns a Tensor from an canvas element.
        let inputTensor = tf.browser.fromPixels(select('canvas').elt).toFloat();

        // normalize [0 ~ 255] -> [-1 ~ 1]
        inputTensor = inputTensor.div(tf.scalar(127.5)).sub(tf.scalar(1.0));

        // expand dimension (HWC -> NHWC)
        inputTensor = inputTensor.expandDims();

        // make a prediction through pix2pix
        let resultTensor = pix2pixModel.predict(inputTensor);

        // to browser [-1 ~ 1] -> [0 ~ 1]
        resultTensor = resultTensor.div(tf.scalar(2)).add(tf.scalar(0.5));

        // squeeze dimension (NHWC -> HWC)
        return resultTensor.squeeze();
    });

    tf.browser.toPixels(resultTensor).then(data => {
        select('#output').html('');
        const canvas = document.createElement('canvas');
        canvas.width = 256; canvas.height = 256;
        const context = canvas.getContext('2d');
        const imageData = context.createImageData(256, 256);
        data.forEach((value, index) => { imageData.data[index] = value });
        context.putImageData(imageData, 0, 0);
        createImg(canvas.toDataURL(), "transfer result").parent('output');
    });
```

TensorFlow.js + p5.js のコードは [tset-tset-tset/p5js-tfjs_for_tensorflow-pix2pix](https://github.com/tset-tset-tset/p5js-tfjs_for_tensorflow-pix2pix/) にあります。

ライブデモは[こちら](https://tset-tset-tset.github.io/pix2pix/) (200MB ほどのモデルを落としてくるので注意です)。