*Interp移植録 - 消しゴムマジック / DeepFill v2 (TflInterp)

0.Prologue

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

前回は、Tensorflow liteの「カスタム演算子組み込み」の手習いとして、旧MediaPapeの Hair-Segmentationを TflInterpに移植して動かしてみた。事前に用意されている３つのカスタム演算子を TflInterpのコードに組み込むだけの作業ではあったが、おかげでカスタム演算子については何をどう調べれば良いのかそこそこの土地勘がついたかなぁと思う。その勢いに乗じて、見よう見まねではあるがカスタム演算ExtractImagePatchesを自作し TflInterpに組み込むところまで辿り着いた。ろくすっぽテストをしていないコードではあるが……たぶん動くだろう

ということで、本丸の DeepFillV2 - 俗に言う「消しゴムマジック」 - の攻略に取り掛かろう。

1.Original Work

DeepFillV2が解決しようとする画像処理タスクは、「大きな欠陥領域(白抜き)を持つ画像を、もっともらしい合成画像で欠陥領域を埋めてあたかも欠陥がない画像にする」ことである。某社のスマホCMで流れている「消しゴムマジック」(MagicEraser)は、"監督"が写っている領域を欠陥と見做し、その領域をあたかもそれらしい背景画像(合成)に書き換えているのだと想像する。

このタスクに対するアプローチは、レトロな画像処理でも行われていたテクスチャ解析&パッチワークや、ディープラーニングによるGANなどがある。前者は画像中の他の領域からパッチ画像を借りてくる手法のため、欠陥部分が顔などの特異なパターンの場合には原理的に合成できない。一方、後者はニューラルネットにエンコードされた画像特徴記憶から合成画像を生成するので、特異なバターンに対してもそれらしい画像が得られるのだが、コンボリューション演算子の局所性のために欠陥領域の境界部分の合成が不連続になる傾向がある。

DeepFillV2は GANをベースとして、欠陥領域の境界部分が不連続になる課題に対して "Contenxtual Attention"や "Gated Convolution"を導入し、コンボリューション演算子の局所性を緩和するアプローチを取っているようだ。TflInterpへの移植作業でネックになっていたExtractImagePatchesは、この "Contenxtual Attention"が使用している。

• "Generative Image Inpainting with Contextual Attention"
  https://arxiv.org/abs/1801.07892

• "Free-Form Image Inpainting with Gated Convolution"
  https://arxiv.org/abs/1806.03589

• GitHub: Generative Image Inpainting
  https://github.com/JiahuiYu/generative_inpainting

2.TflInterp用のLivebookノート

先日、hexにアップした "tfl_interp 0.1.11"にはExtractImagePatchesカスタム演算子を組み込み済みなので、このバージョンを用いて DeepFillV2を遊んでみる。

前回の記事でも触れたが、tfl_interpのダウンロード&ビルドには数十分の時間を要するので、適当なフォルダでtfl_interpを単体ビルドし、アプリケーションからはその成果物を参照する方が利便性が良い。

git clone https://github.com/shoz-f/tfl_interp
cd tfl_interp
mix deps.get
mix compile

DeepFillV2のLivebookノートは下記の位置に置こう。

── tfl_interp - tfl_interpモジュール
    ├─ lib
    ├─ priv
    │   └─ tfl_interp - 実行形式
    └─ demo_deepfillv2
        └─ DeepFillV2.livemd - デモ用Livebookノート

Mix.installの依存リストに記述するモジュールは下記の通り。

File.cd!(__DIR__)
# for windows JP
System.shell("chcp 65001")
System.put_env("SKIP_MAKE_TFLINTERP", "YES")

Mix.install([
  {:tfl_interp, path: ".."},
  {:cimg, "~> 0.1.20"},
  {:kino, "~> 0.7.0"}
])

DeepFillV2モデルの入力には、元画像とマスク画像を連結した画像を渡す。本家の例では、元画像に白抜きの欠陥領域有りの画像を渡しているが、白抜きが無い画像を渡しても期待通りに動作するようだ。出力には、元画像のマスクの領域が描き直された画像が返ってくる。

[モデル・カード]

• inputs:
  [0] f32:{1,512,1360,3} - 元画像とマスク画像を連結し、画素値を{0.0～255.0}に正規化
• outputs:
  [0] u8:{1,512,680,3} - マスクの領域が描き直された画像

defmodule DeepFillV2 do
  @width  680
  @height 512

  alias TflInterp, as: NNInterp
  use NNInterp,
    model: "./model/deepfillv2.tflite",
    url: "https://github.com/shoz-f/tfl_interp/releases/download/0.0.1/deepfillv2.tflite",
    inputs: [f32: {1,@height,2*@width,3}],
    outputs: [u8: {1,@height,@width,3}]

  def apply(img, mask) do
    # preprocess
    input0 = CImg.builder(img)
      |> CImg.append(mask, :x)
      |> CImg.resize({2*@width, @height})
      |> CImg.to_binary(range: {0.0, 255.0})

    # prediction
    output = session()
      |> NNInterp.set_input_tensor(0, input0)
      |> NNInterp.invoke()
      |> NNInterp.get_output_tensor(0)
      |> CImg.from_binary(@width, @height, 1, 3, [{:dtype, "<u1"}, :bgr])

    # postprocess
    CImg.resize(output, img)
  end
end

なんだかなぁと思うぐらい単純なコードで済んでしまった

3.デモンストレーション

DeepFillV2で遊ぶには、画像中の書き換えたい領域を指定するマスク画像が要るので、お好みのペイント・アプリで予め作成しておく。マスク画像は、背景が黒、指定領域は白となるように作る。ペイント・アプリのレイヤー機能を用いればパパッとできるだろう。

DeepFillV2を起動する。

DeepFillV2.start_link([])

元画像とマスク画像を与え、DeepFillV2を実行する。

origin = CImg.load("sample_raw.jpg")
mask   = CImg.load("sample_mask.jpg")

result = DeepFillV2.apply(origin, mask)

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

某スマホの「消しゴムマジック」ほどの精度はないが、まあまあのレベルかな。

4.Epilogue

やっと片付いた。ふうぅ、長かった。
「消しゴムマジック」の様な画像処理をやってみようと思い立ちはや数か月。それらしきDNNモデル DeepFillV2は見つけたものの、その実装が TF1.0と古く、→ TF2.0 → tfliteとコンバートするのに試行錯誤。やっとの思いで用意した tfliteモデルだが、拙作TflInterpで利用するにはカスタム演算子 "ExtractImagePatches"が必要と分かり、どんよりとした挫折感に包まれた。

このテーマは捨てるかと思いつつ、次のネタ仕込みで MadiaPipeをごそごそと触っていると、"Hair-Segmentation"タスクがカスタム演算子を組み込んでおり、TflInterpへの移植も簡単そうだと分かった。これを足掛りにすればもしやと、ひとかけらの希望を胸に Tensorflow liteのソースコードの迷宮を彷徨い、そして起死回生の今日を迎えた。終わり良ければ、すべて良し。大団円である。

さて、手元には tfliteモデルをコンバートした onnxモデルがある。OnnxInterpで動かすには、きっとカスタム演算子が必要だろうな……見なかったことにしよう

Appendix

• TflInterpのノート
  https://github.com/shoz-f/tfl_interp/blob/main/demo_deepfillv2/DeepFillV2.livemd

• DeepFill v2の tfliteモデルの調達手順
  https://github.com/shoz-f/tfl_interp/blob/main/demo_deepfillv2/Export_DeepFillv2_tflite.ipynb
  

• *Interp移植外伝 - Tensorflow liteカスタム演算子の作成例
  https://qiita.com/ShozF/items/19ce312e3438d529e5d6