Ortex で ONNX モデルによる画像分類、物体検出 on Livebook

## はじめに

ONNX 形式のモデルを読み込むのに以前 Axon ONNX を使用しました

https://onnx.ai/

:::note info
Axon ONNX で物体検出を実行した記事は[こちら](https://qiita.com/RyoWakabayashi/items/a550249b4ace1cce11f5)

@piacerex さんによる画像分類の記事は[こちら](https://qiita.com/piacerex/items/0f6115fd2691586b2598)
:::

しかし、 Axon ONNX の README には以下のような記述があります

> See Ortex which provides full-blown compatibility for running ONNX models via ONNX Runtime bindings.

Chat-GPT による和訳

> Ortexは、ONNX Runtimeバインディングを通じてONNXモデルを完全に実行できる互換性を提供します。

https://github.com/mortont/axon_onnx

というわけで、 Ortex を使ってみます

https://github.com/elixir-nx/ortex

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

https://github.com/RyoWakabayashi/elixir-learning/blob/main/livebooks/ortex/ortex_resnet.livemd

https://github.com/RyoWakabayashi/elixir-learning/blob/main/livebooks/ortex/ortex_yolov2.livemd

## 画像分類

### モジュールのインストール

必要なモジュールをインストールします

```elixir
Mix.install(
  [
    {:exla, "~> 0.7"},
    {:stb_image, "~> 0.6"},
    {:req, "~> 0.5"},
    {:kino, "~> 0.13"},
    {:ortex, "~> 0.1"}
  ],
  config: [nx: [default_backend: EXLA.Backend]]
)
```

:::note info
StbImage は画像読込やリサイズに使用していますが、 Evision などで代用可能です
:::

### クラス一覧の取得

今回分類したいクラスの一覧を Web からダウンロードしてきます

```elixir
classes =
  "https://s3.amazonaws.com/deep-learning-models/image-models/imagenet_class_index.json"
  |> Req.get!()
  |> Map.get(:body)
```

### モデルの読込

ONNX のモデルをダウンロードし、読み込みます

```elixir
model_path = "/tmp/resnet.onnx"

unless File.exists?(model_path) do
  "https://media.githubusercontent.com/media/onnx/models/main/validated/vision/classification/resnet/model/resnet18-v1-7.onnx?download=true"
  |> Req.get!(connect_options: [timeout: 300_000], into: File.stream!(model_path))
end

model = Ortex.load(model_path)
```

実行結果（モデルの情報）では入出力形式を確認できます

```elixir
#Ortex.Model<
  inputs: [{"data", "Float32", [nil, 3, 224, 224]}]
  outputs: [{"resnetv15_dense0_fwd", "Float32", [nil, 1000]}]>
```

`Nx.Serving` で AI 処理を提供できるように準備します

```elixir
serving = Nx.Serving.new(Ortex.Serving, model)
```

### 画像の読込

画像を Web からダウンロードして読み込みます

ONNX モデルの入力に合わせて 224 * 224 にリサイズしておきます

```elixir
img_path = "/tmp/shark.jpg"

img_tensor =
  "https://www.collinsdictionary.com/images/full/greatwhiteshark_157273892.jpg"
  |> Req.get!()
  |> then(&StbImage.read_binary!(&1.body))
  |> StbImage.resize(224, 224)
  |> StbImage.to_nx()

Kino.Image.new(img_tensor)
```

実行結果

![shark.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/bd0381d3-6ac3-bf91-aa0c-c561bc3b5abc.png)

画像のチャネル数を取得します

RGB （透過なし）なら 3 、 RGBA （透過あり）なら 4 が返ってきます

```elixir
nx_channels = Nx.axis_size(img_tensor, 2)
```

透過ありの場合は A チャネルを切り離します（今回のサメ画像の場合はチャネル数 3 なのでそのまま）

RGB の値を正規化、標準化し、色の次元が先頭になるよう転置します

```elixir
img_tensor =
  case nx_channels do
    3 -> img_tensor
    4 -> Nx.slice(img_tensor, [0, 0, 0], [224, 224, 3])
  end
  |> Nx.divide(255)
  |> Nx.subtract(Nx.tensor([0.485, 0.456, 0.406]))
  |> Nx.divide(Nx.tensor([0.229, 0.224, 0.225]))
  |> Nx.transpose()
```

画像データをバッチに変換します

```elixir
batch = Nx.Batch.stack([img_tensor])
```

### 推論

`Nx.Serving.run` で推論します

```elixir
serving
|> Nx.Serving.run(batch)
|> Nx.backend_transfer()
|> elem(0)
|> Nx.flatten()
|> Nx.argsort()
|> Nx.reverse()
|> Nx.slice([0], [5])
|> Nx.to_flat_list()
|> Enum.with_index()
|> Enum.map(fn {no, index} -> {index, Map.get(classes, to_string(no))} end)
|> dbg()
```

`Nx.Serving.run` の結果は以下のような形式になっているため、 Nx.BinaryBackend に変換してから Tuple の先頭を取得しています

:::note info
`Nx.argsort` が `Ortex.Backend` では未実装なので `Nx.BinaryBackend` に変換する必要があります
:::


```elixir
{#Nx.Tensor<
   f32[1][1000]
   Ortex.Backend
   [
     [5.107775688171387, 7.190211296081543, 17.523210525512695, ...]
   ]
 >}
```

その後は確信度上位5つを取得してからクラス名に変換しています

最終的な実行結果は以下のようになります

```elixir
[
  {0, ["n01484850", "great_white_shark"]},
  {1, ["n01491361", "tiger_shark"]},
  {2, ["n01494475", "hammerhead"]},
  {3, ["n01496331", "electric_ray"]},
  {4, ["n02074367", "dugong"]}
]
```

## Axon ONNX の場合との差分

インストールするモジュールを ortex に変更します

```diff_elixir
-    {:axon_onnx, "~> 0.4", git: "https://github.com/mortont/axon_onnx/"}
+    {:ortex, "~> 0.1"}
```

モデルの読み込み時、 `Nx.Serving` を用意します

```diff_elixir
- {model, params} = AxonOnnx.import(model_path)
+ model = Ortex.load(model_path)
+ serving = Nx.Serving.new(Ortex.Serving, model)
```

画像をバッチに変換します

```diff_elixir
img_tensor =
  case nx_channels do
    3 -> img_tensor
    4 -> Nx.slice(img_tensor, [0, 0, 0], [224, 224, 3])
  end
  |> Nx.divide(255)
  |> Nx.subtract(Nx.tensor([0.485, 0.456, 0.406]))
  |> Nx.divide(Nx.tensor([0.229, 0.224, 0.225]))
  |> Nx.transpose()
- |> Nx.new_axis(0)
+
+ batch = Nx.Batch.stack([img_tensor])
```

推論結果は OrtexBackend の Tuple になっています

```diff_elixir
- model
- |> Axon.predict(params, img_tensor)
+ serving
+ |> Nx.Serving.run(batch)
+ |> Nx.backend_transfer()
+ |> elem(0)
```

## 物体検出の場合

上記の Axon ONNX の場合との差分と同様に Axon ONNX の場合のコードを書き換えれば良いだけです

https://qiita.com/RyoWakabayashi/items/a550249b4ace1cce11f5

```elixir
Mix.install(
  [
    {:exla, "~> 0.7"},
    {:stb_image, "~> 0.6"},
    {:req, "~> 0.5"},
    {:kino, "~> 0.13"},
    {:ortex, "~> 0.1"}
  ],
  config: [nx: [default_backend: EXLA.Backend]]
)
```

```elixir
classes =
  "https://raw.githubusercontent.com/pjreddie/darknet/master/data/coco.names"
  |> Req.get!()
  |> then(&String.split(&1.body, "\n"))
  |> Enum.map(&String.trim(&1))
  |> Enum.filter(&(String.length(&1) > 0))

num_classes = Enum.count(classes)

model_path = "/tmp/yolov2.onnx"

unless File.exists?(model_path) do
  "https://media.githubusercontent.com/media/onnx/models/main/validated/vision/object_detection_segmentation/yolov2-coco/model/yolov2-coco-9.onnx?download=true"
  |> Req.get!(connect_options: [timeout: 300_000], into: File.stream!(model_path))
end

model = Ortex.load(model_path)

serving = Nx.Serving.new(Ortex.Serving, model)

anchors =
  Nx.tensor([
    [0.57273, 0.677385],
    [1.87446, 2.06253],
    [3.33843, 5.47434],
    [7.88282, 3.52778],
    [9.77052, 9.16828]
  ])

num_anchors =
  anchors
  |> Nx.shape()
  |> elem(0)

anchors_tensor = Nx.reshape(anchors, {1, 1, 1, num_anchors, 2})
```

```elixir
img_tensor =
  "https://raw.githubusercontent.com/pjreddie/darknet/master/data/dog.jpg"
  |> Req.get!()
  |> then(&StbImage.read_binary!(&1.body))
  |> StbImage.resize(416, 416)
  |> StbImage.to_nx()

nx_channels = Nx.axis_size(img_tensor, 2)

img_tensor =
  case nx_channels do
    3 -> img_tensor
    4 -> Nx.slice(img_tensor, [0, 0, 0], [416, 416, 3])
  end
  |> Nx.divide(255)
  |> Nx.transpose(axes: [2, 0, 1])

batch = Nx.Batch.stack([img_tensor])

feats =
  serving
  |> Nx.Serving.run(batch)
  |> Nx.backend_transfer()
  |> elem(0)
```

これ以降は Axon ONNX の場合と同様です

## まとめ

Axon ONNX を使ったコードから一定のルールで変換可能であるため、すぐに実装できました

今後はこちらをメインに使っていこうと思います