Elixir Desktop で YOLOv3 (evision)を動かしてみた

## はじめに

最近は Elixir ばかりですが（他の業務もこなしていますよ）、
とうとう Elixir Desktop で YOLO v3 を動かしてみます

全コードはこちら

https://github.com/RyoWakabayashi/yolo-on-elixir-desktop

参考にした @the_haigo さんのリポジトリーはこちら

https://github.com/thehaigo/axon_desktop

## 実行環境

- macOS Monterey 12.4
- Elixir 1.13.4
- Erlang OTP24

## フォーク

まずは Elixir Desktop 公式のサンプルをフォークしてきます

https://github.com/elixir-desktop/desktop-example-app

フォークしたものをローカルにクローンしてきて

`todo` を `yolo` 、 `Todo` を `Yolo` に一括置換します

## 依存パッケージの変更

@the_haigo さんの記事を参考にして、 mix.exs を以下のように編集します

特筆していない箇所は @the_haigo さんのままなので、そちらを参照してください

https://qiita.com/the_haigo/items/22ca888ab2b19b558828

```elixir
  ...
  defp deps do
    [
      {:desktop, github: "elixir-desktop/desktop", tag: "v1.4.0"},

      # Phoenix
      {:phoenix, "~> 1.6"},
      {:phoenix_live_view, "~> 0.17.4"},
      {:phoenix_html, "~> 3.0"},
      {:phoenix_live_reload, "~> 1.3", only: [:dev]},
      {:gettext, "~> 0.18"},
      {:plug_cowboy, "~> 2.5"},
      {:jason, "~> 1.2"},

      # Assets
      {:esbuild, "~> 0.2", runtime: Mix.env() == :dev},
      {:tailwind, "~> 0.1", runtime: Mix.env() == :dev},
      {:petal_components, "~> 0.17"},

      # Credo
      {:credo, "~> 1.5", only: [:dev, :test], runtime: false},

      # YOLO
      {:evision, "~> 0.1.0-dev", github: "cocoa-xu/evision", branch: "main"},
      {:exla, "~> 0.3.0-dev", github: "elixir-nx/nx", sparse: "exla"},
      {:nx, "~> 0.3.0-dev", [env: :prod, git: "https://github.com/elixir-nx/nx.git", sparse: "nx", override: true]},
      {:stb_image, "~> 0.4.0"}
    ]
  end
end
```

PetalCompoent を使うため、 dart_sass を消して tailwind と petal_components を追加しています

また、機械学習用に evision exla nx stb_image を追加します

ecto_sqlite3 など、使わないものは消しています

その他、全体的に DB (Repo) を使っている箇所を削除します

詳細はこちら

https://github.com/RyoWakabayashi/yolo-on-elixir-desktop/commit/cefd5bd34e7500477e0b3487c22bc7d8892ded07

## モデルのダウンロード

Darknet の YOLOv3 モデルをダウンロードして priv 配下に置きます

```bash
#!/bin/bash

mkdir -p priv/models

wget -c \
  -N https://pjreddie.com/media/files/yolov3.weights \
  -O ./priv/models/yolov3.weights

wget -c \
  -N https://raw.githubusercontent.com/pjreddie/darknet/master/cfg/yolov3.cfg \
  -O ./priv/models/yolov3.cfg

wget -c \
  -N https://raw.githubusercontent.com/pjreddie/darknet/master/data/coco.names \
  -O ./priv/models/labels.txt
```

これらのファイルは .gitignore に追加して git では無視するようにしておきます

.gitignore
```
...
# models
priv/models
```

## モデルのロード

検出実行の度にモデルをロードしたくないので、
モデルを保持しておくための Store を実装します

lib/yolo_app/store.ex
```elixir
defmodule YoloApp.Store do
  use Agent
  alias __MODULE__
  alias Evision, as: OpenCV

  def start_link(opts) do
    # priv ディレクトリー配下から取得する
    priv_path = List.to_string(:code.priv_dir(:yolo_app))

    cfg_path = priv_path <> "/models/yolov3.cfg"
    weights_path = priv_path <> "/models/yolov3.weights"
    labels_path = priv_path <> "/models/labels.txt"

    net = OpenCV.DNN.readNet!(weights_path, config: cfg_path, framework: "")
    out_names = OpenCV.DNN.Net.getUnconnectedOutLayersNames!(net)
    label_list =
      labels_path
      |> File.stream!()
      |> Enum.map(&String.trim/1)

    # Agent に入れておく
    Agent.start_link(fn ->
      %{
        net: net,
        out_names: out_names,
        label_list: label_list,
      }
    end, name: __MODULE__)
  end

  # 使用時に Agent から取り出す
  def get(key) do
    Agent.get(__MODULE__, &Map.get(&1, key))
  end
end
```

DBの初期化を行なっていた箇所をモデルロードに置き換えます

- 変更前

lib/todo_app.ex
```elixir
...
Application.put_env(:todo_app, TodoApp.Repo,
  database: Path.join(config_dir(), "/database.sq3")
)

{:ok, sup} = Supervisor.start_link([TodoApp.Repo], name: __MODULE__, strategy: :one_for_one)
TodoApp.Repo.initialize()
...
```

- 変更後

lib/yolo_app.ex
```elixir
...
    {:ok, sup} = Supervisor.start_link([YoloApp.Store], name: __MODULE__, strategy: :one_for_one)
...
```

## ワーカーの実装

物体検出を実行するワーカーを実装します

lib/yolo_app/worker.ex
```elixir
defmodule YoloApp.Worker do
  use Agent
  alias YoloApp.Store
  alias YoloApp.Worker
  alias Evision, as: OpenCV

  def detect(binary) do
    label_list = Store.get(:label_list)

    mat = to_mat(binary)

    predictions =
      mat
      |> preprocess()
      |> predict()
      |> to_tensor()
      |> filter_predictions(0.8)
      |> format_predictions()
      |> nms(0.8, 0.7)
      |> Enum.map(&Map.put(&1, :class, Enum.at(label_list, &1.class)))

    drawed =
      OpenCV.imencode!(".png", draw_predictions(mat, predictions))
      |> IO.iodata_to_binary()

    {predictions, drawed}
  end

  def measure(function) do
    {time, result} = :timer.tc(function)
    IO.puts "Time: #{time}ms"
    result
  end

  def to_mat(binary) do
    binary
    |> StbImage.from_binary()
    |> elem(1)
    |> StbImage.to_nx()
    |> OpenCV.Nx.to_mat!()
  end

  def preprocess(mat) do
    OpenCV.DNN.blobFromImage!(mat, size: [608, 608], swapRB: true, crop: false)
  end

  def predict(blob) do
    net = Store.get(:net)
    out_names = Store.get(:out_names)

    net
    |> OpenCV.DNN.Net.setInput!(
      blob,
      name: "",
      scalefactor: 1 / 255,
      mean: [0, 0, 0]
    )
    |> OpenCV.DNN.Net.forward!(outBlobNames: out_names)
  end

  def to_tensor(predictions) do
    predictions
    |> Enum.map(fn prediction ->
      OpenCV.Nx.to_nx(prediction)
    end)
    |> Nx.concatenate()
  end

  def filter_predictions(predictions, score_threshold) do
    size =
      predictions
      |> Nx.shape()
      |> elem(0)

    threshold_tensor =
      score_threshold
      |> Nx.tensor()
      |> Nx.broadcast({size})

    index_list =
      Nx.transpose(predictions)[4]
      |> Nx.greater(threshold_tensor)
      |> Nx.to_flat_list()
      |> Enum.with_index()
      |> Enum.filter(fn {value, _} -> value == 1 end)
      |> Enum.map(&elem(&1, 1))
      |> Nx.tensor()

    Nx.take(predictions, index_list)
  end

  def format_predictions(predictions) do
    predictions
    |> Nx.to_batched_list(1)
    |> Enum.map(fn t ->
      class_score_list = t[0][5..-1//1]
      class_id = class_score_list |> Nx.argmax() |> Nx.to_number()
      class_score = class_score_list[class_id] |> Nx.to_number()
      score = t[0][4] |> Nx.to_number() |> Kernel.*(class_score)

      center_x = t[0][0] |> Nx.to_number()
      center_y = t[0][1] |> Nx.to_number()
      box_width = t[0][2] |> Nx.to_number()
      box_height = t[0][3] |> Nx.to_number()
      min_x = center_x - box_width / 2
      min_y = center_y - box_height / 2
      max_x = center_x + box_width / 2
      max_y = center_y + box_height / 2

      box = [min_x, min_y, max_x, max_y]

      %{
        box: box,
        score: score,
        class: class_id
      }
    end)
  end

  def nms(formed_predictions, score_threshold, nms_threshold) do
    box_list = Enum.map(formed_predictions, & &1.box)
    score_list = Enum.map(formed_predictions, & &1.score)

    box_list
    |> OpenCV.DNN.nmsBoxes!(score_list, score_threshold, nms_threshold)
    |> Enum.map(&Enum.at(formed_predictions, &1))
  end

  def draw_predictions(mat, predictions) do
    {height, width, _} = OpenCV.Mat.shape!(mat)

    predictions
    |> Enum.reduce(mat, fn prediction, drawed_mat ->
      box = prediction.box
      left = Enum.at(box, 0) |> Kernel.*(width) |> trunc()
      top = Enum.at(box, 1) |> Kernel.*(height) |> trunc()
      right = Enum.at(box, 2) |> Kernel.*(width) |> trunc()
      bottom = Enum.at(box, 3) |> Kernel.*(height) |> trunc()

      drawed_mat
      |> OpenCV.rectangle!(
        [left, top],
        [right, bottom],
        [255, 0, 0],
        thickness: 4
      ) # 四角形を描画する
      |> OpenCV.putText!(
        prediction.class,
        [left + 6, top + 26],
        OpenCV.cv_FONT_HERSHEY_SIMPLEX,
        0.8,
        [0, 0, 255],
        thickness: 2
      ) # ラベル文字を書く
    end)
    |> OpenCV.cvtColor!(OpenCV.cv_COLOR_BGR2RGB)
  end
end
```

## 結果表示

検出結果の一覧と、結果の枠を書き込まれた画像を取得し、アサインします

lib/yolo_web/live/yolo_live.ex
```elixir
...
  @impl true
  def handle_event("detect",_params,%{assigns: %{upload_file: binary}} = socket) do
    {predictions, drawed} = YoloApp.Worker.detect(binary)
    socket =
      socket
      |> assign(:upload_file, drawed)
      |> assign(:ans, predictions)
    {:noreply, socket}
  end
...
```

lib/yolo_web/live/yolo_live.html.heex
```elixir
        <%= for {ans, index} <- Enum.with_index(@ans) do %>
          <h5><%= "#{index + 1}: " <> to_string(ans.class) %></h5>
        <% end %>
```

## 実行結果

こんな感じで実行できました！

![yolo.gif](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/1f1f8fac-d065-84d6-4aee-8947b34e7e9a.gif)

## 処理速度の向上

実は最初実装したとき、閾値による抽出処理は以下のようになっていました

```elixir
      predictions
      |> Nx.to_batched_list(1)
      |> Enum.filter(fn t ->
        t[0][4]
        |> Nx.to_number()
        |> Kernel.>(score_threshold)
      end)
```

しかし、これだとテンソルの 0 次元軸についてループするため、非常に遅いです

ちょっと計測してみましょう

measure 関数で抽出処理を囲むと、、、

```elixir
    predictions = measure(fn -> filter_predictions(predictions, 0.8) end)
```

なんと 9.21 sec も抽出に掛かっています

```
Time: 9212428
```

なぜこんなに掛かっているのか

YOLOv3 の検出結果は以下のような行列で与えられます

|X座標|Y座標|幅|高さ|スコア|クラス1確信度|クラス２確信度|...|クラス80確信度|
|--:|--:|--:|--:|--:|--:|--:|:--:|--:|
|0.2|0.3|0.05|0.03|0.3|0.8|0.0|...|0.0|
|0.1|0.4|0.12|0.2|0.92|0.0|0.0|...|0.0|
|0.6|0.05|0.3|0.2|0.95|0.0|0.9|...|0.0|
|...|...|...|...|...|...|...|...|...|


入力になっている `predictions` の shape を見てみると

```
{22743, 85}
```

これを普通にループしてしまうと、 2.2 万回も比較を行うことになるので、非常に遅いのです

高速化するためには、抽出をループではなく行列計算で行う必要があります

というわけで、以下のようなことをしています

- shape が {22743, 1}　で、中身が全て閾値の行列を作る

|閾値1|閾値2|閾値3|...|閾値22741|閾値22742|閾値22743|
|--:|--:|--:|:--:|--:|--:|--:|
|0.8|0.8|0.8|...|0.8|0.8|0.8|

```elixir
size =
  predictions
  |> Nx.shape()
  |> elem(0)

threshold_tensor =
  score_threshold
  |> Nx.tensor()
  |> Nx.broadcast({size})
```

- `predictions` からスコアの列だけ取り出す
shape は {22743, 1} になる

|スコア1|スコア2|スコア3|...|スコア22741|スコア22742|スコア22743|
|--:|--:|--:|:--:|--:|--:|--:|
|0.3|0.92|0.95|...|0.2|0.9|0.3|

```elixir
Nx.transpose(predictions)[4]
```

- 閾値の行列とスコアの行列を比較する
スコアが閾値を超えているところだけ 1 , それ以外は 0 の行列になる

|比較結果1|比較結果2|比較結果3|...|比較結果22741|比較結果22742|比較結果22743|
|--:|--:|--:|:--:|--:|--:|--:|
|0|1|1|...|0|1|0|

```elixir
greater = Nx.greater(predictions, threshold_tensor)
```

- 値が1のところのインデックスを取り出す

```
index_list =
    Nx.to_flat_list(greater)
    |> Enum.with_index()
    |> Enum.filter(fn {value, _} -> value == 1 end)
    |> Enum.map(&elem(&1, 1))
    |> Nx.tensor()
```

|インデックス1|インデックス2|インデックス3|...|インデックスn|
|--:|--:|--:|:--:|--:|
|1|2|15|...|22741|

- 指定したインデックスの推論結果だけを取り出す

```elixir
Nx.take(predictions, index_list)
```

このようにした結果、、、

```
Time: 121790
```

9.21 sec -> 0.12 sec なので、圧倒的な差です

如何にループではなく行列演算にすべきか、ということですね

## バックエンド

ちなみに、 main ブランチは EXLA を Nx のバックエンドにしています

Torchx でも問題なく動作しました

https://github.com/RyoWakabayashi/yolo-on-elixir-desktop/tree/feature/torchx

Torchx の場合の抽出処理（行列演算バージョン）の速さは

```
Time: 54847
```

なんと 0.05 sec

十分な回数試していないので、 Torchx のほうが速いとは言えませんが、
実用できる速度で動いています

ついでに、バックエンドを指定していない場合、
つまりバイナリバックエンドを使っている場合も見てみましょう

```
Time: 1909958
```

1.9  sec なので、思ったよりは速いです

行列演算は偉大ですね

## iOS での実行

残念ながら、まだ iOS では動かせていません

EXLA や evision が iOS では動かなさそうです

今後に期待ですね

## おわりに

evision を使った方法は一旦袋小路のようなので、
AxonOnnx を使ってみたいと思います