在Android上进行yolov5目标检测,使用torchscript方式
共 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)
修改完成后就可以进行转换了
(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
下
注意看,app
自带的测试图片也是存放在这里文件夹下,classes.txt
是目标的名称,如果要替换自己训练的模型,这个文件也要记得修改。
准备工具就绪,接下来,我们就使用android studio
来编译并且安装到android
手机上去,测试自带的3张图片
当然,这个app
也可以选择手机内的图片进行检测
除此以外,还可以使用手机摄像头进行目标检测
使用自己的模型
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
是一样的。
最后,重新编译下工程并安装到手机上进行测试
参考资料
借助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 最新论文
点亮 ,告诉大家你也在看