Elixir の evision で YOLOv7 (ONNX) の物体検出を実行する

はじめに

以前、 YOLOv3 を evision で動かす記事を書きました

今回は YOLOv7 (ONNX) を evision で動かします

また、推論後の整形処理を以前は Enum.map を使ったループ処理で実装しましたが、今回はそこを Nx による行列演算で実装します

実行にはいつも通り Livebook を使います

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

evision

evision は OpenCV の Elixir バインディングです

evision を使うことで、 Python で OpenCV を使うのと同じように Elixir で画像処理が実装できます

また、 OpenCV の DNN 機能も利用できるため、 OpenCV で読み込むことができる形式であれば、機械学習モデルの推論を実行することも可能です

今回は PyTorch 形式の YOLOv7x モデルを ONNX 形式に変換し、 evision で読み込みます

変換用のコンテナ、シェルスクリプトはこちら

YOLOv7 の ONNX 形式への変換

こちらの YOLOv7 PyTorch 実装のリポジトリー、モデルを利用します

このリポジトリーにある export.py を使うことで、 YOLOv7 を TorchScript や ONNX や CoreML 形式に変換できます

自分のローカル環境を汚したくないのと、何回でも再現可能にするため、コンテナ上で変換することにします

以下のような配置でファイルを準備します

yolov7/
　├ models/
　├ conversion.sh
　├ docker-compose.yml
　└ Dockerfile

Dockerfile

FROM ubuntu:20.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 \
                libopencv-dev \
                git \
                wget \
                python3.8 \
                python3-pip \
        && apt-get clean \
        && rm -rf /var/lib/apt/lists/* \
        && update-alternatives --install /usr/bin/python python /usr/bin/python3.8 1

WORKDIR /work

# YOLOv7 リポジトリーのクローン
RUN git clone https://github.com/WongKinYiu/yolov7.git

WORKDIR /work/yolov7

# 学習済 YOLOv7x モデルのダウンロード
RUN wget --quiet --continue \
        --timestamping https://github.com/WongKinYiu/yolov7/releases/download/v0.1/yolov7x.pt \
        --output-document ./yolov7x.pt

# 必要な Python パッケージのインストール
# hadolint ignore=DL3006,DL3013
RUN pip install --no-cache-dir --upgrade pip \
    && pip install --no-cache-dir --requirement requirements.txt \
    && pip install --no-cache-dir onnx onnx-simplifier

# 変換用シェルスクリプト
COPY convert.sh /work/yolov7/convert.sh

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

CMD ["./convert.sh"]

実行を簡単にするために docker compose を使います

docker-compose.yml

---

version: "2.3"

services:
  yolov7:
    container_name: yolov7
    build:
      context: .
    tty: true
    volumes:
      - ./models:/tmp/models # 変換したモデルファイルを格納するためのマウント

変換用シェルスクリプトは export.py を所定のパラメータで実行し、結果をマウントしたディレクトリーにコピーしています

変換時のパラメータに --end2end などを指定すると NMS のレイヤーが evision では読み込めずにエラーが発生します

convert.sh

#!/bin/bash

python export.py \
  --weights yolov7x.pt \
  --img-size 640 640 \
  --grid \
  --simplify

cp yolov7x.onnx /tmp/models/

準備したコンテナを実行します

docker compose up

しばらくすると、 models ディレクトリー配下に yolov7x.onnx が出力されます

このファイルを Livebook から参照可能な場所にコピーしておきます

セットアップ

Livebook を起動し、新規ノートブックのセットアップセルに以下のコードを入力して実行します

Mix.install(
  [
    {:nx, "~> 0.5"},
    {:exla, "~> 0.5"},
    {:evision, "~> 0.1"},
    {:req, "~> 0.3"},
    {:kino, "~> 0.9"}
  ],
  config: [nx: [default_backend: EXLA.Backend]]
)

以下のモジュールがインストールされます

• Nx: 行列演算
• EXLA: 行列演算の高速化
• evision: OpenCV バインディング
• Req: Web へのアクセス
• Kino: Livebook の UI/UX

モデルの読み込み

変換しておいた yolov7x.onnx を読み込みます
（パスは各自変更してください）

net = Evision.DNN.readNet("/tmp/yolov7x.onnx")

推論実行時に使うため、出力層の名前を取得しておきます

out_names = Evision.DNN.Net.getUnconnectedOutLayersNames(net)

今回使うモデルは COCO データセットを学習しているため、 COCO　データセットのクラス名（ラベル）一覧を取得します

labels =
  "https://raw.githubusercontent.com/pjreddie/darknet/master/data/coco.names"
  # ダウンロード
  |> Req.get!()
  |> Map.get(:body)
  # バイナリを文字列に変換
  |> then(&Enum.join(for <<c::utf8 <- &1>>, do: <<c::utf8>>))
  # 改行コード区切りで配列に変換
  |> String.split("\n")
  # 空文字（最終行）を除外する
  |> Enum.filter(&(&1 != ""))

クラス数を確認しておくと、 COCO データセットの 80 クラスになっています

Enum.count(labels)

画像の取得

推論対象の画像をダウンロードし、 evision で読み込みます

img =
  "https://raw.githubusercontent.com/pjreddie/darknet/master/data/dog.jpg"
  |> Req.get!()
  |> Map.get(:body)
  |> Evision.imdecode(Evision.Constant.cv_IMREAD_COLOR())

後で使うため、画像のサイズを取得しておきます

{img_height, img_width, _} = Evision.Mat.shape(img)

推論

画像を推論のための形式に変換します

YOLOv7x の入力層 640 * 640 のサイズを指定しています

OpenCV で読み込んだ場合色空間が BGR になっているので、 RGB に変換するため swapRB を true にします

リザイズ時に画像の一部を切り捨てないよう crop は false にします

blob = Evision.DNN.blobFromImage(img, size: {640, 640}, swapRB: true, crop: false)

変換した画像データをモデルに入力し、推論を実行します

色が 0 から 255 だったのを 0 から 1 に変換するため、 scalefactor: 1 / 255 を指定します

predictions =
  net
  |> Evision.DNN.Net.setInput(
    blob,
    name: "",
    scalefactor: 1 / 255,
    mean: {0, 0, 0}
  )
  |> Evision.DNN.Net.forward(outBlobNames: out_names)

実行結果

[
  %Evision.Mat{
    channels: 1,
    dims: 3,
    type: {:f, 32},
    raw_type: 5,
    shape: {1, 25200, 85},
    ref: #Reference<0.4157835992.3598581767.83567>
  }
]

1 (推論した画像の数) * 25200 (検出した領域の数) * 85 (座標情報、座標のスコア、クラス毎のスコア) の行列が配列に入っています

推論結果の整形

スコアが閾値を超える領域だけを抽出します

# スコア閾値
score_threshold = 0.6

推論結果は evision で扱う Matrix (行列)形式になっているので、 Nx で扱う Tensor (テンソル)形式に変換します

演算を高速化するため、 EXLA.Backend を指定します

predictions_tensor =
  predictions
  |> Enum.at(0)
  |> Evision.Mat.to_nx(EXLA.Backend)

まず、各領域の座標のスコア（その位置に物体がある事に対する確信度）を取り出します

1次元目は 1 なので [0] で全て取得できます

2次元目は全部を取りたいので [0..-1//1] （0番目から最後=後ろから1つ目=-1番目まで、一つずつ=//1）で指定します

最後の次元（長さ 85）は以下のような構成になっています

• 0: 領域の中心X座標
• 1: 領域の中心Y座標
• 2: 領域の幅
• 3: 領域の高さ
• 4: 座標のスコア
• 5: クラス0のスコア
• 6: クラス1のスコア
• ...
• 83: クラス78のスコア
• 84: クラス79のスコア

従って、全領域に対する座標のスコアを取り出す場合、範囲指定は [0, 0..-1//1, 4] となります

all_bbox_score_tensor = predictions_tensor[[0, 0..-1//1, 4]]

実行結果

#Nx.Tensor<
  f32[25200]
  EXLA.Backend<host:0, 0.817007110.871235591.8204>
  [2.037998456216883e-6, 3.654646434370079e-7, 4.855778570345137e-7, 1.1191846738256572e-7, 9.531960358799552e-7, 7.281367402356409e-7, 3.001308925831836e-7, 1.6525656576504844e-7, 1.6678282577231585e-7, 4.5229248257783183e-7, 1.539881736789539e-6, 5.390127739701711e-7, 2.3195359517558245e-6, 1.8462549178366316e-7, 2.187121168617523e-7, 6.263160230446374e-7, 4.838346967517282e-7, 1.1154595114248878e-7, 2.4176438273570966e-7, 2.965466023852059e-7, 2.17695159676623e-7, 1.2532504456430615e-7, 1.0492052382460315e-7, 1.8693350511966855e-7, 1.1833714808062723e-7, 3.2386957116159465e-8, 8.214190927446907e-8, 2.7507357458489423e-7, 3.320245696158963e-7, 1.8705850379774347e-7, 8.240031768025347e-9, 3.248266011723899e-8, 4.086846416839762e-8, 4.355639759978658e-8, 3.691250469728402e-7, 5.01677313025084e-8, 1.1205901984112643e-7, 1.0090556088471203e-6, 4.860818989982363e-6, 3.423071746055939e-7, 3.245352786507283e-7, 5.6435549566913323e-8, 7.35682164076934e-8, 2.0417004975570308e-8, 4.975055958311714e-7, 5.4937714821789996e-8, 7.1817254365669214e-9, 4.2377878983757e-9, 1.036702670376144e-7, 5.549345019062457e-7, ...]
>

これらのスコアのうち何番目のものが閾値を超えているか、が取得したいものです

Nx.greater を使い、閾値を超えていれば 1 、閾値以下なら 0 にします

greater_tensor = Nx.greater(all_bbox_score_tensor, score_threshold)

実行結果

#Nx.Tensor<
  u8[25200]
  EXLA.Backend<host:0, 0.817007110.871235591.8048>
  [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, ...]
>

閾値を超えているスコアの数を取得します

0 と 1 になっているので、単純に全ての合計 = 閾値を超えた数です

greater_count =
  greater_tensor
  |> Nx.sum()
  |> Nx.to_number()

実行結果

25

Nx.argsort を使うと、スコアが大きい順に並べたとき、元の順で何番目だったか、が取得できます

その結果の先頭から 25 個を取れば、閾値を超えていたのが何番目の領域だったか、が取得できます

greater_indices = Nx.argsort(greater_tensor, direction: :desc)[[0..(greater_count - 1)]]

実行結果

#Nx.Tensor<
  s64[25]
  EXLA.Backend<host:0, 0.817007110.871235591.8054>
  [21149, 21150, 22749, 22750, 24094, 24095, 24188, 24189, 24245, 24265, 24266, 24494, 24495, 24588, 24589, 24645, 24665, 24666, 24894, 24895, 24988, 24989, 25045, 25065, 25066]
>

Nx.take を使い、取得できたインデックス（何番目か）を指定して、閾値を超えた領域だけを抽出します

greater_predictions_tensor = Nx.take(predictions_tensor[0], greater_indices, axis: 0)

実行結果

#Nx.Tensor<
  f32[25][85]
  EXLA.Backend<host:0, 0.817007110.871235591.8060>
  [
    [482.181884765625, 136.43838500976562, 187.73941040039062, 106.85380554199219, 0.7066338062286377, 5.764405614172574e-6, 2.676756366781774e-6, 0.20150047540664673, 4.91483551741112e-5, 8.143885679601226e-6, 3.5266042686998844e-4, 5.971808150206925e-6, 0.7939861416816711, 3.342183481436223e-5, 1.2045237554048072e-6, 2.811464582919143e-6, 2.4691398721188307e-6, 4.378389348858036e-6, 3.23197127727326e-6, 2.5822807003805792e-8, 2.9516680655206073e-8, 7.304785185624496e-7, 2.7804278488474665e-6, 3.9378952010338253e-7, 2.2409756638808176e-5, 2.255292429254041e-6, 2.3871859866630984e-7, 1.8708020434132777e-5, 1.8235708694191999e-6, 7.666907322345651e-7, 7.7971486689421e-7, 2.251765636174241e-6, 2.4115115593303926e-7, 2.4434334591205698e-6, 1.7157055935967946e-6, 2.187726977354032e-6, 5.534025603992632e-6, 2.3630052510270616e-6, 7.76864089857554e-6, 1.5734760836494388e-6, 2.3528007204731693e-6, 1.0293243576597888e-5, 2.7150198889103194e-7, 5.863289516128134e-7, 4.191058451397112e-6, 1.365727030133712e-6, 8.582446753280237e-6, 2.27843020184082e-6, 2.6452306656210567e-7, 2.59688681580883e-6, ...],
    ...
  ]
>

改めて、抽出された領域の座標スコアを取得します

bbox_score_tensor = greater_predictions_tensor[[0..-1//1, 4]]

実行結果

#Nx.Tensor<
  f32[25]
  EXLA.Backend<host:0, 0.817007110.871235591.8064>
  [0.7066338062286377, 0.8829199075698853, 0.8605510592460632, 0.8959879279136658, 0.7967941164970398, 0.8849864602088928, 0.970268189907074, 0.9656566381454468, 0.9546226263046265, 0.9626724123954773, 0.9419291615486145, 0.8611856698989868, 0.8824902176856995, 0.9743650555610657, 0.9712815284729004, 0.953478991985321, 0.9630793333053589, 0.9396751523017883, 0.6563085913658142, 0.8374518156051636, 0.9759773015975952, 0.9708778858184814, 0.9428169131278992, 0.959770679473877, 0.9291824698448181]
>

どれも閾値(0.6)を超えていることが確認できます

クラスの判定

抽出できた領域のクラス毎のスコアを取得します

クラス毎のスコアが5番目以降の全てなので、 [5..-1//1] で取得きます

class_score_tensor = greater_predictions_tensor[[0..-1//1, 5..-1//1]]

実行結果

#Nx.Tensor<
  f32[25][80]
  EXLA.Backend<host:0, 0.817007110.871235591.8066>
  [
    [5.764405614172574e-6, 2.676756366781774e-6, 0.20150047540664673, 4.91483551741112e-5, 8.143885679601226e-6, 3.5266042686998844e-4, 5.971808150206925e-6, 0.7939861416816711, 3.342183481436223e-5, 1.2045237554048072e-6, 2.811464582919143e-6, 2.4691398721188307e-6, 4.378389348858036e-6, 3.23197127727326e-6, 2.5822807003805792e-8, 2.9516680655206073e-8, 7.304785185624496e-7, 2.7804278488474665e-6, 3.9378952010338253e-7, 2.2409756638808176e-5, 2.255292429254041e-6, 2.3871859866630984e-7, 1.8708020434132777e-5, 1.8235708694191999e-6, 7.666907322345651e-7, 7.7971486689421e-7, 2.251765636174241e-6, 2.4115115593303926e-7, 2.4434334591205698e-6, 1.7157055935967946e-6, 2.187726977354032e-6, 5.534025603992632e-6, 2.3630052510270616e-6, 7.76864089857554e-6, 1.5734760836494388e-6, 2.3528007204731693e-6, 1.0293243576597888e-5, 2.7150198889103194e-7, 5.863289516128134e-7, 4.191058451397112e-6, 1.365727030133712e-6, 8.582446753280237e-6, 2.27843020184082e-6, 2.6452306656210567e-7, 2.59688681580883e-6, 4.199067461740924e-6, 2.094071049896229e-7, 4.377010327516473e-7, 5.298339260662033e-7, 1.1620704754022881e-6, ...],
    ...
  ]
>

分類するクラスが 80 種類なので、 25 領域 * 80 クラスの行列が取得できました

Nx.argmax を使い、各領域についてクラス毎に1番スコアが高かったインデックス = クラス番号を取得します

class_index_tensor = Nx.argmax(class_score_tensor, axis: 1)

実行結果

#Nx.Tensor<
  s64[25]
  EXLA.Backend<host:0, 0.817007110.871235591.8068>
  [7, 7, 7, 7, 7, 7, 1, 1, 16, 16, 16, 7, 7, 1, 1, 16, 16, 16, 7, 7, 1, 1, 16, 16, 16]
>

クラス番号 1, 7, 16 がそれぞれの領域について取得できました

Nx.reduce_max で各領域のトップのクラススコアを取得します

top_class_score_tensor = Nx.reduce_max(class_score_tensor, axes: [1])

実行結果

#Nx.Tensor<
  f32[25]
  EXLA.Backend<host:0, 0.817007110.871235591.8072>
  [0.561057448387146, 0.8754310607910156, 0.8375682234764099, 0.89230877161026, 0.7436654567718506, 0.8773096799850464, 0.9699414372444153, 0.965307354927063, 0.9505174160003662, 0.9617711305618286, 0.9402678608894348, 0.8589267134666443, 0.8795986175537109, 0.9736899733543396, 0.9705929756164551, 0.9472646117210388, 0.9614760875701904, 0.9370947480201721, 0.6508569121360779, 0.8252699971199036, 0.973572313785553, 0.9683232307434082, 0.9321308732032776, 0.9564462900161743, 0.9236267805099487]
>

座標スコアとクラススコアを掛けて、最終的なスコアを計算します

score_tensor = Nx.multiply(bbox_score_tensor, top_class_score_tensor)

実行結果

#Nx.Tensor<
  f32[25]
  EXLA.Backend<host:0, 0.817007110.871235591.8072>
  [0.561057448387146, 0.8754310607910156, 0.8375682234764099, 0.89230877161026, 0.7436654567718506, 0.8773096799850464, 0.9699414372444153, 0.965307354927063, 0.9505174160003662, 0.9617711305618286, 0.9402678608894348, 0.8589267134666443, 0.8795986175537109, 0.9736899733543396, 0.9705929756164551, 0.9472646117210388, 0.9614760875701904, 0.9370947480201721, 0.6508569121360779, 0.8252699971199036, 0.973572313785553, 0.9683232307434082, 0.9321308732032776, 0.9564462900161743, 0.9236267805099487]
>

座標情報の変換

座標情報は中心座標、幅、高さという形式になっているので、これを左上座標、右下座標の形式に変換します

座標情報は 0 から 4 番目に入っているので [0..3] で取得できます

coordinate_tensor = greater_predictions_tensor[[0..-1//1, 0..3]]

実行結果

#Nx.Tensor<
  f32[25][4]
  EXLA.Backend<host:0, 0.817007110.871235591.8074>
  [
    [482.181884765625, 136.43838500976562, 187.73941040039062, 106.85380554199219],
    [482.4459533691406, 136.4070587158203, 188.18247985839844, 106.77376556396484],
    [482.1437683105469, 136.38436889648438, 187.77462768554688, 107.01679992675781],
    [482.49932861328125, 136.36087036132812, 188.29879760742188, 107.01581573486328],
    [482.37872314453125, 136.31161499023438, 190.0236358642578, 106.85448455810547],
    [482.3506164550781, 136.28390502929688, 189.64100646972656, 107.11427307128906],
    [288.2965087890625, 308.01239013671875, 367.0347595214844, 317.332763671875],
    [288.33868408203125, 307.7346496582031, 367.5043029785156, 318.0263366699219],
    [184.40452575683594, 423.7576904296875, 149.1966094970703, 353.48760986328125],
    [184.554443359375, 423.8431701660156, 149.1243896484375, 353.4234619140625],
    [184.81690979003906, 423.9307556152344, 149.1788330078125, 353.64935302734375],
    [482.3323059082031, 136.2632598876953, 190.04119873046875, 106.98236083984375],
    [482.4167785644531, 136.32569885253906, ...],
    ...
  ]
>

中心座標に幅、高さの半分を足し引きすれば左上、右下座標になります

bbox_half_width = Nx.divide(coordinate_tensor[[0..-1//1, 2]], 2)
bbox_half_height = Nx.divide(coordinate_tensor[[0..-1//1, 3]], 2)

また、検出結果は 640 * 640 にリサイズしたときの座標になっているため、元の大きさに合わせるように元画像の幅、高さに対する比を掛けます

min_x_tensor =
  coordinate_tensor[[0..-1//1, 0]]
  |> Nx.subtract(bbox_half_width)
  |> Nx.multiply(img_width / 640)

min_y_tensor =
  coordinate_tensor[[0..-1//1, 1]]
  |> Nx.subtract(bbox_half_height)
  |> Nx.multiply(img_height / 640)

max_x_tensor =
  coordinate_tensor[[0..-1//1, 0]]
  |> Nx.add(bbox_half_width)
  |> Nx.multiply(img_width / 640)

max_y_tensor =
  coordinate_tensor[[0..-1//1, 1]]
  |> Nx.add(bbox_half_height)
  |> Nx.multiply(img_height / 640)

formed_coordinate_tensor =
  [min_x_tensor, min_y_tensor, max_x_tensor, max_y_tensor]
  |> Nx.stack()
  |> Nx.transpose()

実行結果

#Nx.Tensor<
  f32[25][4]
  EXLA.Backend<host:0, 0.817007110.871235591.8113>
  [
    [465.9746398925781, 74.71033477783203, 691.2619018554688, 170.87876892089844],
    [466.0256652832031, 74.71815490722656, 691.8446044921875, 170.81454467773438],
    [465.90777587890625, 74.58837127685547, 691.2373046875, 170.9034881591797],
    [466.0199279785156, 74.56766510009766, 691.978515625, 170.88189697265625],
    [464.8403015136719, 74.59593200683594, 692.8687133789062, 170.7649688720703],
    [465.0361633300781, 74.45408630371094, 692.6054077148438, 170.85693359375],
    [125.7349624633789, 134.41140747070312, 566.1766967773438, 420.0108947753906],
    [125.50384521484375, 133.84933471679688, 566.509033203125, 420.072998046875],
    [131.76747131347656, 222.31248474121094, 310.80340576171875, 540.4512939453125],
    [131.99070739746094, 222.4182891845703, 310.9399719238281, 540.4994506835938],
    [132.2729949951172, 222.3954620361328, 311.28759765625, 540.6798706054688],
    [464.7740783691406, 74.494873046875, 692.823486328125, 170.77899169921875],
    [465.0534973144531, 74.34207916259766, ...],
    ...
  ]
>

Non-Maximun Suppression

抽出した 25 個の領域は複数ほぼ同じ領域に重なり合っています

重複した領域は邪魔なので、 NMS = Non-Maximun Suppression を実行します

Evision.DNN.nmsBoxes にスコア、領域を渡すことで、重なった領域のうち、それぞれ最大スコアの領域だけを抽出できます

score_list = Nx.to_list(score_tensor)

nms_threshold = 0.7

selected_index_tensor =
  formed_coordinate_tensor
  |> Evision.DNN.nmsBoxes(score_list, score_threshold, nms_threshold)
  |> Nx.tensor()

実行結果

#Nx.Tensor<
  s64[3]
  EXLA.Backend<host:0, 0.817007110.871235601.19459>
  [13, 9, 3]
>

13 番目, 9 番目, 3 番目の 3 つまで領域を絞り込めました

この3つの領域について、座標情報を取り出します

selected_bboxes = Nx.take(formed_coordinate_tensor, selected_index_tensor)

実行結果

#Nx.Tensor<
  f32[3][4]
  EXLA.Backend<host:0, 0.817007110.871235591.8118>
  [
    [126.04136657714844, 134.79254150390625, 565.781494140625, 419.6330261230469],
    [131.99070739746094, 222.4182891845703, 310.9399719238281, 540.4994506835938],
    [466.0199279785156, 74.56766510009766, 691.978515625, 170.88189697265625]
  ]
>

同様にクラス番号を取得します

ただし、この後他のテンソルと結合するため、次元を増やして 3 * 1 にしておきます

selected_classes = Nx.take(class_index_tensor, selected_index_tensor) |> Nx.new_axis(1)

実行結果

#Nx.Tensor<
  s64[3][1]
  EXLA.Backend<host:0, 0.817007110.871235591.8251>
  [
    [1],
    [16],
    [7]
  ]
>

スコアも取得します

selected_scores = Nx.take(score_tensor, selected_index_tensor) |> Nx.new_axis(1)

実行結果

#Nx.Tensor<
  f32[3][1]
  EXLA.Backend<host:0, 0.817007110.871235591.8253>
  [
    [0.9736899733543396],
    [0.9617711305618286],
    [0.89230877161026]
  ]
>

座標情報、クラス番号、スコアを一つのテンソルとして結合します

formed_tensor = Nx.concatenate([selected_bboxes, selected_classes, selected_scores], axis: 1)

実行結果

#Nx.Tensor<
  f32[3][6]
  EXLA.Backend<host:0, 0.817007110.871235591.8128>
  [
    [126.04136657714844, 134.79254150390625, 565.781494140625, 419.6330261230469, 1.0, 0.9736899733543396],
    [131.99070739746094, 222.4182891845703, 310.9399719238281, 540.4994506835938, 16.0, 0.9617711305618286],
    [466.0199279785156, 74.56766510009766, 691.978515625, 170.88189697265625, 7.0, 0.89230877161026]
  ]
>

これで推論結果が 3 * 6 のテンソルになりました

推論結果の描画

推論結果のテンソル

formed_tensor
|> Nx.to_list()
|> Enum.reduce(img, fn prediction, drawed_mat ->
  # 座標情報、クラス番号は整数に変換する
  [left, top, right, bottom, class_index] =
    prediction
    |> Enum.slice(0..4)
    |> Enum.map(&trunc(&1))

  # スコアは小数点以下3桁の文字列に変換する
  score =
    prediction
    |> Enum.at(5)
    |> Float.round(3)
    |> Float.to_string()

  # class の値に対応するラベルを取得する
  label = Enum.at(labels, class_index)

  drawed_mat
  # 四角形を描画する
  |> Evision.rectangle(
    {left, top},
    {right, bottom},
    {255, 0, 0},
    thickness: 10
  )
  # ラベル文字を書く
  |> Evision.putText(
    label <> ":" <> score,
    {left + 6, top + 26},
    Evision.Constant.cv_FONT_HERSHEY_SIMPLEX(),
    1,
    {0, 0, 255},
    thickness: 2
  )
end)

自動車がトラックになっているのは微妙ですが、物体の検出、クラスの分類ができていることが確認できました

YOLOv7 のモジュール化

ここまでの処理をモジュールにまとめます

defn を使いたいので Nx.Defn をインポートします

import Nx.Defn

defn を使うと、行列間の四則演算（Nx.subtruct や Nx.divide）を演算子（- や /）で表すことができます

ただし、 defn 内では Nx.to_list や Nx.to_number は使えないので、そのような処理については def にしています

defmodule YOLOv7 do
  def detect(img, net, out_names, score_threshold, nms_threshold) do
    blob = Evision.DNN.blobFromImage(img, size: {640, 640}, swapRB: true, crop: false)

    predictions =
      net
      |> Evision.DNN.Net.setInput(
        blob,
        name: "",
        scalefactor: 1 / 255,
        mean: {0, 0, 0}
      )
      |> Evision.DNN.Net.forward(outBlobNames: out_names)
      |> Enum.at(0)
      |> Evision.Mat.to_nx(EXLA.Backend)

    selected_tensor = filter_predictions(predictions, score_threshold)

    {img_height, img_width, _} = Evision.Mat.shape(img)
    coordinate_tensor = selected_tensor[[0..-1//1, 0..3]]
    formed_coordinate_tensor = format_coordinates(coordinate_tensor, img_width, img_height)

    {class_index_tensor, score_tensor} = get_class_and_score(selected_tensor)

    nms(
      formed_coordinate_tensor,
      class_index_tensor,
      score_tensor,
      score_threshold,
      nms_threshold
    )
  end

  def filter_predictions(predictions, score_threshold) do
    greater_tensor = Nx.greater(predictions[[0, 0..-1//1, 4]], score_threshold)

    greater_count =
      greater_tensor
      |> Nx.sum()
      |> Nx.to_number()

    greater_indices = Nx.argsort(greater_tensor, direction: :desc)[[0..(greater_count - 1)]]

    Nx.take(predictions[0], greater_indices, axis: 0)
  end

  defn format_coordinates(coordinate_tensor, width, height) do
    bbox_half_width = coordinate_tensor[[0..-1//1, 2]] / 2
    bbox_half_height = coordinate_tensor[[0..-1//1, 3]] / 2
    width_ratio = width / 640
    height_ratio = height / 640

    min_x_tensor = (coordinate_tensor[[0..-1//1, 0]] - bbox_half_width) * width_ratio
    min_y_tensor = (coordinate_tensor[[0..-1//1, 1]] - bbox_half_height) * height_ratio
    max_x_tensor = (coordinate_tensor[[0..-1//1, 0]] + bbox_half_width) * width_ratio
    max_y_tensor = (coordinate_tensor[[0..-1//1, 1]] + bbox_half_height) * height_ratio

    [min_x_tensor, min_y_tensor, max_x_tensor, max_y_tensor]
    |> Nx.stack()
    |> Nx.transpose()
  end

  defn get_class_and_score(selected_tensor) do
    bbox_score_tensor = selected_tensor[[0..-1//1, 4]]
    class_score_tensor = selected_tensor[[0..-1//1, 5..-1//1]]
    class_index_tensor = Nx.argmax(class_score_tensor, axis: 1)
    top_class_score_tensor = Nx.reduce_max(class_score_tensor, axes: [1])
    score_tensor = bbox_score_tensor * top_class_score_tensor

    {class_index_tensor, score_tensor}
  end

  def nms(
        formed_coordinate_tensor,
        class_index_tensor,
        score_tensor,
        score_threshold,
        nms_threshold
      ) do
    score_list = score_tensor |> Nx.to_list()

    selected_index_tensor =
      formed_coordinate_tensor
      |> Evision.DNN.nmsBoxes(score_list, score_threshold, nms_threshold)
      |> Nx.tensor()

    selected_bboxes = Nx.take(formed_coordinate_tensor, selected_index_tensor)
    selected_classes = Nx.take(class_index_tensor, selected_index_tensor) |> Nx.new_axis(1)
    selected_scores = Nx.take(score_tensor, selected_index_tensor) |> Nx.new_axis(1)

    Nx.concatenate([selected_bboxes, selected_classes, selected_scores], axis: 1)
  end

  def draw_bbox(img, bbox_tensor, labels) do
    bbox_tensor
    |> Nx.to_list()
    |> Enum.reduce(img, fn prediction, drawed_mat ->
      [left, top, right, bottom, class_index] =
        prediction
        |> Enum.slice(0..4)
        |> Enum.map(&trunc(&1))

      score =
        prediction
        |> Enum.at(5)
        |> Float.round(3)
        |> Float.to_string()

      # class の値に対応するラベルを取得する
      label = Enum.at(labels, class_index)

      drawed_mat
      # 四角形を描画する
      |> Evision.rectangle(
        {left, top},
        {right, bottom},
        {255, 0, 0},
        thickness: 10
      )
      # ラベル文字を書く
      |> Evision.putText(
        label <> ":" <> score,
        {left + 6, top + 26},
        Evision.Constant.cv_FONT_HERSHEY_SIMPLEX(),
        1,
        {0, 0, 255},
        thickness: 2
      )
    end)
  end
end

改めて、モジュールを使って物体検出してみます

bbox_tensor = YOLOv7.detect(img, net, out_names, 0.6, 0.7)

実行結果

#Nx.Tensor<
  f32[3][6]
  EXLA.Backend<host:0, 0.817007110.871235591.8327>
  [
    [126.04136657714844, 134.79254150390625, 565.781494140625, 419.6330261230469, 1.0, 0.9736899733543396],
    [131.99070739746094, 222.4182891845703, 310.9399719238281, 540.4994506835938, 16.0, 0.9617711305618286],
    [466.0199279785156, 74.56766510009766, 691.978515625, 170.88189697265625, 7.0, 0.89230877161026]
  ]
>

YOLOv7.draw_bbox(img, bbox_tensor, labels)

モジュールでも問題なく物体検出できました

おまけ

一応、 Google Colab で GPU を使えるようにして実行してみましたが、速度は向上しませんでした

その際のセットアップは以下のようにしました

```elixir
Mix.install(
  [
    {:nx, "~> 0.5"},
    {:exla, "~> 0.5.1"},
    {:evision, "~> 0.1.31"},
    {:req, "~> 0.3"},
    {:kino, "~> 0.9"}
  ],
  system_env: [
    {"XLA_TARGET", "cuda111"},
    {"CMAKE_OPENCV_OPTIONS",
     "-D WITH_CUDA=ON -D CUDA_FAST_MATH=ON -D WITH_CUBLAS=ON -D WITH_CUDNN=ON -D WITH_NVCUVID=OFF -D OPENCV_DNN_CUDA=ON"},
    {"EVISION_PREFER_PRECOMPILED", "false"},
    {"EVISION_ENABLE_CONTRIB", "true"},
    {"EVISION_ENABLE_CUDA", "true"},
    {"EVISION_CUDA_VERSION", "111"}
  ],
  config: [nx: [default_backend: EXLA.Backend]]
)

OpenCV のビルドに数時間かかるので、実行する場合は覚悟してください

（結果、意味はありませんでしたが、、、）

まとめ

YOLOv3 と同じ形式のテンソルで推論結果が出たため、問題なく整形、描画ができました

以前は Nx.to_batched や Enum.map を使っていましたが、ほぼ全ての整形を行列演算でできるようになった（自分の Nx の理解が進んだ）のは良かったです