Elixir の evision で YOLOv3 の物体検出を実行する #AI

はじめに

以前の記事で、 Elixir の evision を使った画像処理を紹介しました

evision は OpenCV を Elixir で使えるようにしたものです

OpenCV ということは、、、

DNN 機能を使えば機械学習のモデルを動かせるわけです

本当は Axon で動かしたいのですが、まだ YOLO を動かすに至っておらず、、、

evision なら Python と同じ感覚なので、簡単に実装できました

evision の examples でも、他のモデルを使った例があります

実装したコードはこちら

2022/12/13 更新

最新のモジュールを使用する

2023/5/27

YOLOv7 版を公開しました

https://qiita.com/RyoWakabayashi/items/742bb7c172394498875e

推論結果の整形も行列のままできるように改良しました

準備

必要なパッケージをインストールします

httpoison: ファイルダウンロード
evision: OpenCV
Kino: 画像表示
Nx: 行列計算

Mix.install([
  {:httpoison, "~> 1.8"},
  {:evision, "~> 0.1"},
  {:kino, "~> 0.7"},
  {:nx, "~> 0.4"}
])

機械学習モデルのロード

今回は Darknet の YOLOv3 を使います

公式で配布している重み、設定ファイル、COCOのラベル一覧をダウンロードします

weights_path = "/data/yolov3.weights"
cfg_path = "/data/yolov3.cfg"
label_path = "/data/label.txt"

"https://pjreddie.com/media/files/yolov3.weights"
|> HTTPoison.get!(recv_timeout: 300_000)
|> then(&File.write(weights_path, &1.body))

"https://raw.githubusercontent.com/pjreddie/darknet/master/cfg/yolov3.cfg"
|> HTTPoison.get!()
|> then(&File.write(cfg_path, &1.body))

"https://raw.githubusercontent.com/pjreddie/darknet/master/data/coco.names"
|> HTTPoison.get!()
|> then(&File.write(label_path, &1.body))

DNN の機能は Evision.DNN に実装されています

Python でモデルをロードするときと同じように、
readNetModel で重みファイルと設定ファイルを読み込みます

net = Evision.DNN.readNet(weights_path, config: cfg_path, framework: "")

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

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

ラベル一覧ファイルを読み込んでラベル一覧をリストに格納しておきます

label_list =
  label_path
  |> File.stream!() 
  |> Enum.map(&String.trim/1)

画像のロード

Darknet の GitHub ページにあるサンプル画像をダウンロードします

img_path = "/data/dog.jpg"

"https://raw.githubusercontent.com/pjreddie/darknet/master/data/dog.jpg"
|> HTTPoison.get!()
|> then(&File.write(img_path, &1.body))

画像を読み込んで表示します

img = Evision.imread(img_path)

画像サイズを取得しておきます

{height, width, _} = Evision.Mat.shape(img)

推論の実行

画像を Blob に変換し、モデルに入力して推論を実行します

blob = Evision.DNN.blobFromImage(img, size: {608, 608}, 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)

指定した3つの出力層の出力が取得できました

Python だと以下のように書いていましたが、変えないといけないところがありました

blob = cv2.dnn.blobFromImage(img, 1 / 255, (608, 608), 0, True, crop=False)

net.setInput(blob)

predictions = net.forward(out_names)

setInput に name が必須
scalefactor は setInput で指定
mean は setInput で指定

検出結果の抽出、整形

検出結果を使えそうな形に整形します

score_threshold = 0.8

formed_predictions =
  predictions
  # テンソルに変換
  |> Enum.map(fn prediction ->
    Evision.Mat.to_nx(prediction, Nx.BinaryBackend)
  end)
  # くっつける
  |> Nx.concatenate()
  # 配列にする
  |> Nx.to_batched(1)
  # [4] にスコアが入っているので、閾値以下のものを除外する
  |> Enum.filter(fn t ->
    t[0][4]
    |> Nx.to_number()
    |> Kernel.>(score_threshold)
  end)
  |> Enum.map(fn t ->
    # [5] 以降に各クラスに対するスコアが入っているため、トップのものを取得する
    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)

    # [0] から [3] に座標情報が入っている
    # 中央+サイズから、左上右下の値に変換する
    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)

NMS

YOLO には付き物の NMS を実行する

ざっくり説明すると、このままだと同じものが重複して検出されているので、
座標が重なっているものはスコアが高いものだけを残す

OpenCV には NMS が実装されているので、それを利用する

box_list = Enum.map(formed_predictions, & &1.box)
score_list = Enum.map(formed_predictions, & &1.score)

nms_threshold = 0.7

index_list = Evision.DNN.nmsBoxes(box_list, score_list, score_threshold, nms_threshold)

生き残った検出結果のインデックスが返ってくるので、
当該インデックスのものだけを抽出する

selected_predictions = Enum.map(index_list, &Enum.at(formed_predictions, &1))

検出結果の可視化

検出結果を元画像に描き込みます

selected_predictions
|> Enum.reduce(img, fn prediction, drawed_mat ->
  # 座標に画像サイズを掛ける
  box = Tuple.to_list(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()

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

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

おわりに

Python と同じやり方で Elixir でも YOLO が実行できました

これでアレやコレやできそうですね