在移动设备上使用 YOLOv8 进行对象检测和姿态估计

了解如何在移动设备上使用内置的预处理和后处理来构建和运行 ONNX 模型,以进行对象检测和姿态估计。

目录

使用 YOLOv8 进行对象检测

您可以在 ONNX Runtime 推理示例仓库的 Android 应用程序中找到完整的源代码。

构建具有内置预处理和后处理的 ONNX 模型

此步骤是可选的,因为模型在上面应用程序文件夹中的示例仓库中可用。如果您有兴趣,以下步骤将向您展示如何自行构建模型。

创建一个 Python 环境并安装以下软件包。

pip install --upgrade onnx onnxruntime onnxruntime-extensions pillow

下载以下脚本以构建模型。

curl https://raw.githubusercontent.com/microsoft/onnxruntime-extensions/main/tutorials/yolo_e2e.py > yolo_e2e.py

运行脚本。

python yolo_e2e.py [--test_image <image to test on>]

脚本运行后,您将看到一个 PyTorch 模型和两个 ONNX 模型

  • yolov8n.pt:原始 YOLOv8 PyTorch 模型
  • yolov8n.onnx:导出的 YOLOv8 ONNX 模型
  • yolov8n.with_pre_post_processing.onnx:包含预处理和后处理的 ONNX 模型
  • <test image>.out.jpg:带有边界框的测试图像。

例如,扩展仓库中的狼测试图像

Image of three white wolves with red bounding boxes

构建 Android 应用程序

将 Android 应用程序加载到 Android Developer Studio 中。

您可以在 ObjectDetector.kt 中看到主要的推理代码。它就像将图像加载到字节数组中,并通过 ONNX Runtime 运行模型以获得带有框的原始图像一样简单。

    fun detect(inputStream: InputStream, ortEnv: OrtEnvironment, ortSession: OrtSession): Result {
        // Step 1: convert image into byte array (raw image bytes)
        val rawImageBytes = inputStream.readBytes()

        // Step 2: get the shape of the byte array and make ort tensor
        val shape = longArrayOf(rawImageBytes.size.toLong())

        val inputTensor = OnnxTensor.createTensor(
            ortEnv,
            ByteBuffer.wrap(rawImageBytes),
            shape,
            OnnxJavaType.UINT8
        )
        inputTensor.use {
            // Step 3: call ort inferenceSession run
            val output = ortSession.run(Collections.singletonMap("image", inputTensor),
                setOf("image_out","scaled_box_out_next")
            )

            // Step 4: output analysis
            output.use {
                val rawOutput = (output?.get(0)?.value) as ByteArray
                val boxOutput = (output?.get(1)?.value) as Array<FloatArray>
                val outputImageBitmap = byteArrayToBitmap(rawOutput)

                // Step 5: set output result
                var result = Result(outputImageBitmap,boxOutput)
                return result
            }
        }
    }

Image of person with bicycle

使用 YOLOv8 进行姿态估计

构建姿态估计模型

注意:本教程的这一部分使用 Python。Android 和 iOS 示例即将推出!

创建一个 Python 环境并安装以下软件包。

pip install --upgrade onnx onnxruntime onnxruntime-extensions pillow

下载以下脚本以构建模型。

curl https://raw.githubusercontent.com/microsoft/onnxruntime-extensions/main/tutorials/yolov8_pose_e2e.py > yolov8_pose_e2e.py

运行脚本。

python yolov8_pose_e2e.py 

脚本运行后,您将看到一个 PyTorch 模型和两个 ONNX 模型

  • yolov8n-pose.pt:原始 YOLOv8 PyTorch 模型
  • yolov8n-pose.onnx:导出的 YOLOv8 ONNX 模型
  • yolov8n-pose.with_pre_post_processing.onnx:包含预处理和后处理的 ONNX 模型

运行姿态估计示例

您可以使用相同的脚本来运行模型,并提供您自己的图像来检测姿势。

python yolov8_pose_e2e.py --test_image person.jpg --run_model

输出将绘制在原始图像上!

Person with pose drawn

开发您的移动应用程序

您可以使用 Python 推理代码作为开发移动应用程序的基础。

其他资源

ONNX Runtime 示例仓库 ONNX Runtime 扩展仓库