欢迎访问本站,希望站内信息可以帮助到您!感兴趣的,可以加入我的知识星球,一对一解答!

在Android上进行yolov5目标检测,使用torchscript方式

YOLO 迷途小书童 17评论

环境

  • windows 10 64bit
  • android studio 4.1.2
  • yolov5 4.0
  • pytorch 1.6+cu101

前言

前文 在Android上运行YOLOv5目标检测 我们介绍过使用 ncnn 的方式在 android 设备上进行 yolov5 的目标检测。本篇介绍另一种方式,即 torchscript

代码实践

这个 demo 来自 pytorch 官方,地址是: https://github.com/pytorch/android-demo-app/tree/master/ObjectDetection 下载下来使用 android studio 打开备用。

接下来需要准备 torchscript 模型,这里使用 yolov5 3.0的版本进行转换,来到源码目录,修改 model/export.py 文件

model.model[-1].export = True

改为

model.model[-1].export = False

然后,还是在这个文件,在代码块

ts = torch.jit.trace(model, img)
ts.save(f)

的中间,加入下面两句

from torch.utils.mobile_optimizer import optimize_for_mobile
ts = optimize_for_mobile(ts)

yolov5 android torchscript

修改完成后就可以进行转换了

(pytorch1.6) xugaoxiang@1070Ti:~/workshop/yolov5-3.0$ python models/export.py --weights weights/yolov5s.pt
Namespace(batch_size=1, img_size=[640, 640], weights='weights/yolov5s.pt')
Fusing layers...
Model Summary: 224 layers, 7266973 parameters, 0 gradients

Starting TorchScript export with torch 1.6.0+cu101...
/home/xugaoxiang/workshop/yolov5-3.0/models/yolo.py:49: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if self.grid[i].shape[2:4] != x[i].shape[2:4]:
/home/xugaoxiang/anaconda3/envs/pytorch1.6/lib/python3.7/site-packages/torch/jit/_trace.py:940: TracerWarning: Encountering a list at the output of the tracer might cause the trace to be incorrect, this is only valid if the container structure does not change based on the module's inputs. Consider using a constant container instead (e.g. for `list`, use a `tuple` instead. for `dict`, use a `NamedTuple` instead). If you absolutely need this and know the side effects, pass strict=False to trace() to allow this behavior.
  _force_outplace,
TorchScript export success, saved as weights/yolov5s.torchscript.pt

Starting ONNX export with onnx 1.8.1...
ONNX export success, saved as weights/yolov5s.onnx

Starting CoreML export with coremltools 4.1...
CoreML export failure: Unknown type __torch__.torch.classes.xnnpack.Conv2dOpContext encountered in graph lowering. This type is not supported in ONNX export.

Export complete (13.42s). Visualize with https://github.com/lutzroeder/netron.

命令的最后是导出 onnx 的报错,由于我们使用的是 torchscript 而没有用到 onnx,所以这个错误可以忽略。将 weights/yolov5s.torchscript.pt 拷贝 android 工程中的目录 app/src/main/assets

yolov5 android torchscript

注意看,app 自带的测试图片也是存放在这里文件夹下,classes.txt 是目标的名称,如果要替换自己训练的模型,这个文件也要记得修改。

准备工具就绪,接下来,我们就使用 android studio 来编译并且安装到 android 手机上去,测试自带的3张图片

yolov5 android torchscript

yolov5 android torchscript

yolov5 android torchscript

当然,这个 app 也可以选择手机内的图片进行检测

yolov5 android torchscript

除此以外,还可以使用手机摄像头进行目标检测

yolov5 android torchscript

使用自己的模型

yolov5 的模型训练请参考这篇 https://xugaoxiang.com/2020/07/02/yolov5-training/,作为测试,我们也使用上文中训练出来的口罩检测模型

torchscript 转换的步骤和上面是一样的,这里就省略了,把生成的 torchscript.pt 放到 assets 目录下,重命名为 yolov5s.torchscript.pt

接着修改 classes.txt

mask
no-mask

这里的值必须与你训练时的保持一致

然后修改文件 app/src/main/java/org/pytorch/demo/objectdetection/MainActivity.java

mModule = PyTorchAndroid.loadModuleFromAsset(getAssets(), "best.torchscript.pt");
BufferedReader br = new BufferedReader(new InputStreamReader(getAssets().open("classes.txt")));

根据自己的情况,修改这2个文件名

接着修改文件 ObjectDetectionActivity.java

private static int mOutputColumn = 7; // left, top, right, bottom, score and 80 class probability

这个值是 5+class 数量,针对口罩这个就是 5+2=7,

然后,是 PrePostProcessor.java 这个文件,往下拉到最后部分

Result result = new Result(cls, outputs[i*7+4], rect);

修改 i* 后面的值,与 mOutputColumn 是一样的。

最后,重新编译下工程并安装到手机上进行测试

yolov5 android torchscript

yolov5 android torchscript

yolov5 android torchscript

yolov5 android torchscript

参考资料

喜欢 (14)

您必须 登录 才能发表评论!

(17)个小伙伴在吐槽
  1. 小白想问问是不是只要安装好pytorch和yolo后,在cmd的yolo环境下安装您上面的步骤来就可以完成了?
    匿名2022-07-08 18:03
    • 在pc端将pt转换成torchscript,然后替换android工程中的torchscript,最后修改其它参数
      迷途小书童2022-07-09 11:31
  2. 请问我不清楚torchscript这些内容,我应该去学习什么知识才知道怎么操作这个呢?
    匿名2022-07-08 13:14
    • 参考官方文档 https://pytorch.org/tutorials/beginner/Intro_to_TorchScript_tutorial.html?highlight=torchscript 和 https://pytorch.org/mobile/android/
      迷途小书童2022-07-08 14:17
  3. results = model(img) 博主,我看了下4.0版本和6.0版本模型输出不同,4.0直接是[16,25200,85] ,6.0是一个tensorlist[16,3,80,80,85][16,3,40,40,85][16,3,80,80,85],因此就不能运行通,您这边知道应如何修改吗,谢谢
    匿名2021-11-05 01:56
  4. 博主,请问一下4.0版本的yoko v5的权重文件可以运行,但最近尝试了一下6.0版本的,就不可以了,在推理(按detect)的时候就会闪退,您知道是什么原因吗,谢谢~
    匿名2021-11-04 18:19
  5. 谢谢博主,这里提醒一下5月30号以后版本的最后一步需要修改的地方是PrePostProcessor.java
    匿名2021-08-12 21:32
  6. 博主,请问一下模型输出为什么是mOutputRow=25200呢,谢谢
    匿名2021-08-12 16:39
  7. 想知道博主使用的是git上几月几号的objectdetection
    匿名2021-08-11 12:27
    • 应该是上一个commit版本。最新的是加入了pytorch1.9的支持。
      迷途小书童2021-08-11 20:56
  8. 博主你好为什么我的apk会闪退,我用的是自己训练的模型
    匿名2021-07-31 18:59
    • 在线调试一下,先定位问题
      迷途小书童2021-08-04 11:21
  9. 不好意思我知道了live怎么退出,但是感觉自己的模型live时很难看到效果
    匿名2021-07-14 14:53
  10. 您好,我在使用live时,所获取的视频会变形,另外,请问一下live怎么退出呢
    匿名2021-07-14 14:51
  11. 您好,为什么我运行avd时,模拟器的桌面是空白的,没有yolov5的图标(应用程序),请问怎么解决呢?
    匿名2021-06-23 22:58
    • 桌面没有的话,你到应用抽屉里看看。首先确保,app已经安装成功了? app的图标是pytorch的图标
      迷途小书童2021-06-25 10:50
  12. 请问如何解决No module named 'utils'的问题?我看utils文件夹下是有__init__.py文件的
    匿名2021-06-20 18:00