Elixir Livebook で画像分割・結合

はじめに

Elixir で画像を分割・結合してみます

この記事は @zacky1972 さんが ElixirConf US 2022 で発表した内容の一部を Livebook 上で実行したものです

ElixirConf US 2022 の @zacky1972 さんの発表動画はこちら

参考にした @zacky1972 さんの Gist はこちら

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

次回は Flow を使って並列処理する予定です

実行環境

以下のリポジトリーのコンテナ上で実行しています

準備

ノートブックを起動して、以下のコードを実行してセットアップします

Mix.install([
  {:download, "~> 0.0.4"},
  {:evision, "~> 0.1"},
  {:kino, "~> 0.7"},
  {:nx, "~> 0.4"}
])

セットアップ対象

• download: データダウンロード
• evision: 画像処理
• kino: 出力可視化
• nx: 行列演算

処理する画像をダウンロードしてきます

# 再実行時、Download.from()でeexistエラーになるのを防止
File.rm("Lenna_%28test_image%29.png")

lenna =
  Download.from("https://upload.wikimedia.org/wikipedia/en/7/7d/Lenna_%28test_image%29.png")
  |> elem(1)

画像を読み込みます

mat = Evision.imread(lenna)

後で使用する値を取得しておきます

# 画像サイズ
shape = mat.shape

# 分割サイズ
div_size = 64

# ファイル名
dst_file_ext = Path.extname(lenna)
dst_file_basename = Path.basename(lenna, dst_file_ext)

縦方向に分割

縦方向に分割します

div_img =
  mat
  |> Evision.Mat.to_nx()
  # 分割
  |> Nx.to_batched(div_size)
  |> Enum.map(&Evision.Mat.from_nx_2d(&1))
  |> Enum.map(&Kino.render(&1))
  |> dbg()

Evision.Mat.to_nx() で Nx のテンソルに変換します

この時点で縦512、横512、色3のテンソルになっています

Nx.to_batched(div_size) でテンソルを div_size で縦方向に分割します

これにより、縦64、横512、色3のテンソル8個になっています

Enum.map(&Evision.Mat.from_nx_2d(&1)) で Nx のテンソルから Evision のマトリックスに戻します

Enum.map(&Kino.render(&1)) で各マトリックスを出力結果に描画します

縦方向に８分割されました

以下のコードを実行し、分割した画像をファイルに保存します

# 分割したファイルを保存
dst_files =
  Stream.unfold(0, fn counter -> {counter, counter + 1} end)
  |> Stream.map(&"#{dst_file_basename}_h_#{&1}#{dst_file_ext}")

dst_file_paths =
  div_img
  |> Enum.zip(dst_files)
  |> Enum.map(fn {img, dst_file} ->
    Evision.imwrite(dst_file, img)
    dst_file
  end)

　分割した画像の各ファイルを読み込んでを結合します

単に結合しても面白くないので、偶数番目の色を反転させています

# 分割したファイルを取得
imgs =
  Stream.unfold(0, fn counter -> {counter, counter + 1} end)
  |> Stream.map(&{&1, "#{dst_file_basename}_h_#{&1}#{dst_file_ext}"})
  |> Stream.take_while(fn {_, f} -> File.exists?(f) end)
  |> Enum.map(fn {index, file_name} ->
    new_tensor =
      file_name
      |> Evision.imread()
      # Evisionバックエンドでは concatenate できないため、バイナリバックエンドを指定
      |> Evision.Mat.to_nx(Nx.BinaryBackend)

    # 偶数の場合は色を反転
    case rem(index, 2) do
      0 ->
        Nx.reverse(new_tensor, axes: [2])
      _ ->
        new_tensor
    end
  end)
  # 結合
  |> Nx.concatenate()
  # トリミング
  |> Nx.slice([0, 0, 0], Tuple.to_list(shape))
  |> then(&Evision.Mat.from_nx_2d(&1))
  |> dbg()

横方向に分割

次は横方向に分割してみます

Nx.to_batched は縦方向にしか分割してくれないため、
まず縦横を入れ替える必要があります

縦横の入れ替えは以下のようにします

mat
|> Evision.Mat.to_nx()
|> Nx.transpose(axes: [1, 0, 2])
|> Evision.Mat.from_nx_2d()

この状態で Nx.to_batched を実行し、分割結果の縦横をもとに戻せばOKです

div_img =
  mat
  |> Evision.Mat.to_nx()
  #　縦横入れ替え
  |> Nx.transpose(axes: [1, 0, 2])
  # 分割
  |> Nx.to_batched(div_size)
  #　縦横入れ替え
  |> Enum.map(&Nx.transpose(&1, axes: [1, 0, 2]))
  |> Enum.map(&Evision.Mat.from_nx_2d(&1))
  |> Enum.map(&Kino.render(&1))
  |> dbg()

こちらも保存します

# 分割したファイルを保存
dst_file_ext = Path.extname(lenna)
dst_file_basename = Path.basename(lenna, dst_file_ext)

dst_files =
  Stream.unfold(0, fn counter -> {counter, counter + 1} end)
  |> Stream.map(&"#{dst_file_basename}_v_#{&1}#{dst_file_ext}")

dst_file_paths =
  div_img
  |> Enum.zip(dst_files)
  |> Enum.map(fn {img, dst_file} ->
    Evision.imwrite(dst_file, img)
    dst_file
  end)

また結合します

Nx.concatenate に axis: 1 を指定することで横方向に結合が可能です

# 分割したファイルを取得
imgs =
  Stream.unfold(0, fn counter -> {counter, counter + 1} end)
  |> Stream.map(&{&1, "#{dst_file_basename}_v_#{&1}#{dst_file_ext}"})
  |> Stream.take_while(fn {_, f} -> File.exists?(f) end)
  |> Enum.map(fn {index, file_name} ->
    new_tensor =
      file_name
      |> Evision.imread()
      |> Evision.Mat.to_nx(Nx.BinaryBackend)

    # 偶数の場合は色を反転
    case rem(index, 2) do
      0 ->
        Nx.reverse(new_tensor, axes: [2])
      _ ->
        new_tensor
    end
  end)
  # 結合
  |> Nx.concatenate(axis: 1)
  # トリミング
  |> Nx.slice([0, 0, 0], Tuple.to_list(shape))
  |> then(&Evision.Mat.from_nx_2d(&1))
  |> dbg()

タイル状に分割

最後に縦横両方分割してタイル状にします

div_img =
  mat
  |> Evision.Mat.to_nx()
  # 縦方向に分割
  |> Nx.to_batched(div_size)
  |> Enum.map(&Nx.transpose(&1, axes: [1, 0, 2]))
  # 横方向に分割
  |> Enum.flat_map(&Nx.to_batched(&1, div_size))
  |> Enum.map(&Nx.transpose(&1, axes: [1, 0, 2]))
  |> Enum.map(&Evision.Mat.from_nx_2d(&1))
  |> Enum.map(&Kino.render(&1))

まず縦方向に分割した後、それらを横方向に分割します

8 * 8 で 64 個のタイルに分割されました

これも個別に保存します

# 分割したファイルを保存
dst_file_ext = Path.extname(lenna)
dst_file_basename = Path.basename(lenna, dst_file_ext)

dst_files =
  Stream.unfold(0, fn counter -> {counter, counter + 1} end)
  |> Stream.map(&"#{dst_file_basename}_t_#{&1}#{dst_file_ext}")

dst_file_paths =
  div_img
  |> Enum.zip(dst_files)
  |> Enum.map(fn {img, dst_file} ->
    Evision.imwrite(dst_file, img)
    dst_file
  end)

タイルを結合します

今度はチェック模様に反転してみましょう

# 横方向の分割数を取得
{width, _, _} = shape
h_size = div(width, div_size)

# 分割したファイルを取得
imgs =
  Stream.unfold(0, fn counter -> {counter, counter + 1} end)
  |> Stream.map(&{&1, "#{dst_file_basename}_t_#{&1}#{dst_file_ext}"})
  |> Stream.take_while(fn {_, f} -> File.exists?(f) end)
  |> Enum.map(fn {t_index, file_name} ->
    new_tensor =
      file_name
      |> Evision.imread()
      |> Evision.Mat.to_nx(Nx.BinaryBackend)

    # 何番目のタイルなのか分かるようにタイル番号を保持
    {new_tensor, t_index}
  end)
  |> Enum.chunk_every(h_size)
  |> Enum.map(fn new_tensor_list ->
    new_tensor_list
    |> Enum.with_index()
    |> Enum.map(fn {{new_tensor, t_index}, v_index} ->
      cond do
        rem(v_index, 2) == rem(div(t_index, h_size), 2) ->
          Nx.reverse(new_tensor, axes: [2])
        true ->
          new_tensor
      end
    end)
    # 横方向に結合
    |> Nx.concatenate(axis: 1)
  end)
  # 縦方向に結合
  |> Nx.concatenate()
  # トリミング
  |> Nx.slice([0, 0, 0], Tuple.to_list(shape))
  |> then(&Evision.Mat.from_nx_2d(&1))
  |> dbg()

64個の画像を 8 * 8 に結合するため、まず横方向が8個であることを計算します

タイルを読み込むとき {new_tensor, t_index} として何番目のタイルなのかを保持しておきます

Enum.chunk_every(h_size) で 64 個の配列を 8 * 8 の二次元配列に変換します

先に横方向に8個ずつ結合します

チェック模様にするため rem(v_index, 2) == rem(div(t_index, h_size), 2) で、反転するか判定しています

まとめ

Nx.to_batched で axis が指定できればもっと分かりやすいのに