最近需要在手机端实现一个目标检测的功能,于是选择了小巧又在目标检测方面表现很好的yolov5s,官网下载yolov5代码,用自己做的数据集进行了训练,然后把模型转换成torchscript格式,这些过程网上都有很多讲解,不再赘述。主要讲一下在安卓上推理的代码。
pytorch在安卓上的使用官方demo,主要代码如下:
Bitmap bitmap = null;Module module = null;try {// creating bitmap from packaged into app android asset 'image.jpg',// app/src/main/assets/image.jpgbitmap = BitmapFactory.decodeStream(getAssets().open("image.jpg"));// loading serialized torchscript module from packaged into app android asset model.pt,// app/src/model/assets/model.ptmodule = LiteModuleLoader.load(assetFilePath(this, "model.pt"));} catch (IOException e) {Log.e("PytorchHelloWorld", "Error reading assets", e);finish();}// showing image on UIImageView imageView = findViewById(R.id.image);imageView.setImageBitmap(bitmap);// preparing input tensorfinal Tensor inputTensor = TensorImageUtils.bitmapToFloat32Tensor(bitmap,TensorImageUtils.TORCHVISION_NORM_MEAN_RGB, TensorImageUtils.TORCHVISION_NORM_STD_RGB, MemoryFormat.CHANNELS_LAST);// running the modelfinal Tensor outputTensor = module.forward(IValue.from(inputTensor)).toTensor();// getting tensor content as java array of floatsfinal float[] scores = outputTensor.getDataAsFloatArray();// searching for the index with maximum scorefloat maxScore = -Float.MAX_VALUE;int maxScoreIdx = -1;for (int i = 0; i < scores.length; i++) {if (scores[i] > maxScore) {maxScore = scores[i];maxScoreIdx = i;}}String className = ImageNetClasses.IMAGENET_CLASSES[maxScoreIdx];
但是这段代码中用的模型不是yolov5,直接用于yolov5的模型是跑不通的,首先计算outputTensor的时候直接把模型输出toTensor(),这个会报错,报错讲说期望Tensor类型但是给了个Tuple,由此可知模型的输出IValue其内置类型是Tuple,于是toTuple然后取第一个元素再toTensor()就可以了。原因是yolov5的输出在Tensor外面又包装了一层,组成了一个Tuple。
然后是结果scores的解析方法,对于yolov5,当有n个目标类别的时候,这个scores的含义是[x,y,w,h,conf,type1score,type2score,......typenscore,x,y,w,h,conf,type1score,type2score,....typenscore......],一直重复25200次,其中x,y是目标框的中心坐标,w,h是目标框的宽高,conf是框的置信度,后面分别是n个类别的得分。所以自然不能用上述代码中的方法取结果。
等我修改完这两处之后,代码可以正常运行,但奇怪的是在python上运行训练好的模型,结果是非常好的,基本95%的时候都可以获取到目标物体在图像中的最小外接矩形,其它5%也只是偏移一点点,但到了手机上,这个结果常常不准确,检测框没有包住目标物体的所有部分是很大概率的事,一开始我怀疑是模型转换的时候丢失了精度,但后来发现转换成torchscript并没有量化,并且在不量化的情况下,模型没必要把一些参数进行修改,这不是努力降精度吗?不合常理。于是仔细看了下yolov5源码中的推理部分,发现图片在进入模型之前,进行了/255的归一化操作。于是乎问题聚集到了原来代码中的这一行:
TensorImageUtils.bitmapToFloat32Tensor(bitmap,
TensorImageUtils.TORCHVISION_NORM_MEAN_RGB, TensorImageUtils.TORCHVISION_NORM_STD_RGB, MemoryFormat.CHANNELS_LAST);
经过了多次调试,终于发现这个函数其实是对bitmap的像素值进行了/255的归一化后,再使用传入的均值数组和标准差数组对归一化过的数值进行了Z-score归一化。Z-score归一化的目的原本是为了让数据符合标准正态分布,但是进入TensorImageUtils类可以看到:
public static float[] TORCHVISION_NORM_MEAN_RGB = new float[]{0.485F, 0.456F, 0.406F}; public static float[] TORCHVISION_NORM_STD_RGB = new float[]{0.229F, 0.224F, 0.225F};
即使用了事先固定的均值和标准差,而不是传入数据的均值和标准差,所以不一定可以得到符合标准正态分布的数据。但是这不重要,因为我要的是直接不作Z-score归一化,只/255就可以了,于是我自定义了一个值为0的均值数组,和值为1的标准差数组,然后传入这个函数,就保证了结果相当于没有做Z-score归一化。原因是Z-score归一化公式如下:
x* = ( x − μ ) / σ
我的最终关键代码如下:注意处理结果的部分,因为我是图片中一定只有0或1个目标检测框,所以我没有使用NMS(非极大值抑制)算法。如果你的图片中有多个检测框,则必须用NMS。我只有两个类别,所以idcnt计算是score.length/7,也就是score.length/(4+1+类别数)。
model = Module.load(path);float[] TORCHVISION_NORM_MEAN_RGB = new float[]{0F, 0F, 0F};float[] TORCHVISION_NORM_STD_RGB = new float[]{1F, 1F, 1F};Tensor inputTensor = TensorImageUtils.bitmapToFloat32Tensor(newBitmap, TORCHVISION_NORM_MEAN_RGB, TORCHVISION_NORM_STD_RGB);// running the modelIValue value = IValue.from(inputTensor);Tensor outputTensor_ori = model.forward(value).toTuple()[0].toTensor();// getting tensor content as java array of floatsfloat[] scores = outputTensor_ori.getDataAsFloatArray();// searching for the index with maximum scorefloat maxScore = 0.85F;int maxScoreIdx = -1;int idcnt = scores.length / 7;for (int i = 0; i < idcnt; i++) {int exist = i*7+4;int j = exist+1+type;if (scores[exist] > 0.25F && scores[j] > maxScore) {maxScore = scores[j];maxScoreIdx = i;}}if (maxScoreIdx == -1) {return false;}float tx = scores[maxScoreIdx*7];float ty = scores[maxScoreIdx*7+1];float tw = scores[maxScoreIdx*7+2];float th = scores[maxScoreIdx*7+3];float ltx = (tx-tw/2);float lty = (ty-th/2);float rbx = (tx+tw/2);float rby = (ty+th/2);drawROI(newBitmap, (int)ltx, (int)lty, (int)rbx, (int)rby);