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

更新日期: 2023-04-03 阅读次数: 3068 字数: 808 分类: AI

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

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

识别效果

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

yolov5 Android 识别小麦穗

导出 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

  • tflite 毕竟是 Android 的亲儿子,模型导入非常方便,省去了手写代码导入的繁琐步骤。但是对应的概念也很多,metadata,tensor 对新手都不太友好。我感觉已经算是有一点基础了,依然觉得上手困难。而且在 github 中搜素关键代码,居然一行实际代码都没有。
  • pytorch mobile 配合 yolov5 简直就是菜鸟神器,轻松上手。但是导出的 torchscript lite model 在 google colab 上无法验证,执行 detect.py 报错,不如 yolov5 tflite 模型方便验证。

Faster R-CNN 与 yolov5 对比

  • Faster R-CNN 在准确率方面表现更出色,适合对精度要求高的场景
  • yolov5 检测速度更快,适合对实时性要求高的场景

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

演示视频

参考

  • 完整代码 1k star,有效果图 https://github.com/pytorch/android-demo-app/tree/master/ObjectDetection

tags: pytorch yolov5

关于作者 🌱

我是来自山东烟台的一名开发者,有敢兴趣的话题,或者软件开发需求,欢迎加微信 zhongwei 聊聊, 查看更多联系方式