使用 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
:带有提供的边界框的测试图像。
例如,扩展仓库中的狼测试图像
构建 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
}
}
}
使用 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
结果输出绘制在原始图像上!
开发您的移动应用
您可以使用 Python 推理代码作为开发移动应用程序的基础。