yolov5 pytorch mobile 模型导入 Android,实现图片目标检测

发布时间: 2023-03-29 21:12:17 作者: 大象笔记

由于 yolov5 相关的 TensorFlow Lite 示例代码太少,而且都是旧版本的示例,导致我完全不知道怎么上手。 所以转投 pytorch mobile 的方案,主要是 facebook 官方的示例代码非常完整,且清晰明了,容易修改。

再就是国内 Android 手机普遍不能使用 google play 的服务,所以 TensorFlow lite 安装包小的优势也就没有了。至于性能,暂时没有要求,而且实测来看,pytorch mobile 的识别速度就非常快,几乎瞬间完成,没有做计时,估计也是几十毫秒的范围,已经很满意了。

识别效果

以识别照片中小麦穗数量为例:

导出 pytorch mobile 模型

在 Google Colab 的免费 GPU 服务器上执行:

!python export.py --weights /content/drive/MyDrive/yolo/wheat_head_best.pt --include torchscript --optimize

并不需要按照 pytorch mobile 官方文档中说的修改 export.py 的代码,因为最新的 yolov5 支持了 optimize 参数,自动转换为 torchscript lite 版本的模型。

if optimize:  # https://pytorch.org/tutorials/recipes/mobile_interpreter.html
    optimize_for_mobile(ts)._save_for_lite_interpreter(str(f), _extra_files=extra_files)

模型文件大小涨了一倍

https://discuss.pytorch.org/t/optimize-for-mobile-exports-are-double-the-size-of-normal-torchscript/128368

原来的 pt 模型只有 14M,但是转换后的 torchscript lite 版本却有 27M。暂时不纠结这个问题了,先在 Android 上跑通再说。

Lite Interpreter version number does not match

com.facebook.jni.CppException: Lite Interpreter version number does not match. The model version must be between 3 and 7 but the model version is 8 ()

1.10.0 升级到 1.13.1

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

但是国内镜像好像更新慢,可以换成 1.13.0 版本。

目标标注原理

下层是一个 image view,上层是一个等宽高的 canvas,在其中标注方框, class name 及 conf 值。不过对于单种类识别,描述没啥用,我就去掉了。

java 与 koltin 混合开发初尝试

pytorch mobile demo 里用的是 java 实现的,但是我的工程是基于 kotlin 创建的。在复制功能模块时,我对 android studio 的自动代码转换还是不放心,于是就直接复制的 java 文件。 没想到 java 与 kotlin 文件间调用如丝般顺滑。

除了一处包内 array list type 的小问题,其他均正常工作。

pytorch mobile 对比 TensorFlow Lite

Faster R-CNN 与 yolov5 对比

我感觉 yolov5 资料更多,且上手毫无门槛,更适合新手。

演示视频

参考

我是一名山东烟台的开发者,联系作者