MacBook 上の Livebook で YOLO Elixir による物体検出を実行する

はじめに

Google Colaboratory 上で YOLO Elixir を実行しました

これを MacBook Pro M2 上で実行したいと思います

モデルの変換

YOLOv8 の .pt ファイルを ONNX 形式に変換します

macOS の環境を汚したくないので、 Docker を使います

実装したコンテナはこちら

以下のようなディレクトリー構成になっています

yolov8/
　├ models/
　├ convert.sh
　├ docker-compose.yml
　└ Dockerfile

Dockerfile の内容は以下のとおりです

FROM ubuntu:22.04

ENV DEBIAN_FRONTEND noninteractive

ENV TZ=Asia/Tokyo

# hadolint ignore=DL3008
RUN apt-get update \
        && apt-get install --no-install-recommends --no-install-suggests -y \
                ca-certificates \
                git \
                libopencv-dev \
                python3-pip \
        && apt-get clean \
        && rm -rf /var/lib/apt/lists/*

WORKDIR /work

RUN git clone https://github.com/poeticoding/yolo_elixir.git

WORKDIR /work/yolo_elixir

# hadolint ignore=DL3006,DL3013
RUN pip install --no-cache-dir --upgrade pip \
    && pip install --no-cache-dir --requirement python/requirements.txt

COPY convert.sh /work/yolo_elixir/convert.sh

RUN chmod +x /work/yolo_elixir/convert.sh

CMD ["./convert.sh"]

Ubuntu 22.04 上に YOLO Elixir のリポジトリーをクローンしてきて、変換を実行しています

convert.sh の内容は以下のとおり

#!/bin/bash

python3 python/yolov8_to_onnx.py n
python3 python/yolov8_to_onnx.py x

cp -r ./models/* /tmp/models/

YOLOv8 と YOLOv8x の変換を実行し、変換結果を /tmp/models/ にコピーしています

docker-compose.yml の内容は以下のとおり

---
services:
  yolov8:
    container_name: yolov8
    build:
      context: .
    tty: true
    volumes:
      - ./models:/tmp/models

ローカルの ./models をコンテナ内の /tmp/models にマウントしているため、変換結果のモデルファイルは ./models ディレクトリー内に出力されます

以下のコマンドで変換を実行します

docker compose up --build

実行後、以下のファイルが models ディレクトリー内に出力されます

• yolov8n_classes.json
• yolov8n.onnx
• yolov8n.pt
• yolov8x_classes.json
• yolov8x.onnx
• yolov8x.pt

Livebook の起動

MacBook 上で起動するため、 Livebook をソースコードから実行します

git clone -b v0.14.5 https://github.com/livebook-dev/livebook.git
cd livebook
mix setup
ELIXIR_ERL_OPTIONS="-epmd_module Elixir.Livebook.EPMD" mix phx.server

セットアップ

新しいノートブックを開き、セットアップセルで必要なモジュールをインストールします

Mix.install(
  [
    {:yolo, ">= 0.0.0"},
    {:yolo_fast_nms, "~> 0.1"},
    {:emlx, github: "elixir-nx/emlx"},
    {:evision, "~> 0.2.0"},
    {:kino, "~> 0.14.2"}
  ],
  config: [
    nx: [default_backend: EMLX.Backend]
  ]
)

実行環境が MacBook なので、 Apple Silicon を有効活用できる EMLX を使います

後述しますが、 YOLO Elixir に関しては EMLX の効果はあまりありません

EXLA を使っても同じ速度になります

モデルの読込

入力フォームを作り、変換したモデルファイルとクラス定義ファイルを選択します

onnx_input = Kino.Input.file("ONNX")
classes_json_input = Kino.Input.file("CLASSES JSON")

Kino.Layout.grid([onnx_input, classes_json_input], columns: 2)

モデルを読み込みます

model = YOLO.load([
  model_path: onnx_input |> Kino.Input.read() |> Map.get(:file_ref) |> Kino.Input.file_path(), 
  classes_path: classes_json_input |> Kino.Input.read() |> Map.get(:file_ref) |> Kino.Input.file_path()
])

画像の読込

画像選択のフォームを作成します

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

今回は YOLO Elixir のガイドに使用されていた画像を使用します

画像を行列に変換します

image =
  image_input
  |> Kino.Input.read()
  |> Map.get(:file_ref)
  |> Kino.Input.file_path()
  |> File.read!()

mat = Evision.imdecode(image, Evision.Constant.cv_IMREAD_COLOR())

実行結果

物体検出の実行

物体検出を実行し、結果を整形します

nms_fun: &YoloFastNMS.run/3 により、 NMS (Non-Max Suppression) を高速化しています

objects =
  model
  |> YOLO.detect(mat, nms_fun: &YoloFastNMS.run/3)
  |> YOLO.to_detected_objects(model.classes)

実行結果

[
  %{
    class: "person",
    prob: 0.7682794332504272,
    bbox: %{h: 147, w: 69, cx: 699, cy: 579},
    class_idx: 0
  },
  %{
    class: "person",
    prob: 0.7557899951934814,
    bbox: %{h: 213, w: 81, cx: 609, cy: 777},
    class_idx: 0
  },
  %{
    class: "person",
    prob: 0.7420753836631775,
    bbox: %{h: 225, w: 90, cx: 468, cy: 849},
    class_idx: 0
  },
  ...
]

検出結果の描画

検出結果を画像上に描画します

四角形の色は適当にクラス毎に違う色にしています

draw_objects = fn mat, objects ->
  objects
  |> Enum.reduce(mat, fn %{class: class, prob: prob, bbox: bbox, class_idx: class_idx}, drawed_mat ->
    %{w: w, h: h, cx: cx, cy: cy} = bbox
    left = cx - div(w, 2)
    top = cy - div(h, 2)
    right = left + w
    bottom = top + h
  
    score = round(prob * 100) |> Integer.to_string()
  
    color = {
      case rem(class_idx, 3) do
        0 -> 0
        1 -> 128
        2 -> 255
      end,
      case rem(80 - class_idx, 4) do
        0 -> 0
        1 -> 30
        2 -> 60
        3 -> 90
      end,
      case rem(40 + class_idx, 5) do
        0 -> 255
        1 -> 196
        2 -> 128
        3 -> 64
        4 -> 0
      end
    }
  
    text = class <> ":" <> score
    font = Evision.Constant.cv_FONT_HERSHEY_SIMPLEX()
    font_scale = 1
    font_thickness = 2
    {{tw, th}, _} = Evision.getTextSize(text, font, font_scale, font_thickness)
  
    drawed_mat
    |> Evision.rectangle(
      {left, top},
      {right, bottom},
      color,
      thickness: 10
    )
    |> Evision.rectangle(
      {left - 5, top - th - 10},
      {left + tw + 5, top},
      color,
      thickness: -1
    )
    |> Evision.putText(
      text,
      {left, top - 5},
      font,
      font_scale,
      {255, 255, 255},
      thickness: font_thickness
    )
  end)
end

draw_objects.(mat, objects)

実行結果

MacBook 上でも問題なく YOLO Elixir が実行できました

EMLX と EXLA の速度比較

せっかくなので、 Nx のバックエンドに EMLX を使った場合と EXLA を使った場合でどの程度差があるのか検証してみます

新しいノートブックを開き、以下のコードをセットアップセルで実行します

Mix.install([
  {:yolo, ">= 0.0.0"},
  {:yolo_fast_nms, "~> 0.1"},
  {:exla, "~> 0.9"},
  {:emlx, github: "elixir-nx/emlx"},
  {:evision, "~> 0.2"},
  {:kino_benchee, "~> 0.1"}
])

変換した YOLOv8x のモデルを使用します

v8x_onnx_input = Kino.Input.file("YOLOv8x ONNX")
v8x_classes_json_input = Kino.Input.file("YOLOv8x CLASSES JSON")

Kino.Layout.grid([v8x_onnx_input, v8x_classes_json_input], columns: 2)

v8x_model = YOLO.load([
  model_path: v8x_onnx_input |> Kino.Input.read() |> Map.get(:file_ref) |> Kino.Input.file_path(), 
  classes_path: v8x_classes_json_input |> Kino.Input.read() |> Map.get(:file_ref) |> Kino.Input.file_path()
])

画像を選択し、行列に変換します

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

mat =
  image_input
  |> Kino.Input.read()
  |> Map.get(:file_ref)
  |> Kino.Input.file_path()
  |> File.read!()
  |> Evision.imdecode(Evision.Constant.cv_IMREAD_COLOR())

ベンチマーク用のモジュールを定義します

defmodule MyBenchmark do
  def detect(model, mat) do
    model
    |> YOLO.detect(mat, nms_fun: &YoloFastNMS.run/3)
    |> YOLO.to_detected_objects(model.classes)
  end

  def run(model, mat) do
    Benchee.run(
      %{
        "binary" => fn ->
          Nx.default_backend(Nx.BinaryBackend)
          detect(model, mat)
        end,
        "exla" => fn ->
          Nx.default_backend(EXLA.Backend)
          detect(model, mat)
        end,
        "emlx" => fn ->
          Nx.default_backend(EMLX.Backend)
          detect(model, mat)
        end
      },
      memory_time: 2,
      reduction_time: 2
    )
  end
end

Nx.default_backend で使用するバックエンドを切り替えています

ベンチマークを実行します

MyBenchmark.run(v8x_model, mat)

バイナリーバックエンドは遅いですが、 EMLX と EXLA にはほぼ差がありません

実は YOLO Elixir は ortex を使用しており、 ortex は推論時、明示的に Ortex.Backend を使っています

そのため、 Nx のデフォルトバックエンドを指定していても、前処理や後処理でしか使用されません

まとめ

MacBook 上でも YOLO Elixir による物体検出を実行できました

ただし、 EMLX による高速化の効果はあまりありませんでした