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

AI算法与图像处理

共 5391字,需浏览 11分钟

 ·

2021-06-24 22:03

点击下方AI算法与图像处理”,关注一下

重磅干货,第一时间送达

环境

  • windows 10 64bit
  • android studio 4.1.2
  • yolov5 3.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,同样是这个文件,往下拉

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

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

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

yolov5 android torchscript

yolov5 android torchscript

yolov5 android torchscript

yolov5 android torchscript

参考资料

  • 借助NCNN,在Android上运行YOLOv5目标检测
  • windows 10安装cuda和cudnn
  • yolov5模型训练
  • yolov5 5.0版本
  • Android studio gradle构建失败的解决方法
  • https://pytorch.org/tutorials/recipes/script_optimized.html
  • https://github.com/pytorch/android-demo-app/issues/102

努力分享优质的计算机视觉相关内容,欢迎关注:


个人微信(如果没有备注不拉群!
请注明:地区+学校/企业+研究方向+昵称



下载1:何恺明顶会分享


AI算法与图像处理」公众号后台回复:何恺明,即可下载。总共有6份PDF,涉及 ResNet、Mask RCNN等经典工作的总结分析


下载2:终身受益的编程指南:Google编程风格指南


AI算法与图像处理」公众号后台回复:c++,即可下载。历经十年考验,最权威的编程规范!



下载3 CVPR2021

AI算法与图像处公众号后台回复:CVPR即可下载1467篇CVPR 2020论文 和 CVPR 2021 最新论文

点亮 ,告诉大家你也在看




浏览 64
点赞
评论
收藏
分享

手机扫一扫分享

分享
举报
评论
图片
表情
推荐
点赞
评论
收藏
分享

手机扫一扫分享

分享
举报