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

文章目录

    由于 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

    关于作者 🌱

    我是来自山东烟台的一名开发者,有感兴趣的话题,或者软件开发需求,欢迎加微信 zhongwei 聊聊,或者关注我的个人公众号“大象工具”, 查看更多联系方式