在移动设备上使用 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
:带有边界框的测试图像。
例如,扩展仓库中的狼测试图像
构建 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 推理代码作为开发移动应用程序的基础。