0.Prologue

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

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

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

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

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)を設定する。

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カラー

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を起動する。

AotGan.start_link([])

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

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

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)