使用 OpenCV DNN 模块从 YOLOv8 onnx 模型获取边界框、置信度分数和类标签

Question

我正在开发一个 Android 应用程序，我已经在使用 OpenCV，转换后我从 YOLOv8 获得了一个 onnx 格式的模型。这是它的输出元数据。

• 名称 - 输出0
• 类型 - float32[1,5,8400]

到目前为止，我成功运行了模型，但最终我得到的输出我无法理解。

这是输出的打印语句

Mat [ 1* 5* 8400*CV_32FC1, isCont=true, isSubmat=true, nativeObj=0x72345b4840, dataAddr=0x723076b000 ]

class Detector(private val context: Context) {
    private var net: Net? = null

    fun detect(frame: Bitmap) {
        // preprocess image
        val mat = Mat()
        Utils.bitmapToMat(resizedBitmap, mat)
        Imgproc.cvtColor(mat, mat, Imgproc.COLOR_RGBA2RGB)
        val inputBlob = Dnn.blobFromImage(mat, 1.0/255.0, Size(640.0, 640.0), Scalar(0.0), true, false)
        net?.setInput(inputBlob)
        val outputBlob = net?.forward() ?: return
        println(outputBlob)
    }

    fun setupDetector() {
        val modelFile = File(context.cacheDir, MODEL_NAME)
        if (!modelFile.exists()) {
            try {
                val inputStream = context.assets.open(MODEL_NAME)
                val size = inputStream.available()
                val buffer = ByteArray(size)
                inputStream.read(buffer)
                inputStream.close()
                val outputStream = FileOutputStream(modelFile)
                outputStream.write(buffer)
                outputStream.close()
                net = Dnn.readNetFromONNX(modelFile.absolutePath)
            } catch (e: Exception) {
                throw RuntimeException(e)
            }
        } else {
            net = Dnn.readNetFromONNX(modelFile.absolutePath)
        }
    }

    companion object {
        private const val MODEL_NAME = "model.onnx"
        private const val TENSOR_WIDTH = 640
        private const val TENSOR_HEIGHT = 640
    }
}

获取边界框、置信度得分和类别标签的一般方法是什么？如果您有任何使用 OpenCV 的 onnx 模型解决方案，那么您也可以提供。另外这个问题不是 Android 特有的。

Answer 1

根据评论中的建议，我深入研究了 YOLOv8，这就是我想出的解决方案。

val mat = Mat()
Utils.bitmapToMat(croppedBitmap, mat)
Imgproc.cvtColor(mat, mat, Imgproc.COLOR_RGBA2RGB)

val inputBlob = Dnn.blobFromImage(
    mat,
    1.0/255.0,
    Size(TENSOR_WIDTH_DOUBLE, TENSOR_HEIGHT_DOUBLE),
    Scalar(0.0),
    false,
    false
)

net?.setInput(inputBlob)

val outputBlob = net?.forward() ?: return

val strip = outputBlob.reshape(1, outputBlob.size(1))

val transposedMat = Mat()
Core.transpose(strip, transposedMat)

val boundingBoxes = mutableListOf<BoundingBox>()
for (i in 0 until transposedMat.rows()) {
    if (transposedMat.get(i, 4)[0] > CONFIDENCE_THRESHOLD) {
        boundingBoxes.add(
            BoundingBox(
                transposedMat.get(i, 0)[0],
                transposedMat.get(i, 1)[0],
                transposedMat.get(i, 2)[0],
                transposedMat.get(i, 3)[0],
                transposedMat.get(i, 4)[0]
            )
        )
    }
}

data class BoundingBox(
    val centerX: Double,
    val centerY: Double,
    val width: Double,
    val height: Double,
    val confidence: Double
)