LoginSignup
5

More than 1 year has passed since last update.

Elixir の AxonOnnx で YOLOv2 による物体検出を実行する

Last updated at Posted at 2022-08-04

はじめに

以前の記事で、 Elixir で YOLOv3 を実行しました

ただし、こちらは OpenCV のラッパーである evision を利用しています

今回は @piacerex さんの紹介してくださった AxonOnnx を使って、 YOLOv2 を動かしたいと思います
※ YOLOv3 はまだ AxonOnnx で読み込めないらしい

@piacerex さんの記事

実装したコード(Livebook)はこちら

セットアップ

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

ここは @piacerex さんのものをコピーしただけです

Mix.install([
  {:exla, "~> 0.2"},
  {:axon_onnx, "~> 0.1"},
  {:stb_image, "~> 0.5"},
  {:download, "~> 0.0"},
  {:jason, "~> 1.3"},
  {:kino, "~> 0.6"},
])

EXLA を Nx バックエンドとして使うよう指定します
※こうしないとすごく遅いです

EXLA.set_as_nx_default([:tpu, :cuda, :rocm, :host])

ヘルパー

ファイルダウンロードと画像表示用のヘルパーです

defmodule Helper do
  def download!(url, save_as) do
    unless File.exists?(save_as) do
      Download.from(url, path: save_as)
    end

    save_as
  end

  def show_image(image_path) do
    image_path
    |> File.read!()
    |> Kino.Image.new(:jpeg)
  end
end

モデルのロード

今回使うモデルは COCO データセットを学習しているので、ラベルの一覧をダウンロードしてきて読み込みます

classes_path = "/data/coco.names"

classes =
  "https://raw.githubusercontent.com/pjreddie/darknet/master/data/coco.names"
  |> Helper.download!(classes_path)
  |> File.stream!()
  |> Enum.map(&String.trim/1)

スクリーンショット 2022-08-04 11.32.52.png

クラス数(分類できる物体の種類)を取得しておきます

COCO データセットの場合は 80 種類です

num_classes = Enum.count(classes)

スクリーンショット 2022-08-04 11.34.08.png

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

入力が 1 * 3 * 416 * 416 なことが分かりますね

model_path = "/data/yolov2.onnx"

{model, params} =
  "https://media.githubusercontent.com/media/onnx/models/main/vision/object_detection_segmentation/yolov2-coco/model/yolov2-coco-9.onnx"
  |> Helper.download!(model_path)
  |> AxonOnnx.import()

スクリーンショット 2022-08-04 11.35.18.png

アンカーの指定

YOLO で使うアンカーボックスの値を指定しておきます

YOLOv2 の COCO データセットの場合、以下の値になります

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

Darknet のこちらに定義している値ですね

anchors =  0.57273, 0.677385, 1.87446, 2.06253, 3.33843, 5.47434, 7.88282, 3.52778, 9.77052, 9.16828

アンカーボックスの数も必要になるので取得しておきます

見ての通り、アンカーボックスは 5 種類です

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

スクリーンショット 2022-08-04 11.41.48.png

後々の計算のために、形を変えておきます

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

スクリーンショット 2022-08-04 11.43.09.png

画像の読み込み

YOLOv3 のときにも使った犬の画像を読み込みます

入力層の形に合わせて 416 * 416 にリサイズしておきます

image_path = "/data/dog.jpg"

image_tensor =
  "https://raw.githubusercontent.com/pjreddie/darknet/master/data/dog.jpg"
  |> Helper.download!(image_path)
  |> File.read!()
  |> StbImage.read_binary!()
  |> StbImage.resize(416, 416)
  |> StbImage.to_nx()

Helper.show_image(image_path)

スクリーンショット 2022-08-04 11.47.30.png

RGBA だったときの回避をしたあと、 255 で割ります

入力層の形に合わせるため、以下のように転置します

高さ、幅、色 -> 色、高さ、幅

nx_channels = Nx.axis_size(image_tensor, 2)

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

推論実行

推論を実行します

出力は 1 * 425 * 13 * 13 で得られました

feats = Axon.predict(model, params, image_tensor)

スクリーンショット 2022-08-04 11.52.33.png

これがそのまま使える出力だったらいいのですが、
YOLOv2 の場合はここから色々しないと使える形にはなりません

13 はアンカーボックスの配置されている数です
縦横に 13 * 13 のアンカーボックスが配置されています

x 0 1 2 3 4 5 6 7 8 9 10 11 12
0
1
2
3
4
5
6
7
8
9
10
11
12

425 というのは

アンカーの種類 5 * (Y座標・X座標・高さ・幅・検出スコアの5つ + クラス毎の分類スコア 80)
= 5 * 85 = 425

のことです

アンカーボックスの 13 * 13 がアンカーボックスの種類毎にあって、
それぞれに座標情報と分類結果の85項目を持っている、ということです(分かりにくい)

結果の整形

前述の推論結果をうまく変形していくと、どの物体がどこにあるかを示す情報になるのですが、
何の頼りもなしに実装できないので、 Python 実装から翻訳することにしました

このリポジトリーが5年以上も前のものだったので、
検証のためにこれを最近の TensorFlow で動くようにするのもちょっと苦労しました
(これは別途リポジトリーと記事を書きます)

整形の処理は yad2k/models/keras_yolo.py の中の yolo_head と yolo_eval に書かれています

翻訳元

    num_anchors = len(anchors)
...
    anchors_tensor = K.reshape(K.variable(anchors), [1, 1, 1, num_anchors, 2])

まずアンカーの種類数とアンカーテンソルを取得しています

これは既にやっていますね

次に、畳み込みのためにアンカーボックスの縦横の数を取得しています

翻訳元

    conv_dims = K.shape(feats)[1:3]  # assuming channels last

ここで注目するのが [1:3] です

Axon の出力は 1 * 425 * 13 * 13 だったので、これの [1:3] をとると 425, 13 になります
しかし、実際にとりたいのはアンカーボックスの縦横の数なので、 13, 13 です

実行してみればわかるのですが、
実はこのPython実装では、出力の形が 1 * 13 * 13 * 425 になっており、
Axon のときと軸が入れ替わっているのです

最初はそれに気づかず、おかしな結果になっていました :sweat:

なので、 Elixir ではこうなります

翻訳後

{_, _, num_y_blocks, num_x_blocks} = Nx.shape(feats)

そして、後の実装を合わせるために転置しておきます

翻訳後

feats =
  feats
  |> Nx.transpose(axes: [0, 2, 3, 1])

次にアンカーボックスの座標情報を補正するためのテンソルを作ります

翻訳元

    # In YOLO the height index is the inner most iteration.
    conv_height_index = K.arange(0, stop=conv_dims[0])
    conv_width_index = K.arange(0, stop=conv_dims[1])
    conv_height_index = K.tile(conv_height_index, [conv_dims[1]])

    # TODO: Repeat_elements and tf.split doesn't support dynamic splits.
    # conv_width_index = K.repeat_elements(conv_width_index, conv_dims[1], axis=0)
    conv_width_index = K.tile(K.expand_dims(conv_width_index, 0), [conv_dims[0], 1])
    conv_width_index = K.flatten(K.transpose(conv_width_index))
    conv_index = K.transpose(K.stack([conv_height_index, conv_width_index]))
    conv_index = K.reshape(conv_index, [1, conv_dims[0], conv_dims[1], 1, 2])
    conv_index = K.cast(conv_index, K.dtype(feats))

翻訳後

conv_height_index =
  Nx.iota({num_y_blocks})
  |> Nx.tile([num_x_blocks])

conv_width_index =
  Nx.iota({num_x_blocks})
  |> Nx.reshape({1, num_x_blocks})
  |> Nx.tile([num_y_blocks, 1])
  |> Nx.transpose()
  |> Nx.flatten()

conv_index =
  Nx.stack([conv_height_index, conv_width_index])
  |> Nx.transpose()
  |> Nx.reshape({1, num_x_blocks, num_y_blocks, 1, 2})
  |> Nx.as_type({:f, 32})

実行すると、こんな感じのテンソルができます

スクリーンショット 2022-08-04 13.35.13.png

...

スクリーンショット 2022-08-04 13.35.33.png

...

そして、推論結果の方も変形させます

翻訳元

    feats = K.reshape(feats, [-1, conv_dims[0], conv_dims[1], num_anchors, num_classes + 5])

翻訳先

feats =
  feats
  |> Nx.reshape({1, num_y_blocks, num_x_blocks, num_anchors, num_classes + 5})

スクリーンショット 2022-08-04 13.47.19.png

これで 1 * 13 * 13 * 5 * 85 の形になりました

もう一つ、座標情報をブロック数で割るためのテンソルを作ります

翻訳元

    conv_dims = K.cast(K.reshape(conv_dims, [1, 1, 1, 1, 2]), K.dtype(feats))

翻訳先

conv_tensor =
  Nx.tensor([num_y_blocks, num_x_blocks])
  |> Nx.reshape({1, 1, 1, 1, 2})
  |> Nx.as_type({:f, 32})

スクリーンショット 2022-08-04 13.50.41.png

さあ、必要なものが揃ったので座標情報を補正しましょう

翻訳元

    box_xy = K.sigmoid(feats[..., :2])
    box_wh = K.exp(feats[..., 2:4])
    box_confidence = K.sigmoid(feats[..., 4:5])
    box_class_probs = K.softmax(feats[..., 5:])

    # Adjust preditions to each spatial grid point and anchor size.
    # Note: YOLO iterates over height index before width index.
    box_xy = (box_xy + conv_index) / conv_dims
    box_wh = box_wh * anchors_tensor / conv_dims

翻訳先

box_xy =
  feats[[0..0, 0..(num_y_blocks - 1), 0..(num_x_blocks - 1), 0..(num_anchors - 1), 0..1]]
  |> Nx.logistic()
  |> Nx.add(conv_index)
  |> Nx.divide(conv_tensor)

box_wh =
  feats[[0..0, 0..(num_y_blocks - 1), 0..(num_x_blocks - 1), 0..(num_anchors - 1), 2..3]]
  |> Nx.exp()
  |> Nx.multiply(anchors_tensor)
  |> Nx.divide(conv_tensor)

box_confidence =
  Nx.logistic(
    feats[[
      0..0,
      0..(num_y_blocks - 1),
      0..(num_x_blocks - 1),
      0..(num_anchors - 1),
      4..4
    ]]
  )

box_class_probs =
  feats[[
    0..0,
    0..(num_y_blocks - 1),
    0..(num_x_blocks - 1),
    0..(num_anchors - 1),
    5..(num_classes + 4)
  ]]

exp =
  box_class_probs
  |> Nx.exp()

exp_sum =
  box_class_probs
  |> Nx.exp()
  |> Nx.sum(axes: [4])
  |> Nx.reshape({1, num_y_blocks, num_x_blocks, num_anchors, 1})
  |> Nx.broadcast({1, num_y_blocks, num_x_blocks, num_anchors, num_classes})

box_class_probs = Nx.divide(exp, exp_sum)

途中、 Nx に softmax がなかったので exp と sum で計算しています

多分もっとスマートにできると思いますが、ちゃんと softmax になっているので良しとします

続いて、中心座標と高さ幅になっているのを、左上右下の座標に変換します

翻訳元

    """Convert YOLO box predictions to bounding box corners."""
    box_mins = box_xy - (box_wh / 2.)
    box_maxes = box_xy + (box_wh / 2.)

    return K.concatenate([
        box_mins[..., 1:2],  # y_min
        box_mins[..., 0:1],  # x_min
        box_maxes[..., 1:2],  # y_max
        box_maxes[..., 0:1]  # x_max
    ])

翻訳先

box_mins =
  box_xy
  |> Nx.subtract(Nx.divide(box_wh, 2.0))

box_maxes =
  box_xy
  |> Nx.add(Nx.divide(box_wh, 2))

box_list =
  Nx.concatenate([
    box_mins[[0..0, 0..(num_y_blocks - 1), 0..(num_x_blocks - 1), 0..(num_anchors - 1), 0..0]],
    box_mins[[0..0, 0..(num_y_blocks - 1), 0..(num_x_blocks - 1), 0..(num_anchors - 1), 1..1]],
    box_maxes[[0..0, 0..(num_y_blocks - 1), 0..(num_x_blocks - 1), 0..(num_anchors - 1), 0..0]],
    box_maxes[[0..0, 0..(num_y_blocks - 1), 0..(num_x_blocks - 1), 0..(num_anchors - 1), 1..1]]
  ])
  |> Nx.transpose(axes: [4, 1, 2, 3, 0])

分類結果のクラスを取得し、座標のスコア * 分類のスコアが閾値以上のデータだけを抽出します

翻訳元

    box_scores = box_confidence * box_class_probs
    box_classes = K.argmax(box_scores, axis=-1)
    box_class_scores = K.max(box_scores, axis=-1)
    prediction_mask = box_class_scores >= threshold

    # TODO: Expose tf.boolean_mask to Keras backend?
    boxes = tf.boolean_mask(boxes, prediction_mask)
    scores = tf.boolean_mask(box_class_scores, prediction_mask)
    classes = tf.boolean_mask(box_classes, prediction_mask)

翻訳先

box_scores = Nx.multiply(box_confidence, box_class_probs)

box_classes =
  box_scores
  |> Nx.argmax(axis: -1)
  |> Nx.reshape({1, num_y_blocks, num_x_blocks, num_anchors, 1})

box_class_scores =
  box_scores
  |> Nx.reduce_max(axes: [-1])
  |> Nx.reshape({1, num_y_blocks, num_x_blocks, num_anchors, 1})

score_threshold = 0.5

prediction_mask = Nx.greater(box_class_scores, score_threshold)

joined =
  Nx.concatenate(
    [
      prediction_mask,
      box_confidence,
      box_class_scores,
      box_classes,
      box_list
    ],
    axis: 4
  )
  |> Nx.reshape({num_y_blocks * num_x_blocks * num_anchors, 8})

masked_index_list =
  prediction_mask
  |> Nx.to_flat_list()
  |> Enum.with_index()
  |> Enum.filter(fn {value, _} -> value == 1 end)
  |> Enum.map(&elem(&1, 1))
  |> Nx.tensor()

selected_predictions = Nx.take(joined, masked_index_list)

formed_predictions =
  selected_predictions
  |> Nx.to_batched_list(1)
  |> Enum.map(fn t ->
    %{
      box: t[0][[4..7]] |> Nx.to_flat_list(),
      score: t[0][2] |> Nx.to_number(),
      class: t[0][3] |> Nx.to_number() |> trunc()
    }
  end)

スクリーンショット 2022-08-04 14.01.52.png

やっと使えそうな形になりました

NMS

が、このままだと様々なアンカーで検出しているため、重複(ほぼ同じ場所で重なっている推論結果)が多く存在します

そのため、 YOLO では Non-Maximum Suppression を行い、重複している場合はスコアが高いものだけを残します

前回の evision の YOLOv3 では、 evision に実装されている NMS を使いましたが、
今回は evision 抜きで実装したいため、自作しようと思います

ちなみに、先ほどまで見ていた翻訳元の Python では TensorFlow の NMS を使っていたので、
ここは参考にできません

翻訳元

    nms_index = tf.image.non_max_suppression(
        boxes, scores, max_boxes_tensor, iou_threshold=iou_threshold)

また、 Nx にも実装されていません

なので、これまた Python 実装を見つけて翻訳します

翻訳元の記事

まず、重複を見つけるために重なり具合(=IOU)を計算する関数を作ります

翻訳元

def iou_np(a, b, a_area, b_area):
    abx_mn = np.maximum(a[0], b[:,0]) # xmin
    aby_mn = np.maximum(a[1], b[:,1]) # ymin
    abx_mx = np.minimum(a[2], b[:,2]) # xmax
    aby_mx = np.minimum(a[3], b[:,3]) # ymax
    w = np.maximum(0, abx_mx - abx_mn + 1)
    h = np.maximum(0, aby_mx - aby_mn + 1)
    intersect = w*h

    iou_np = intersect / (a_area + b_area - intersect)
    return iou_np

翻訳先

iou_nx = fn (a, b, a_area, b_area) ->
  num_b = Nx.shape(b) |> elem(0)
  abx_mn = Nx.max(a[0], b[[0..(num_b - 1), 0]]) # xmin
  aby_mn = Nx.max(a[1], b[[0..(num_b - 1), 1]]) # ymin
  abx_mx = Nx.min(a[2], b[[0..(num_b - 1), 2]]) # xmax
  aby_mx = Nx.min(a[3], b[[0..(num_b - 1), 3]]) # ymax
  w = Nx.subtract(abx_mx, abx_mn)
  h = Nx.subtract(aby_mx, aby_mn)
  intersect = Nx.multiply(w, h)

  Nx.divide(intersect, Nx.subtract(Nx.add(a_area, b_area), intersect))
end

そして、 NMS を実行する関数を作ります

翻訳元

def nms_fast(bboxes, scores, classes, iou_threshold=0.5):
    areas = (bboxes[:,2] - bboxes[:,0] + 1) \
             * (bboxes[:,3] - bboxes[:,1] + 1)
    
    sort_index = np.argsort(scores)
    
    i = -1 # 未処理の矩形のindex
    while(len(sort_index) >= 2 - i):
        max_scr_ind = sort_index[i]
        ind_list = sort_index[:i]
        iou = iou_np(bboxes[max_scr_ind], bboxes[ind_list], \
                     areas[max_scr_ind], areas[ind_list])
        del_index = np.where(iou >= iou_threshold)
        sort_index = np.delete(sort_index, del_index)
        #print(len(sort_index), i, flush=True)
        i -= 1 # 未処理の矩形のindexを1減らす

    bboxes = bboxes[sort_index]
    scores = scores[sort_index]
    classes = classes[sort_index]
    
    return bboxes, scores, classes

翻訳先

nms = fn (bboxes, scores, iou_threshold) ->
  num_boxes = Nx.shape(bboxes) |> elem(0)

  areas = Nx.multiply(
    Nx.subtract(bboxes[[0..(num_boxes - 1), 2]], bboxes[[0..(num_boxes - 1), 0]]),
    Nx.subtract(bboxes[[0..(num_boxes - 1), 3]], bboxes[[0..(num_boxes - 1), 1]])
  )

  bboxes
  |> Nx.to_batched_list(1)
  |> Enum.with_index()
  |> Enum.map(fn {box, index} ->
    box[0]
    |> iou_nx.(bboxes, areas[index], areas)
    |> Nx.greater(iou_threshold)
    |> Nx.to_flat_list()
  end)
  |> Enum.uniq()
  |> Enum.map(fn mask_list ->
    duplicated_index_list =
      mask_list
      |> Enum.with_index()
      |> Enum.filter(fn {value, _} -> value == 1 end)
      |> Enum.map(&elem(&1, 1))

    max_index =
      scores
      |> Nx.take(duplicated_index_list |> Nx.tensor())
      |> Nx.argmax()
      |> Nx.to_number()

    Enum.at(duplicated_index_list, max_index)
  end)
  |> Enum.uniq()
end

Elixir には while が存在しないため、とりあえず全部の組み合わせを計算したあとユニークにしています
ここも間違いなくもっとスマートな実装があるはずです

NMS を適用させます

翻訳元

        nms_bboxes, nms_scores, nms_classes = \
            nms_fast(bboxes_np, scores_np, classes_np, \
                     iou_threshold=iou_threshold)

翻訳先

box_list =
  formed_predictions
  |> Enum.map(& &1.box)
  |> Nx.tensor()

score_list =
  formed_predictions
  |> Enum.map(& &1.score)
  |> Nx.tensor()

nms_threshold = 0.7

index_list = nms.(box_list, score_list, nms_threshold)

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

スクリーンショット 2022-08-04 14.14.07.png

推論結果が3つだけに絞られました

推論結果の視覚化

前回の evision の YOLOv3 では、 OpenCV の機能で簡単に矩形を描画できました

今回は Nx だけで画像を切り抜きたいと思います

まず、リサイズなどしていない元画像をテンソルにします

image_tensor =
  image_path
  |> File.read!()
  |> StbImage.read_binary!()
  |> StbImage.to_nx()

スクリーンショット 2022-08-04 14.16.44.png

576 * 768 * 3 のテンソルになっていますね

推論結果の座標は画像内の幅、高さを1としたときの値なので、
描画時には画像の幅、高さをかける必要があります

{height, width, _} = Nx.shape(image_tensor)

Kino.render() を入れることで、最終結果ではない画像も描画されます

croped_list =
  selected_predictions
  |> Enum.map(fn prediction ->
    classes
    |> Enum.at(prediction.class)
    |> IO.inspect()

    box = 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()

    image_tensor[[top..bottom, left..right, 0..2]]
    |> StbImage.from_nx()
    |> StbImage.to_binary(:jpg)
    |> Kino.Image.new(:jpeg)
    |> Kino.render()
  end)

スクリーンショット 2022-08-04 14.20.26.png

まとめ

Axon が読み込みさえすれば、その後は気合で何とかなりますね

これで Axon が iOS で動いてくれれば、 Elixir で AI アプリが実装できますが、、、

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
5