Elixir Livebook で画像分割・結合

## はじめに

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

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

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

https://youtu.be/RkMzCQm-Ws4

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

https://gist.github.com/zacky1972/7777c34c89ae797c581a946c08cba6c4

https://gist.github.com/zacky1972/d3f590b7ce9ba9de1f8653b4fdffbf95

https://gist.github.com/zacky1972/ab48d2cea7a9aecd8d8cefa00c2b2ba5

https://gist.github.com/zacky1972/c029e2c1999d36f67ac2d408af82bd86

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

https://github.com/RyoWakabayashi/elixir-learning/blob/main/livebooks/distributed_image_processing/image_division.livemd

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

## 実行環境

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

https://github.com/RyoWakabayashi/elixir-learning

## 準備

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

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

セットアップ対象

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

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

```elixir
# 再実行時、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)
```

画像を読み込みます

```elixir
mat = Evision.imread(lenna)
```

![スクリーンショット 2022-11-11 13.16.30.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/cc889613-2b7c-73d2-d5df-ed9bfdf82abd.png)

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

```elixir
# 画像サイズ
shape = mat.shape

# 分割サイズ
div_size = 64

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

## 縦方向に分割

縦方向に分割します

```elixir
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))` で各マトリックスを出力結果に描画します

![スクリーンショット 2022-11-11 13.26.05.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/bb3f997d-6f55-db51-519d-701013f8903e.png)

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

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

```elixir
# 分割したファイルを保存
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)
```

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

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

```elixir
# 分割したファイルを取得
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()
```

![スクリーンショット 2022-11-11 13.31.31.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/70ee4d95-fb20-2eac-0eed-45697354f5b3.png)

## 横方向に分割

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

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

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

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

![スクリーンショット 2022-11-11 13.38.34.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/ac83e29d-9c3d-1e04-f945-bc940ffef48c.png)

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

```elixir
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()
```

![スクリーンショット 2022-11-11 13.41.37.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/7b569448-f743-7803-bffb-6bf66c7a4ede.png)

こちらも保存します

```elixir
# 分割したファイルを保存
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` を指定することで横方向に結合が可能です

```elixir
# 分割したファイルを取得
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()
```

![スクリーンショット 2022-11-11 13.44.56.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/77086c1c-2b5d-3c94-648e-f246f6ca5ec6.png)

## タイル状に分割

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

```elixir
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 個のタイルに分割されました

![スクリーンショット 2022-11-11 13.48.19.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/a68fb7ee-bc39-9d48-f24f-abfd09492dc7.png)

これも個別に保存します

```elixir
# 分割したファイルを保存
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)
```

タイルを結合します

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

```elixir
# 横方向の分割数を取得
{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)` で、反転するか判定しています

![スクリーンショット 2022-11-11 13.59.47.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/2e738848-1a0d-ffae-644e-9bfc1d226c88.png)

## まとめ

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