*Interp移植録 - 物体検出 / YOLOv8 (OnnxInterp)

0.Prologue

暇つぶしに、興味を引いた DNNアプリを *Interpに移植して遊んでいる。
本稿はその雑記&記録。

またもや新たなYOLOが登場した。YOLOv8である。まだ詳細な技術情報は公開されていないが、GitHubにはPythonプロジェクトが登録されている。それを Elixirで動かしてみようと思う。たぶん、年末にお遊びで作った「怒涛のよろず祭り」を流用すれば、大して時間は掛からないだろう。

1.Original Work

YOLOv8は、元祖DarkNet系モデルではなく、YOLOv5の流れを汲む Ultralytcs系。先にも述べたが、まだ正式な技術情報は公開されていないため詳細は分からない。GitHubに掲載のグラフをよると、精度が改善されたのかなぁ?

• Ultralytics YOLOv8
  https://github.com/ultralytics/ultralytics

2.準備

YOLOv8の ONNXモデルは、本家プロジェクトに同梱のCLIyoloを用いて作成する。Python環境に "ultralytics"をインストールし、

pip install ultralytics

下記のコマンド・ラインを実行すれば良い。デフォルトでは、入力画像 640x640のモデルが作成される。

yolo task=detect mode=export model=yolov8n.pt format=onnx

YOLOv8は、ONNXモデルのほかに Tensorflow Saved model, Tensorflow lite, OpenVINOなどのモデルにも変換できるようだ。詳しくは本家プロジェクトに同梱の jupyterノートブックを参照のこと。

3.OnnxInterp用のLivebookノート

Mix.installの依存リストに記述するモジュールは下記の通り。

File.cd!(__DIR__)
# for windows JP
System.shell("chcp 65001")

Mix.install([
  {:onnx_interp, "~> 0.1.8"},
  {:cimg, "~> 0.1.14"},
  {:postdnn, "~> 0.1.5"},
  {:kino, "~> 0.7.0"}
])

モデルの入出力はオーソドックスな仕様。
後処理では、モデルの出力をBBoxとクラス評価値に分割し、NMSでふるい掛けして推論結果を得る。BBoxの座標値/サイズが 640x640の座標系で表されているので、これを後続の表示処理で扱い易いように画像サイズに対する比率に変換する(ratio_box/1)。

[モデル・カード]
inputs:
　[0] f32:{1,3,640,640} - RGB画像,NCHWレイアウト,画素値を{0.0～1.0}に正規化
outputs:
　[0] f32:{1,84,8400} - 対象物のBBox[,0:4,]とそのクラス毎の評価値[,4:,],BBoxの座標値/サイズは 640x640の座標系で表記

defmodule YOLOv8 do
  @moduledoc """
  Original work:
    Ultralytics YOLOv8 - https://github.com/ultralytics/ultralytics
  """

  @width 640
  @height 640

  alias OnnxInterp, as: NNInterp
  use NNInterp, label: "./model/coco.label",
    model: "./model/yolov8n.onnx",
    url: "https://github.com/shoz-f/onnx_interp/releases/download/models/yolov8n.onnx",
    inputs: [f32: {1, 3, @width, @height}],
    outputs: [f32: {1, 84, 8400}]

  def apply(img) do
    # preprocess
    input0 = CImg.builder(img)
      |> CImg.resize({@width, @height})
      |> CImg.to_binary([{:range, {0.0, 1.0}}, :nchw])

    # prediction
    output0 = session()
      |> NNInterp.set_input_tensor(0, input0)
      |> NNInterp.invoke()
      |> NNInterp.get_output_tensor(0)
      |> Nx.from_binary(:f32) |> Nx.reshape({84, 8400})

    # postprocess
    scores = Nx.transpose(output0[4..-1//1])
    boxes = Nx.transpose(output0[0..3])

    NNInterp.non_max_suppression_multi_class( __MODULE__,
      Nx.shape(scores), Nx.to_binary(boxes), Nx.to_binary(scores)
    )
    |> ratio_box()
  end

  defp ratio_box({:ok, result}) do
    clamp = fn x -> min(max(x, 0.0), 1.0) end

    {
      :ok,
      Enum.reduce(Map.keys(result), result, fn key, map ->
        Map.update!(map, key, &Enum.map(&1, fn [score, x1, y1, x2, y2, index] ->
            [score, clamp.(x1 / @width), clamp.(y1 / @height), clamp.(x2 / @width), clamp.(y2 / @height), index ]
          end)
        )
      end)
    }
  end
end

デモ・モジュール LiveYOLOv8は、YOLOv8から受け取ったBBoxを入力画像上に描画する。

defmodule LiveYOLOv8 do
  @palette CImg.Util.rand_palette("./model/coco.label")

  def run(path) do
    img = CImg.load(path)

    with {:ok, res} <- YOLOv8.apply(img) do
      IO.inspect(res)
      Enum.reduce(res, CImg.builder(img), &draw_item(&1, &2))
      |> CImg.display_kino(:jpeg)
    end
  end

  defp draw_item({name, boxes}, canvas) do
    color = @palette[name]

    Enum.reduce(boxes, canvas, fn [_score, x1, y1, x2, y2, _index], canvas ->
      [x1, y1, x2, y2] = PostDNN.clamp([x1, y1, x2, y2], {0.0, 1.0})

      CImg.fill_rect(canvas, x1, y1, x2, y2, color, 0.4)
    end)
  end
end

4.デモンストレーション

YOLOv8を起動する。

YOLOv8.start_link([])

画像を与え、YOLOv8を実行する。

DemoYOLOv8.run("bus.jpg")

5.Epilogue

新たなYOLOが発表されたので、早速OnnxInterpに移植してみた。
目を見張るほど速いという感じはないが、対モデル・サイズの精度は良いように思う。どのようなアイデアが組み込まれたのだろうか? 技術公開に期待したい。

話は変わるが、年末に作った「よろず祭り」が良い仕事をしてくれた。YOLOv8を動かす為に書いたコードはほんの数行で済んだ。この移植録もそんな風にサクサクと書ければなぁ

Appendix

OnnxInterpのノート
https://github.com/shoz-f/onnx_interp/blob/main/demo_yolov8/YOLOv8.livemd