在Android上進(jìn)行yolov5目標(biāo)檢測，使用torchscript方式

環(huán)境

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

前言

前文 在Android上運(yùn)行YOLOv5目標(biāo)檢測 我們介紹過使用ncnn的方式在android設(shè)備上進(jìn)行yolov5的目標(biāo)檢測。本篇介紹另一種方式，即torchscript。

代碼實(shí)踐

這個(gè)demo來自pytorch官方，地址是: https://github.com/pytorch/android-demo-app/tree/master/ObjectDetection 下載下來使用android studio打開備用。

接下來需要準(zhǔn)備torchscript模型，這里使用yolov5 3.0的版本進(jìn)行轉(zhuǎn)換，來到源碼目錄，修改model/export.py文件

model.model[-1].export = True

改為

model.model[-1].export = False

然后，還是在這個(gè)文件，在代碼塊

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

的中間，加入下面兩句

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

修改完成后就可以進(jìn)行轉(zhuǎn)換了

(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.

命令的最后是導(dǎo)出onnx的報(bào)錯(cuò)，由于我們使用的是torchscript而沒有用到onnx，所以這個(gè)錯(cuò)誤可以忽略。將weights/yolov5s.torchscript.pt拷貝android工程中的目錄app/src/main/assets下

注意看，app自帶的測試圖片也是存放在這里文件夾下，classes.txt是目標(biāo)的名稱，如果要替換自己訓(xùn)練的模型，這個(gè)文件也要記得修改。

準(zhǔn)備工具就緒，接下來，我們就使用android studio來編譯并且安裝到android手機(jī)上去，測試自帶的3張圖片

當(dāng)然，這個(gè)app也可以選擇手機(jī)內(nèi)的圖片進(jìn)行檢測

除此以外，還可以使用手機(jī)攝像頭進(jìn)行目標(biāo)檢測

使用自己的模型

yolov5的模型訓(xùn)練請參考這篇 https://xugaoxiang.com/2020/07/02/yolov5-training/，作為測試，我們也使用上文中訓(xùn)練出來的口罩檢測模型

torchscript轉(zhuǎn)換的步驟和上面是一樣的，這里就省略了，把生成的torchscript.pt放到assets目錄下，重命名為yolov5s.torchscript.pt

接著修改classes.txt

mask
no-mask

這里的值必須與你訓(xùn)練時(shí)的保持一致

然后修改文件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")));

根據(jù)自己的情況，修改這2個(gè)文件名

接著修改文件ObjectDetectionActivity.java

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

這個(gè)值是5+class數(shù)量，針對口罩這個(gè)就是5+2=7，同樣是這個(gè)文件，往下拉

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

修改i*后面的值，與mOutputColumn是一樣的。

最后，重新編譯下工程并安裝到手機(jī)上進(jìn)行測試

參考資料

• 借助NCNN，在Android上運(yùn)行YOLOv5目標(biāo)檢測
• windows 10安裝cuda和cudnn
• yolov5模型訓(xùn)練
• yolov5 5.0版本
• Android studio gradle構(gòu)建失敗的解決方法
• https://pytorch.org/tutorials/recipes/script_optimized.html
• https://github.com/pytorch/android-demo-app/issues/102
  
  

努力分享優(yōu)質(zhì)的計(jì)算機(jī)視覺相關(guān)內(nèi)容，歡迎關(guān)注：

個(gè)人微信（如果沒有備注不拉群！）
請注明：地區(qū)+學(xué)校/企業(yè)+研究方向+昵稱



下載1：何愷明頂會分享

在「AI算法與圖像處理」公眾號后臺回復(fù)：何愷明，即可下載。總共有6份PDF，涉及 ResNet、Mask RCNN等經(jīng)典工作的總結(jié)分析

下載2：終身受益的編程指南：Google編程風(fēng)格指南

在「AI算法與圖像處理」公眾號后臺回復(fù)：c++，即可下載。歷經(jīng)十年考驗(yàn)，最權(quán)威的編程規(guī)范！

下載3 CVPR2021

在「AI算法與圖像處理」公眾號后臺回復(fù)：CVPR，即可下載1467篇CVPR 2020論文 和 CVPR 2021 最新論文