ONNXに変換したYOLOv5のバウンディングボックスの処理をPythonで記述したい

Question

前提・実現したいこと

YOLOV5のカスタム学習モデルを使ってTensorFlow.jsで物体検出をやってみた
を参考にして，ONNXに変換したYOLOv5を用いて推論させた後，バウンディングボックスを表示させたいです．
推論まではできましたが，バウンディングボックスの処理が分かりません．
このリンクではJavaScriptで書かれていますが，Pythonで記述したいです．

PyTorchは使えないので，YOLOv5の元のコードをコピペすることはできません．

該当のソースコード

var output = model.executeAsync(imageTensor).then(output=> {
                    const o0 = output[0].arraySync();###ここがわからない

                    const OBJECT_TH = 0.5;
                    const IOU_TH = 0.5;
                    const bairitu_w = 1008/640;
                    const bairitu_h = 756/640;

                    let a=0;
                    var list = new Array();

                    for (let i = 0; i < o0[0].length; i++) {
                        if((o0[0][i][4]*o0[0][i][5])>OBJECT_TH){
                            a = a+1;
                            const dx = o0[0][i][2]*bairitu_w/2;
                            const dy = o0[0][i][2]*bairitu_h/2;

                            var ary = new Array();
                            ary.push(o0[0][i][0]*bairitu_w - dx);
                            ary.push(o0[0][i][1]*bairitu_h - dy);
                            ary.push(o0[0][i][2]*bairitu_w);
                            ary.push(o0[0][i][3]*bairitu_h);
                            ary.push(o0[0][i][4]*o0[0][i][5]);

                            list.push(ary);
                        }
                    }

                    list.sort(function(a,b){return(a[4] - b[4]);});

                    for (let i = 0; i < list.length; i++) {
                        let aa = 0;
                        for(let j = i+1; j< list.length; j++) {
                            var iou = bbox_iou(list[i], list[j]);
                            if(iou > IOU_TH) {
                                aa = 1;
                                break;
                            }
                        }
                   }
             

試したこと(一部省略)

import cv2
import onnx
import onnxruntime
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image, ImageDraw

# 推論を実行する
for i in range(3):
    output_name = session.get_outputs()[i+1].name

    output[i] = session.run([output_name], {input_name: image_data})[0]

###var output = model.executeAsync(imageTensor).then(output=> {
                    #const o0 = output[0].arraySync();###Pythonに書き換えたい

o0 = [] #

OBJECT_TH = 0.5 #閾値
IOU_TH = 0.5
W_RATIO = 1000/640
H_RATIO = 756/640

a = 0
var_list = []

for i in range(len(o0)):
    if o0[0][i][4] * o0[0][i][5] > OBJECT_TH:
        a = a + 1
        DX = o0[0][i][2] * W_RATIO / 2
        DY = o0[0][i][2] * H_RATIO / 2

        ary = []
        ary.append(o0[0][i][0] * W_RATIO - DX)
        ary.append(o0[0][i][1] * H_RATIO - DY)
        ary.append(o0[0][i][2] * W_RATIO)
        ary.append(o0[0][i][3] * H_RATIO)
        ary.append(o0[0][i][4] * o0[0][i][5])

        var_list.append(ary)

var_list.sort(key=lambda x:x[4],reverse=True)

for i in range(len(list)):
    aa = 0
    for j in range(i+1, len(list)):
        iou = bbox_iou(var_list[i], var_list[j])
        if iou > IOU_TH:
            aa = 1
            break

追記

def main():
    
    img = get_image()

    # リサイズ
    image_data = preprocess(img)
    image_size = np.array([img.size[1], img.size[0]], dtype=np.int32).reshape(1, 2)

    # 学習済みモデルからセッションを作る
    session = onnxruntime.InferenceSession('yolov5s.onnx')

    # #元の画像の表示
    # plt.imshow(img)
    # plt.show()

    #input output 名前取得
    print("\n")
    print("input:")
    for session_input in session.get_inputs():
        input_name = session_input.name
        input_name_img_shape = session_input.shape
        print(session_input.name, session_input.shape, session_input.type)
        

    print("\n")
    # 出力のラベル名の確認
    print("output:")
    for session_output in session.get_outputs():
        print(session_output.name, session_output.shape, session_output.type)

を実行したところ，

input:
images [1, 3, 640, 640] tensor(float)


output:
output [1, 25200, 85] tensor(float)
351 [1, 3, 80, 80, 85] tensor(float)
420 [1, 3, 40, 40, 85] tensor(float)
489 [1, 3, 20, 20, 85] tensor(float)

と出ました．

ここから

outputs_index = session.run([output_name_boxes, output_name_scores, output_name_indices],
      {input_name: image_data, input_name_img_shape: image_size})

に落とし込みたいのですが，どうすればよいか分かりません．

Comments

[jbpb0]

このあたりはPythonでやってますけど、参考になりませんか？
https://qiita.com/enya314/items/1bd053d6a81a156ff814
https://www.alpha.co.jp/blog/202108_02

[PJwnOI]

私は学習や推論ではなく，推論結果からバウンディングボックスの座標計算と画像に出力させたいのです．
貼っていただいたリンクはとても参考になりづらいです，すみません．

[jbpb0]

> 推論結果からバウンディングボックスの座標計算と画像に出力させたい

https://github.com/ultralytics/yolov5/blob/master/detect.py
をながめると、上記の処理は
https://github.com/ultralytics/yolov5/blob/master/utils/plots.py
の「class Annotator:」の中で定義されてるのを使ってごにょごにょしてるので、そのあたりが参考になるかな、と思って紹介したのですが、求めてるものはそういうのではないようですね
失礼しました

[PJwnOI]

そのリンクについては私も閲覧しましたが，ONNXで得られたデータをどうするのかがわからず，断念しました．
また，PyTorchを使っている部分があるので難しいのです...

[jbpb0]

https://cpp-learning.com/onnx-runtime_yolo/#ONNXRuntime_YoloV3py
のコードは、ONNX Runtimeでの推論結果を、画像に重ね書きしてます
PyTorchは使ってません

yolo v3なので、v5とは細部で違いがあるかもしれませんが、下記に詳しく解説されてるので、それを参考にしながらv5に対応させるのは、それほど難しくはないように思います
https://cpp-learning.com/wp-content/uploads/2019/07/ONNX-Runtime%E3%81%A8YoloV3%E3%81%A7%E7%89%A9%E4%BD%93%E6%A4%9C%E5%87%BA.html

[PJwnOI]

ありがとうございます．ONNXで調べていましたが，onnxruntimeで調べるとよかったんですね．
後でやってみます．

Answer 1

YOLOv5のページを見て解決していただきました．