yolov5部署到android studio

环境

Ubuntu22.10
Pytorch2.0.1+cu117
Android Studio Flamingo | 2022.2.1 Patch 1

获取demo

git clone https://github.com/pytorch/android-demo-app.git

打开android-demo-app/ObjectDectection项目。
打开app/build.gradle文件，修改pytorch版本
这里遇到最大的坑，就是训练的pytorch版本和android的包版本相差过大，网上大多数安卓这里都是1.10.0

implementation 'org.pytorch:pytorch_android_lite:1.13.1'
implementation 'org.pytorch:pytorch_android_torchvision_lite:1.13.1'

将pt文件导出为ptl文件

修改yolov5中的export.py，其中export_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)')
        print(f'{prefix} export success, saved as {fl} ({file_size(fl):.1f} MB)')
    except Exception as e:
        print(f'{prefix} export failure: {e}')
    return fl, None

运行导出ptl文件

python export.py --weights <your pt file path> --include torchscript

在pt文件的目录下会发现这样的文件，第二个最后后缀为ptl的是我们所需要的文件

修改demo

修改PrePostProcessor

在PrePostProcessor文件中将mOutputColumn修改为网络模型类别数+5，我这里网络模型类别数是2,所以写7

增加ptl文件并增加类别文件

在assets文件夹中放入之前导出的ptl文件并且添加一个类别文件，类别文件中填写网络模型类别的name，以回车间隔，如下图

修改MainActivity

这里修改为ptl文件名和classes文件名

大功告成