前言
本篇博文用来研究YOLOv5在Android上部署的例程
主要参考的是Pytorch官方提供的Demo：https://github.com/pytorch/android-demo-app/tree/master/PyTorchDemoApp
功能简述
App主页如下图所示：
主要功能：
-
切换测试图片
在程序中直接指定三张(或任意张)图片，点击测试图片，可以切换图片 -
选择图片
点击选择图片，可以在相册中选择一张图片，也可以直接进行拍照 -
实时视频
点击实时视频，可以开启摄像头，直接在摄像预览中显示检测结果 -
切换模型(我添加的功能)
点击切换模型，可以选择不同的模型进行检测
快速上手
首先来跑通官方Demo，首先下载官方提供的yolov5s.torchscript.ptl
下载链接：https://pytorch-mobile-demo-apps.s3.us-east-2.amazonaws.com/yolov5s.torchscript.ptl
下载完放到assets文件夹下
直接运行，从相册中选择图片时会报错：
Unable to decode stream: java.io.FileNotFoundException:/…/open failed: EACCES (Permission denied)
此时需要在AndroidManifest.xml的application标签中添加一句：
android:requestLegacyExternalStorage="true"
然后就可以正常运行了
训练自己的模型
下面用YOLOv5-6.0版本训练自己的模型，怎么训练不做赘述，可以参考本专栏的往期博文。
然后修改export.py中的export_torchscript函数，主要添加三行代码，用以导出.torchscript.ptl后缀模型。
def export_torchscript(model, im, file, optimize, prefix=colorstr('TorchScript:')):
# YOLOv5 TorchScript model export
try:
print(f'\n{
prefix} starting export with torch {
torch.__version__}...')
f = file.with_suffix('.torchscript.pt')
f = str(f)
fl = file.with_suffix('.torchscript.ptl')
ts = torch.jit.trace(model, im, strict=False)
(optimize_for_mobile(ts) if optimize else ts).save(f)
(optimize_for_mobile(ts) if optimize else ts)._save_for_lite_interpreter(str(fl))
print(f'{
prefix} export success, saved as {
f} ({
file_size(f):.1f} MB)')
except Exception as e:
print(f'{
prefix} export failure: {
e}')
然后在终端运行：
python export.py --weights runs/train/exp/weights/best.pt --include torchscript
运行完得到best.torchscript.ptl模型
切换自己的模型
下面来添加一个切换模型的功能，并使用自己训练的模型。
首先修改pytorch依赖版本，修改build.gradle中的依赖：
implementation 'org.pytorch:pytorch_android_lite:1.9.0'
implementation 'org.pytorch:pytorch_android_torchvision:1.9.0'
这里的版本尽量和后面训练用的pytorch版本对应，比如后面自己用的pytorch版本是1.9.0，这里就写1.9.0。
然后修改ObjectDetectionActivitys,java，这里将mOutputColumn的private修饰符去掉，使其可以在外部访问：
接下来修改xml界面，在activity_main.xml中添加切换模型按钮，并调整布局
<Button
android:id="@+id/select"
android:layout_width="100dp"
android:layout_height="wrap_content"
android:layout_marginTop="32dp"
android:textAllCaps="false"
android:text="@string/select_model"
app:layout_constraintEnd_toEndOf="parent"
app:layout_constraintEnd_toStartOf="@+id/selectButton"
app:layout_constraintHorizontal_bias="0.5"
app:layout_constraintStart_toStartOf="parent"
app:layout_constraintTop_toBottomOf="@+id/detectButton"
android:background="@drawable/button_selector"/>
<Button
android:id="@+id/testButton"
android:layout_width="wrap_content"
android:layout_height="wrap_content"
android:layout_marginTop="180dp"
android:textAllCaps="false"
app:layout_constraintEnd_toEndOf="parent"
app:layout_constraintHorizontal_bias="0.5"
app:layout_constraintStart_toStartOf="parent"
app:layout_constraintTop_toBottomOf="@+id/imageView"
android:background="@drawable/button_selector"/>
然后修改MainActivity.java，添加以下三个属性
private String model_name = "yolov5s.torchscript.ptl";
private String model_class = "classes.txt";
private int num_class = 80;
添加选择模型按钮响应：
private void ShowChoise()
{
AlertDialog.Builder builder = new AlertDialog.Builder(MainActivity.this);
// builder.setIcon(R.drawable.ic_launcher_foreground);
builder.setTitle("选择一个模型");
// 指定下拉列表的显示数据
final String[] cities = {
"YOLOv5s", "王者荣耀模型"};
// 设置一个下拉的列表选择项
builder.setItems(cities, new DialogInterface.OnClickListener()
{
@Override
public void onClick(DialogInterface dialog, int which)
{
Toast.makeText(MainActivity.this, "选择的模型为：" + cities[which], Toast.LENGTH_SHORT).show();
if (which==0){
model_name = "yolov5s.torchscript.ptl";
model_class = "classes.txt";
num_class = 80;
}
else {
model_name = "mymodel.ptl";
model_class = "classes_wzry.txt";
num_class = 10;
}
// 重新加载
try {
mModule = LiteModuleLoader.load(MainActivity.assetFilePath(getApplicationContext(), model_name));
BufferedReader br = new BufferedReader(new InputStreamReader(getAssets().open(model_class)));
String line;
List<String> classes = new ArrayList<>();
while ((line = br.readLine()) != null) {
classes.add(line);
}
PrePostProcessor.mClasses = new String[classes.size()];
PrePostProcessor.mOutputColumn = num_class + 5;
classes.toArray(PrePostProcessor.mClasses);
} catch (IOException e) {
Log.e("Object Detection", "Error reading assets", e);
finish();
}
}
});
builder.show();
}
这里选择的模型数量添加if分支，model_class为模型对应的类别标签，需要仿照classes.txt单独创建，num_class为类别数量。
最后将之上一步得到的best.torchscript.ptl复制到assets文件夹下，注意需要手动修改文件名mymodel.ptl，这里不改名会发生文件找不到的报错，最后再运行即可。
完整代码
除了上面这部分，还对界面进行了汉化，图片加载做了微调，几个修改过的文件的完整源码如下：
activity_main.xml
<?xml version="1.0" encoding="utf-8"?>
<androidx.constraintlayout.widget.ConstraintLayout xmlns:android="http://schemas.android.com/apk/res/android"
xmlns:app="http://schemas.android.com/apk/res-auto"
xmlns:tools="http://schemas.android.com/tools"
android:layout_width="match_parent"
android:layout_height="match_parent"
tools:context="org.pytorch.demo.objectdetection.MainActivity">
<ImageView
android:id="@+id/imageView"
android:layout_width="0dp"
android:layout_height="0dp"
android:layout_marginTop="0dp"
android:background="#FFFFFF"
android:contentDescription="@string/image_view"
app:layout_constraintDimensionRatio="1:1"
app:layout_constraintEnd_toEndOf="parent"
app:layout_constraintStart_toStartOf="parent"
app:layout_constraintTop_toTopOf="parent" />
<org.pytorch.demo.objectdetection.ResultView
android:id="@+id/resultView"
android:layout_width="0dp"
android:layout_height="0dp"
android:layout_marginTop="0dp"
app:layout_constraintDimensionRatio="1:1"
app:layout_constraintEnd_toEndOf="parent"
app:layout_constraintStart_toStartOf="parent"
app:layout_constraintTop_toTopOf="parent" />
<Button
android:id="@+id/detectButton"
android:layout_width="wrap_content"
android:layout_height="wrap_content"
android:layout_marginTop="20dp"
android:text="@string/detect"
android:textAllCaps="false"
app:layout_constraintEnd_toEndOf="parent"
app:layout_constraintHorizontal_bias="0.498"
app:layout_constraintStart_toStartOf="parent"
app:layout_constraintTop_toBottomOf="@+id/imageView"
android:background="@drawable/button_selector"/>
<ProgressBar
android:id="@+id/progressBar"
android:layout_width="wrap_content"
android:layout_height="wrap_content"
android:layout_marginTop="20dp"
android:visibility="invisible"
app:layout_constraintEnd_toEndOf="parent"
app:layout_constraintHorizontal_bias="0.498"
app:layout_constraintStart_toStartOf="parent"
app:layout_constraintTop_toBottomOf="@+id/imageView" />
<Button
android:id="@+id/selectButton"
android:layout_width="100dp"
android:layout_height="wrap_content"
android:text="@string/select"
android:textAllCaps="false"
app:layout_constraintEnd_toEndOf="parent"
app:layout_constraintEnd_toStartOf="@+id/liveButton"
app:layout_constraintHorizontal_bias="0.5"
app:layout_constraintStart_toEndOf="@+id/select"
app:layout_constraintTop_toTopOf="@+id/select"
android:background="@drawable/button_selector"/>
<Button
android:id="@+id/liveButton"
android:layout_width="100dp"
android:layout_height="wrap_content"
android:text="@string/live"
android:textAllCaps="false"
app:layout_constraintEnd_toEndOf="parent"
app:layout_constraintHorizontal_bias="0.5"
app:layout_constraintStart_toEndOf="@+id/selectButton"
app:layout_constraintTop_toTopOf="@+id/selectButton"
android:background="@drawable/button_selector"/>
<Button
android:id="@+id/select"
android:layout_width="100dp"
android:layout_height="wrap_content"
android:layout_marginTop="32dp"
android:textAllCaps="false"
android:text="@string/select_model"
app:layout_constraintEnd_toEndOf="parent"
app:layout_constraintEnd_toStartOf="@+id/selectButton"
app:layout_constraintHorizontal_bias="0.5"
app:layout_constraintStart_toStartOf="parent"
app:layout_constraintTop_toBottomOf="@+id/detectButton"
android:background="@drawable/button_selector"/>
<Button
android:id="@+id/testButton"
android:layout_width="wrap_content"
android:layout_height="wrap_content"
android:layout_marginTop="180dp"
android:textAllCaps="false"
app:layout_constraintEnd_toEndOf="parent"
app:layout_constraintHorizontal_bias="0.5"
app:layout_constraintStart_toStartOf="parent"
app:layout_constraintTop_toBottomOf="@+id/imageView"
android:background="@drawable/button_selector"/>
</androidx.constraintlayout.widget.ConstraintLayout>
MainActivity.java
// Copyright (c) 2020 Facebook, Inc. and its affiliates.
// All rights reserved.
//
// This source code is licensed under the BSD-style license found in the
// LICENSE file in the root directory of this source tree.
package org.pytorch.demo.objectdetection;
import androidx.appcompat.app.AlertDialog;
import androidx.appcompat.app.AppCompatActivity;
import androidx.core.app.ActivityCompat;
import androidx.core.content.ContextCompat;
import android.Manifest;
import android.content.Context;
import android.content.DialogInterface;
import android.content.Intent;
import android.content.pm.PackageManager;
import android.database.Cursor;
import android.graphics.Bitmap;
import android.graphics.BitmapFactory;
import android.graphics.Matrix;
import android.net.Uri;
import android.os.Bundle;
import android.provider.MediaStore;
import android.util.Log;
import android.view.View;
import android.widget.Button;
import android.widget.ImageView;
import android.widget.ProgressBar;
import android.widget.Toast;
import org.pytorch.IValue;
import org.pytorch.LiteModuleLoader;
import org.pytorch.Module;
import org.pytorch.Tensor;
import org.pytorch.torchvision.TensorImageUtils;
import java.io.BufferedReader;
import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.io.OutputStream;
import java.util.ArrayList;
import java.util.List;
public class MainActivity extends AppCompatActivity implements Runnable {
private int mImageIndex = 0;
private String[] mTestImages = {
"test1.png", "test2.jpg", "test3.png"};
private ImageView mImageView;
private ResultView mResultView;
private Button mButtonDetect;
private Button mButtonSelect;
private ProgressBar mProgressBar;
private Bitmap mBitmap = null;
private Module mModule = null;
private float mImgScaleX, mImgScaleY, mIvScaleX, mIvScaleY, mStartX, mStartY;
private String model_name = "yolov5s.torchscript.ptl";
private String model_class = "classes.txt";
private int num_class = 80;
public static String assetFilePath(Context context, String assetName) throws IOException {
File file = new File(context.getFilesDir(), assetName);
if (file.exists() && file.length() > 0) {
return file.getAbsolutePath();
}
try (InputStream is = context.getAssets().open(assetName)) {
try (OutputStream os = new FileOutputStream(file)) {
byte[] buffer = new byte[4 * 1024];
int read;
while ((read = is.read(buffer)) != -1) {
os.write(buffer, 0, read);
}
os.flush();
}
return file.getAbsolutePath();
}
}
private void ShowChoise()
{
AlertDialog.Builder builder = new AlertDialog.Builder(MainActivity.this);
// builder.setIcon(R.drawable.ic_launcher_foreground);
builder.setTitle("选择一个模型");
// 指定下拉列表的显示数据
final String[] cities = {
"YOLOv5s", "王者荣耀模型"};
// 设置一个下拉的列表选择项
builder.setItems(cities, new DialogInterface.OnClickListener()
{
@Override
public void onClick(DialogInterface dialog, int which)
{
Toast.makeText(MainActivity.this, "选择的模型为：" + cities[which], Toast.LENGTH_SHORT).show();
if (which==0){
model_name = "yolov5s.torchscript.ptl";
model_class = "classes.txt";
num_class = 80;
}
else {
model_name = "mymodel.ptl";
model_class = "classes_wzry.txt";
num_class = 10;
}
// 重新加载
try {
mModule = LiteModuleLoader.load(MainActivity.assetFilePath(getApplicationContext(), model_name));
BufferedReader br = new BufferedReader(new InputStreamReader(getAssets().open(model_class)));
String line;
List<String> classes = new ArrayList<>();
while ((line = br.readLine()) != null) {
classes.add(line);
}
PrePostProcessor.mClasses = new String[classes.size()];
PrePostProcessor.mOutputColumn = num_class + 5;
classes.toArray(PrePostProcessor.mClasses);
} catch (IOException e) {
Log.e("Object Detection", "Error reading assets", e);
finish();
}
}
});
builder.show();
}
@Override
protected void onCreate(Bundle savedInstanceState) {
super.onCreate(savedInstanceState);
if (ContextCompat.checkSelfPermission(this, Manifest.permission.READ_EXTERNAL_STORAGE) != PackageManager.PERMISSION_GRANTED) {
ActivityCompat.requestPermissions(this, new String[]{
Manifest.permission.READ_EXTERNAL_STORAGE}, 1);
}
if (ContextCompat.checkSelfPermission(this, Manifest.permission.CAMERA) != PackageManager.PERMISSION_GRANTED) {
ActivityCompat.requestPermissions(this, new String[]{
Manifest.permission.CAMERA}, 1);
}
setContentView(R.layout.activity_main);
try {
mBitmap = BitmapFactory.decodeStream(getAssets().open(mTestImages[mImageIndex]));
} catch (IOException e) {
Log.e("Object Detection", "Error reading assets", e);
finish();
}
mImageView = findViewById(R.id.imageView);
mImageView.setImageBitmap(mBitmap);
mResultView = findViewById(R.id.resultView);
mResultView.setVisibility(View.INVISIBLE);
final Button buttonTest = findViewById(R.id.testButton);
buttonTest.setText(("测试图片 1/3"));
buttonTest.setOnClickListener(new View.OnClickListener() {
public void onClick(View v) {
mResultView.setVisibility(View.INVISIBLE);
mImageIndex = (mImageIndex + 1) % mTestImages.length;
buttonTest.setText(String.format("测试图片 %d/%d", mImageIndex + 1, mTestImages.length));
try {
mBitmap = BitmapFactory.decodeStream(getAssets().open(mTestImages[mImageIndex]));
mImageView.setImageBitmap(mBitmap);
} catch (IOException e) {
Log.e("Object Detection", "Error reading assets", e);
finish();
}
}
});
final Button buttonSelect = findViewById(R.id.selectButton);
buttonSelect.setOnClickListener(new View.OnClickListener() {
public void onClick(View v) {
mResultView.setVisibility(View.INVISIBLE);
final CharSequence[] options = {
"从相册选择", "拍照", "取消" };
AlertDialog.Builder builder = new AlertDialog.Builder(MainActivity.this);
builder.setTitle("新测试图片");
builder.setItems(options, new DialogInterface.OnClickListener() {
@Override
public void onClick(DialogInterface dialog, int item) {
if (options[item].equals("拍照")) {
Intent takePicture = new Intent(android.provider.MediaStore.ACTION_IMAGE_CAPTURE);
startActivityForResult(takePicture, 0);
}
else if (options[item].equals("从相册选择")) {
Intent pickPhoto = new Intent(Intent.ACTION_PICK, android.provider.MediaStore.Images.Media.INTERNAL_CONTENT_URI);
startActivityForResult(pickPhoto , 1);
}
else if (options[item].equals("取消")) {
dialog.dismiss();
}
}
});
builder.show();
}
});
final Button buttonLive = findViewById(R.id.liveButton);
buttonLive.setOnClickListener(new View.OnClickListener() {
public void onClick(View v) {
final Intent intent = new Intent(MainActivity.this, ObjectDetectionActivity.class);
startActivity(intent);
}
});
mButtonDetect = findViewById(R.id.detectButton);
mProgressBar = (ProgressBar) findViewById(R.id.progressBar);
mButtonDetect.setOnClickListener(new View.OnClickListener() {
public void onClick(View v) {
mButtonDetect.setEnabled(false);
mProgressBar.setVisibility(ProgressBar.VISIBLE);
mButtonDetect.setText(getString(R.string.run_model));
mImgScaleX = (float)mBitmap.getWidth() / PrePostProcessor.mInputWidth;
mImgScaleY = (float)mBitmap.getHeight() / PrePostProcessor.mInputHeight;
mIvScaleX = (mBitmap.getWidth() > mBitmap.getHeight() ? (float)mImageView.getWidth() / mBitmap.getWidth() : (float)mImageView.getHeight() / mBitmap.getHeight());
mIvScaleY = (mBitmap.getHeight() > mBitmap.getWidth() ? (float)mImageView.getHeight() / mBitmap.getHeight() : (float)mImageView.getWidth() / mBitmap.getWidth());
mStartX = (mImageView.getWidth() - mIvScaleX * mBitmap.getWidth())/2;
mStartY = (mImageView.getHeight() - mIvScaleY * mBitmap.getHeight())/2;
Thread thread = new Thread(MainActivity.this);
thread.start();
}
});
// 新增选择模型按钮
mButtonSelect = findViewById(R.id.select);
mButtonSelect.setOnClickListener(new View.OnClickListener() {
public void onClick(View v) {
ShowChoise();
}
});
try {
mModule = LiteModuleLoader.load(MainActivity.assetFilePath(getApplicationContext(), model_name));
BufferedReader br = new BufferedReader(new InputStreamReader(getAssets().open(model_class)));
String line;
List<String> classes = new ArrayList<>();
while ((line = br.readLine()) != null) {
classes.add(line);
}
PrePostProcessor.mClasses = new String[classes.size()];
PrePostProcessor.mOutputColumn = num_class;
classes.toArray(PrePostProcessor.mClasses);
} catch (IOException e) {
Log.e("Object Detection", "Error reading assets", e);
finish();
}
}
@Override
protected void onActivityResult(int requestCode, int resultCode, Intent data) {
super.onActivityResult(requestCode, resultCode, data);
if (resultCode != RESULT_CANCELED) {
switch (requestCode) {
case 0:
if (resultCode == RESULT_OK && data != null) {
mBitmap = (Bitmap) data.getExtras().get("data");
Matrix matrix = new Matrix();
//matrix.postRotate(90.0f);
matrix.postRotate(0);
mBitmap = Bitmap.createBitmap(mBitmap, 0, 0, mBitmap.getWidth(), mBitmap.getHeight(), matrix, true);
mImageView.setImageBitmap(mBitmap);
}
break;
case 1:
if (resultCode == RESULT_OK && data != null) {
Uri selectedImage = data.getData();
String[] filePathColumn = {
MediaStore.Images.Media.DATA};
if (selectedImage != null) {
Cursor cursor = getContentResolver().query(selectedImage,
filePathColumn, null, null, null);
if (cursor != null) {
cursor.moveToFirst();
int columnIndex = cursor.getColumnIndex(filePathColumn[0]);
String picturePath = cursor.getString(columnIndex);
mBitmap = BitmapFactory.decodeFile(picturePath);
Matrix matrix = new Matrix();
//matrix.postRotate(90.0f);
matrix.postRotate(0);
mBitmap = Bitmap.createBitmap(mBitmap, 0, 0, mBitmap.getWidth(), mBitmap.getHeight(), matrix, true);
mImageView.setImageBitmap(mBitmap);
cursor.close();
}
}
}
break;
}
}
}
@Override
public void run() {
Bitmap resizedBitmap = Bitmap.createScaledBitmap(mBitmap, PrePostProcessor.mInputWidth, PrePostProcessor.mInputHeight, true);
final Tensor inputTensor = TensorImageUtils.bitmapToFloat32Tensor(resizedBitmap, PrePostProcessor.NO_MEAN_RGB, PrePostProcessor.NO_STD_RGB);
IValue[] outputTuple = mModule.forward(IValue.from(inputTensor)).toTuple();
final Tensor outputTensor = outputTuple[0].toTensor();
final float[] outputs = outputTensor.getDataAsFloatArray();
final ArrayList<Result> results = PrePostProcessor.outputsToNMSPredictions(outputs, mImgScaleX, mImgScaleY, mIvScaleX, mIvScaleY, mStartX, mStartY);
runOnUiThread(() -> {
mButtonDetect.setEnabled(true);
mButtonDetect.setText(getString(R.string.detect));
mProgressBar.setVisibility(ProgressBar.INVISIBLE);
mResultView.setResults(results);
mResultView.invalidate();
mResultView.setVisibility(View.VISIBLE);
});
}
}
strings.xml
<resources>
<string name="app_name">YOLOv5</string>
<string name="image_view">Image View</string>
<string name="detect">检测</string>
<string name="run_model">正在运行，请稍后</string>
<string name="restart">Restart</string>
<string name="select">选择图片</string>
<string name="live">实时视频</string>
<string name="select_model">切换模型</string>
</resources>
button_selector.xml
<?xml version="1.0" encoding="utf-8"?>
<selector xmlns:android="http://schemas.android.com/apk/res/android">
<item android:state_pressed="true">
<shape>
<solid android:color="#64AFFA"/>
<corners android:radius="10dp"/>
<padding
android:bottom="2dp"
android:left="3dp"
android:right="3dp"
android:top="2dp">
</padding>
</shape>
</item>
<item android:state_pressed="false">
<shape>
<solid android:color="#99CCFF"/>
<corners android:radius="10dp"/>
<padding
android:bottom="2dp"
android:left="3dp"
android:right="3dp"
android:top="2dp">
</padding>
</shape>
</item>
</selector>
总结
经过实测，整个APK文件打包出来有1点多G，由此可见pytorch框架一加进去体积就会变得很大，后续轻量化还有研究空间。同时，视频实时检测，帧率很低，基本卡成PPT，可能是受限于手机的算力不足，后续也有待研究优化。