使用 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 模型
  • <测试图像>.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 扩展仓库