YOLOv5 是一个非常流行的实时目标检测框架，其高效性和灵活性使其在移动设备上的部署变得越来越重要。本文将详细介绍如何将 YOLOv5 部署到 Android 端，并提供具体的实战步骤。
YOLOv5 提供了多种预训练模型（如 yolov5s.pt、yolov5m.pt 等），这些模型可以在官方 GitHub 仓库中下载。选择适合移动端性能需求的模型（通常推荐使用较小的模型如 yolov5s）。
Android 端通常不直接支持 PyTorch 的 .pt 模型格式，因此需要将其转换为 ONNX 格式。以下是具体步骤：
安装依赖
确保安装了最新版本的 PyTorch 和 ONNX。
pip install torch torchvision onnx
转换代码
使用以下 Python 脚本将 .pt 模型转换为 .onnx：
import torch
from models.experimental import attempt_load
# 加载模型
model = attempt_load('yolov5s.pt', map_location=torch.device('cpu')) # 替换为你的模型路径
model.eval()
# 导出为 ONNX 格式
input_size = (1, 3, 640, 640) # 输入尺寸
dummy_input = torch.randn(input_size)
torch.onnx.export(
model,
dummy_input,
"yolov5s.onnx",
opset_version=12,
input_names=['images'],
output_names=['output']
)
print("ONNX 文件导出成功！")
为了进一步优化性能，可以将 ONNX 模型转换为 TensorRT 或 TensorFlow Lite 格式。以下是两种方式的简要介绍：
本文将以 TFLite 为例进行详细说明。
安装依赖
pip install tf-nightly
转换代码
import tensorflow as tf
# 加载 ONNX 模型并转换为 TensorFlow SavedModel
model = tf.saved_model.load('yolov5s_saved_model')
# 转换为 TFLite 格式
converter = tf.lite.TFLiteConverter.from_saved_model('yolov5s_saved_model')
tflite_model = converter.convert()
# 保存 TFLite 文件
with open('yolov5s.tflite', 'wb') as f:
f.write(tflite_model)
print("TFLite 文件导出成功！")
在 Android 项目中，确保添加了 TensorFlow Lite 的依赖项。编辑 build.gradle 文件：
dependencies {
implementation 'org.tensorflow:tensorflow-lite:2.10.0'
implementation 'org.tensorflow:tensorflow-lite-gpu:2.10.0' // 可选：GPU 加速
}
将生成的 yolov5s.tflite 文件放入 assets 目录中，并通过以下代码加载模型：
import org.tensorflow.lite.Interpreter;
// 加载模型
InputStream inputStream = getAssets().open("yolov5s.tflite");
File file = new File(getCacheDir(), "yolov5s.tflite");
FileOutputStream fileOutputStream = new FileOutputStream(file);
byte[] buffer = new byte[1024];
int read;
while ((read = inputStream.read(buffer)) != -1) {
fileOutputStream.write(buffer, 0, read);
}
fileOutputStream.close();
inputStream.close();
Interpreter interpreter = new Interpreter(file);
YOLOv5 模型要求输入图像为固定大小（如 640x640），并且需要归一化处理。以下是一个简单的预处理方法：
public float[][][] preprocess(Bitmap bitmap) {
int inputSize = 640; // 模型输入尺寸
Bitmap resizedBitmap = Bitmap.createScaledBitmap(bitmap, inputSize, inputSize, false);
float[][][] input = new float[1][inputSize][inputSize][3];
for (int y = 0; y < inputSize; y++) {
for (int x = 0; x < inputSize; x++) {
int pixel = resizedBitmap.getPixel(x, y);
input[0][y][x][0] = (Color.red(pixel) / 255.0f); // R
input[0][y][x][1] = (Color.green(pixel) / 255.0f); // G
input[0][y][x][2] = (Color.blue(pixel) / 255.0f); // B
}
}
return input;
}
通过 Interpreter.run() 方法执行推理，并对输出结果进行解析和绘制边界框。
推理代码
float[][][] input = preprocess(bitmap);
float[][] output = new float[1][25200]; // 输出维度根据模型调整
interpreter.run(input, output);
后处理 解析输出结果并绘制边界框。YOLOv5 的输出包含多个候选框，需要通过非极大值抑制（NMS）筛选最终结果。以下是 NMS 的伪代码流程图：
graph TD
A[开始] --> B[获取所有候选框]
B --> C[计算每个框的置信度]
C --> D[按置信度排序]
D --> E[初始化空列表]
E --> F[遍历候选框]
F --> G{是否与已有框重叠？}
G --是--> H[跳过当前框]
G --否--> I[加入最终结果]
I --> J[返回最终结果]
模型压缩
使用量化技术（如 INT8 量化）减少模型大小和推理时间。
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_types = [tf.int8]
GPU 加速
在支持 GPU 的设备上启用 TensorFlow Lite 的 GPU 委托：
Interpreter.Options options = new Interpreter.Options();
options.addDelegate(new GpuDelegate());
Interpreter interpreter = new Interpreter(file, options);
内存管理
确保在应用退出时释放模型资源，避免内存泄漏。