Nx でカラー画像の黒背景を透過する

はじめに

Discord の「elixir と見習い錬金術師」サーバーにて、 @SF-28 さんから質問がりました

カラー画像の黒背景を透過する処理についてです。
とりあえず以下コードで処理自体は実現出来たんですが、あまりスマートじゃないなと・・・
EvisionやNxでこのあたり上手く扱う方法をご存知の方いらっしゃいませんか？

bgr_img
|> Evision.Mat.to_nx()
|> Nx.to_list()
|> Enum.map(fn x ->
  Enum.map(x, fn
    [0, 0, 0] -> [0, 0, 0, 0]
    y -> y ++ [255]
  end)
end)
|> Nx.tensor(type: :u8)
|> Evision.Mat.from_nx_2d()

というわけで記事に答えを書きます

実装したノートブックはこちら

実行環境

Livebook を使います

• Elixir: 1.14.2
• Livebook: 0.8.1

セットアップ

必要なモジュールをインストールします

Mix.install([
  {:evision, "~> 0.1"},
  {:kino, "~> 0.8"},
  {:nx, "~> 0.5"}
])

• evision: 画像処理
• kino: Livebook の UI/UX
• nx: 行列演算

画像の準備

今回は以下の画像を使います

piyopiyo.ex のキャラクターを殻から出しています

以下のコードを実行するとアップロード用のウィジェットが表示されます

img_input = Kino.Input.image("IMAGE", format: :png)

ドラッグ&ドロップ等で画像を選択し、アップロードします

画像を Nx で扱えるテンソルに変換します

Evision.imdecode でバイナリデータを evision (OpenCV) の行列データに変換します

Evision.Mat.to_nx で evision の行列データを Nx のテンソルに変換します

evision で読み込んだ画像は RGB ではなく BGR の順に色データが並んでいるので bgr_img にしています

Kino.Image.new に画像のテンソルを渡すと、Livebook 上に画像として表示してくれます

bgr_img =
  img_input
  |> Kino.Input.read()
  |> then(& &1.data)
  |> Evision.imdecode(Evision.Constant.cv_IMREAD_COLOR)
  |> Evision.Mat.to_nx(Nx.BinaryBackend)

Kino.Image.new(bgr_img)

背景透過

BGR に A （アルファ）チャネルを付けると透過画像になります

A が 0 のとき完全な透明、 A が 255 のときは完全な不透明です

alpha =
  bgr_img
  |> Nx.reduce_max(axes: [2])
  |> Nx.greater(0)
  |> Nx.multiply(255)
  |> Nx.new_axis(2)

bgra_img =
  [bgr_img, alpha]
  |> Nx.concatenate(axis: 2)

Kino.Image.new(bgra_img)

Nx.reduce_max(axes: [2]) により、画像内の各ピクセルについて BGR の最大を計算します

黒は [0, 0, 0] なので最大が 0 で、それ以外の色は 0 より大きくなります

Nx.greater(0) で、最大が 0 より大きければ 1　、 そうでなければ 0 にします

これによって黒のピクセルは 0 、 それ以外のピクセルは 1 になります

Nx.multiply(255) で各ピクセルの値を 255 倍にします

黒は 0 （透明）、それ以外は 255 （不透明）になります

Nx.new_axis(2) でテンソルの形を {<横ピクセル数>, <縦ピクセル数>} から {<横ピクセル数>, <縦ピクセル数>, 1} に変形します

後で BGR とくっつけるためです

Nx.concatenate(axis: 2) で BGR に A をくっつけます

まとめ

Enum.map などで処理するよりも、 Nx の行列演算の方が圧倒的に速いので、可能な限り行列演算するようにしましょう