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

はじめに

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

全コードはこちら

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

実行環境

• macOS Monterey 12.4
• Elixir 1.13.4
• Erlang OTP24

フォーク

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

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

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

依存パッケージの変更

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

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

  ...
  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) を使っている箇所を削除します

詳細はこちら

モデルのダウンロード

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

#!/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

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

...
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

...
    {:ok, sup} = Supervisor.start_link([YoloApp.Store], name: __MODULE__, strategy: :one_for_one)
...

ワーカーの実装

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

lib/yolo_app/worker.ex

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

...
  @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

        <%= for {ans, index} <- Enum.with_index(@ans) do %>
          <h5><%= "#{index + 1}: " <> to_string(ans.class) %></h5>
        <% end %>

実行結果

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

処理速度の向上

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

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

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

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

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

    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

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

Nx.transpose(predictions)[4]

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

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

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

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

Nx.take(predictions, index_list)

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

Time: 121790

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

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

バックエンド

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

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

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

Time: 54847

なんと 0.05 sec

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

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

Time: 1909958

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

行列演算は偉大ですね

iOS での実行

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

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

今後に期待ですね

おわりに

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