*Interp移植録 - 消しゴムマジック / AOT-GAN (NNInterp/torchscript)

# 0.Prologue
暇つぶしに、興味を引いた DNNアプリを *Interpに移植して遊んでいる。
本稿はその雑記&記録。

前回の ["DeepFillV2"](https://qiita.com/ShozF/items/218a49841c47b2f360d7)に引き続いて「消しゴムマジック」な DNNをもう一つ試してみようと思う。今回取り上げるDNNモデルは、"AOT-GAN"と言うモデルだ。"DeepFillV2"と基本的なアプローチは似ているのだが、モデルの具体的な作りはよりシンプルになっていそうだ。

ちなみに、この機会に「消しゴムマジック」の技術をあれこれと探求・比較・修得しようという訳ではない。今回の移植は、拙作*Interpの統合バージョンNNInterp用にデモを書き溜めておこうという邪な動機によっている。実は、torchsrciptに簡単にコンバート出来るモデルならば何でも良いのだ:stuck_out_tongue_winking_eye:

# 1.Original Work
AOT-GANが解決しようとする画像処理タスクは、DeepFillV2のそれとほぼ同じだ。特に、高解像度画像（例：512x512）において、もっともらしく精細なテクスチャが合成できる画像インペインティング手法にフォーカスしている。

彼らのアプローチでは、Aggregated COntextual-Transformation(AOT)ブロックと呼ぶ特徴抽出器を、シンプルに複数積み重ねてAOT-GANモデルを構成している。AOTブロックは、異なる大きさの受容野を束ねた構造を持つ特徴抽出器で、注目点回りのテクスチャの特徴および注目点から少し離れた領域との相関特徴を同時に効率よく学習できるように工夫したモノの様だ。

- "Aggregated Contextual Transformations for High-Resolution Image Inpainting"
https://arxiv.org/abs/2104.01431

- GitHub: AOT-GAN for High-Resolution Image Inpainting
https://github.com/researchmm/AOT-GAN-for-Inpainting

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/14158/e6df6d11-b280-d532-d593-1dca4f6c9c27.png)

# 2.NNInterp用のLivebookノート
では、拙作のNNInterpを利用して "AOT-GAN"を Livebookで動かし遊んでみよう。NNInterpは、まだhexに公開していないモジュールなのだが、一連の*Interpを統合したバージョンに当たり、バックエンドのDNNエンジンを "Tensorflow lite", "Onnx Runtime", "torch script(Libtorch)"の中から選べるようにしている。本家の"AOT-GAN"は Pytorchで実装されているので、このモデルを torch scriptにコンバートすれば、NNInterpでサクッと動くはずとの目論見だ。

Mix.installの依存リストの記述は次の通り。NNInterpモジュールは、小生のGitHubから引っ張ってくる。DNNエンジンに"torch script(Libtorch)"を選択したいので、環境変数 `NNINTERP`に "LibTorch"(大文字/小文字は don't care)を設定する。

```elixir
File.cd!(__DIR__)
# for windows JP
System.shell("chcp 65001")

Mix.install(
  [
    {:nn_interp, github: "shoz-f/nn-interp"},
    {:cimg, "~> 0.1.20"},
    {:kino, "~> 0.7.0"}
  ],
  system_env: [{"NNINTERP", "LibTorch"}]
)
```

AOT-GANのモデルは、書き換え領域が白塗りされた画像とマスク画像を受け取り、白塗り領域がそれらしく合成された画像を返す。

> [モデル・カード]
> - inputs:
> [0] f32:{1,3,512,512} - マスク領域が白塗りされた元画像。各RGB値を{-1.0～1.0}に正規化。BGRカラー。
> [1] f32:{1,1,512,512} - マスク画像。画素値を{0.0～1.0}に正規化。マスク部の値は "1.0"。
> - outputs:
> [0] f32:{1,3,512,512} - 白塗り領域が描き直された画像、各RGB値は{-1.0～1.0}に正規化されている。BGRカラー

```elixir
defmodule AotGan do
  @width 512
  @height 512

  use NNInterp,
    model: "./model/AOT-GAN_CELLEBA-HQ.pt",
    url: "https://github.com/shoz-f/nn-interp/releases/download/0.0.1/AOT-GAN_CELLEBA-HQ.pt",
    inputs: [f32: {1, 3, @height, @width}, f32: {1, 1, @height, @width}],
    outputs: [f32: {1, 3, @height, @width}]

  def apply(img, mask) do
    # preprocess
    input0 =
      CImg.builder(img)
      |> CImg.resize({@width, @height})
      |> CImg.to_binary([{:range, {-1.0, 1.0}}, :nchw, :bgr])

    input1 =
      CImg.builder(mask)
      |> CImg.resize({@width, @height})
      |> CImg.to_binary([{:range, {0.0, 1.0}}, :nchw])

    # prediction
    session()
    |> NNInterp.set_input_tensor(0, input0)
    |> NNInterp.set_input_tensor(1, input1)
    |> NNInterp.invoke()
    |> NNInterp.get_output_tensor(0)
    |> CImg.from_binary(@width, @height, 1, 3, [{:range, {-1.0, 1.0}}, :nchw, :bgr])
    |> CImg.resize(img)
  end
end
```

# 3.デモンストレーション
`AotGan`を起動する。

```elixir
AotGan.start_link([])
```

AotGanで遊ぶには、先に述べたように書き換え領域を白塗りにした入力画像`masked_img`とそのマスク画像`mask`が要る。お好みのペイント・アプリで予め作成しておこう。ここでは、本家のサンプル画像 - 眼鏡を掛けたおじさんの画像 - を借用して、おじさんの眼鏡を消すようにマスク画像を作ってみた。

入力画像`masked_img.jpg`とマスク画像`mask.jpg`を与え、`AotGan`を実行する。

```elixir
origin = CImg.load("origin.jpg")
input = CImg.load("masked_img.jpg")
mask = CImg.load("mask.jpg")

result = AotGan.apply(input, mask)

Enum.map([input, mask, origin, result], &CImg.display_kino(&1, :jpeg))
|> Kino.Layout.grid(columns: 2)
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/14158/6cc31cfa-e6dc-7b37-95de-14ffa77aba30.png)

なんだか目の回りに少し隈が出ているお疲れ顔になったが、眼鏡はちゃんと消えているね:wink:

# 4.Epilogue

Elixirで「消しゴムマジック」な DNNモデル "AOT-GAN"を動かしてみた。

う～む、「これを使えるアプリにするには、マスクを編集するお絵描きソフトが必要だなぁ…でも作るの面倒くさいなぁ」と言う心境にある。NNInterpに添付するデモとしては不適かも:sweat_smile:

一方でNNInterpは、英文資料を用意するのが億劫で、なかなか hex公開に辿り着けない。あ～～あっ:weary:

(気晴らしに Libtensorflow/saved-modelの組み込みでも検討してみようか?)

# Appendix
- NNInterpのノート
https://github.com/shoz-f/nn-interp/blob/main/demo_aotgan/AOT-GAN.livemd

- AOT-GANの torch scriptモデルの調達手順
https://github.com/shoz-f/nn-interp/blob/main/demo_aotgan/Export_AOT_GAN_torchscript.ipynb