安卓利用tensorflow-lite使用yolov5训练的模型

前言

作为使用yolov5后一次简单的尝试

准备工作

• 通过yolov5训练出自己所需要的模型查看模型训练教程
• 将模型通过tensorflow的python版转换，使用yolov5 6.1以上版本
• 安卓端引入tensorflow远端依赖，并置入模型文件在项目工程里
  这里我并没有去看tensorflow的api，而是直接参考了yolov5-android

注意事项

• 模型我全部用的是demo默认模型
• 运行demo时请打开悬浮窗权限

配置gradle

需要执行cmake来生成通过jni调用的资源
在app目录下的CmakeList.txt文本可以查看具体配置
这里不过多描述，只描述构建过程
远程引入tensorflow依赖

    externalNativeBuild {
        cmake {
            path "CMakeLists.txt"
        }
    }
      defaultConfig {......
        ndk {
            abiFilters 'armeabi-v7a', 'arm64-v8a'
        }
    }
    dependencies {......
	 implementation 'org.tensorflow:tensorflow-lite:2.4.0'
 }

迁移文件

将在modle_tflite下的文件夹复制到如下图所示
将在main/cpp的文件夹也复制到如下图所示

postprocess.cpp文件提供了供java调用的检测函数，这里可以去参考Tensorflow官方网站，是将结果通过jni给回调至客户端

编写悬浮窗

使用了XToast,我是以library形式引入工程的

执行录制屏幕代码

使用的是官方api，MediaProjection，感兴趣可以去查找相关资料
注意在安卓10级以上需要启动一个前台服务才可以使用
开启录屏部分代码

	// 顶部初始化
    private var mediaProjection: MediaProjection?=null
    private lateinit var projectionManager:MediaProjectionManager
    ......
    // 点击事件请求开始录屏
      binding.btnRecord.setOnClickListener {
             projectionManager =getSystemService(MEDIA_PROJECTION_SERVICE) as MediaProjectionManager
            val captureIntent= projectionManager.createScreenCaptureIntent();
            startActivityForResult(captureIntent,20);
        }
    ......
    // 和运行时权限申请一样，在onActivityResult回调里处理开始录屏事件
     override fun onActivityResult(requestCode: Int, resultCode: Int, data: Intent?) {
        super.onActivityResult(requestCode, resultCode, data)
        if (requestCode == 20 && resultCode == RESULT_OK) {
           FloatUtil.showFloatWindow(application).apply {
               contentImageView = this
           }
            mediaProjection = projectionManager.getMediaProjection(resultCode, data!!);
            //执行查看视图函数
            seeContent()
        }
    }    

开启录屏后，需要用到另外一个东西，就是mediaProjection的createVirtualDisplay函数，其实就是VirtualDisplay,创建一个虚拟显示器。mediaProjection的使用(转载)和关于VirtualDisplay的使用

    private fun seeContent() {
        mediaProjection?.apply {
            createDetector()//创建检测器
            createImageReader() //创建imageReader
            registerCallback(object :MediaProjection.Callback(){
                override fun onStop() {
                    super.onStop()
                }
            },handler)
			//使用createVirtuaDisplay函数，需要制定名称，宽和高，dpi，显示模式，surface，剩余两个参数可以为null
            dispalyD = createVirtualDisplay("ScreenImageReader",640,640,1000
            , DisplayManager.VIRTUAL_DISPLAY_FLAG_AUTO_MIRROR,imageReader.surface,object :VirtualDisplay.Callback(){
                    override fun onResumed() {
                        super.onResumed()
                    }

                    override fun onPaused() {
                        super.onPaused()
                    }

                    override fun onStopped() {
                        super.onStopped()
                    }
                },handler)
        }
    }

这里这里使用imageReader,imageReader它是包含一个surface的，并且它有一个setOnImageAvailableListener可以进行监听视图刷新，在刷新的时候去获取image对象，关于图片的边框添加和识别操作，可以看下面代码片段

  /**
     * 图片更新
     * 此处为回调函数
     * 交由handler处理
     */
    override fun onImageAvailable(p0: ImageReader?) {
        //利用handler
            p0?.apply {
                val message = Message()
                message.what = 0xdd
                message.obj = this
                handlerDelayImage.sendMessage(message)
                handlerDelayImage.obtainMessage()
            }
    }
   /**
     * 开始截屏并识别
     * 这个才是具体的处理函数
     */
    private fun startCacheAndSetInput(p0: ImageReader) {
        var nowImage:Image?=null
        try {
            nowImage = p0.acquireLatestImage()

        }catch (e:Exception){
            nowImage = p0.acquireLatestImage()
        }
        nowImage?.apply {
            val width = 640 //可选
            val height = 640
            val planes = planes
            val buffer: ByteBuffer = planes[0].buffer
            val pixelStride = planes[0].pixelStride
            val rowStride = planes[0].rowStride
            val rowPadding = rowStride - pixelStride * width
            var bitmap: Bitmap =
                Bitmap.createBitmap(width + rowPadding / pixelStride, height, Bitmap.Config.ARGB_8888)
            bitmap.copyPixelsFromBuffer(buffer)
            //这就是初始截图
            bitmap = Bitmap.createBitmap(bitmap, 0, 0, width, height)
            //开始检测
            detector.setInput(bitmap)
            //获取到检测结果
            val bboxes: List<TfliteRunner.Recognition> = detector.runInference()
            //绘制识别框到bitmap里
            val resBitmap: Bitmap =ImageProcess.drawBboxes(bboxes, bitmap, 640)
            if (bboxes.size>0){
                //对人点击
                if (bboxes[bboxes.lastIndex].title=="person"){
                    if (!personClick){
                        //检测到人之后对人进行点击
                        //ScreenUtil.click(bboxes[bboxes.lastIndex].location.centerX(), bboxes[bboxes.lastIndex].location.centerY())
                        personClick = true //防止重复点击
                    }
                }else{
                    personClick = false
                }
            }else{
                personClick = false
            }

            runOnUiThread {
                contentImageView?.apply {
                    setImageBitmap(resBitmap)
                }

            }
            close()
        }


    }

关键使用类

悬浮窗 -FloatUtil
模型检测-TfliteRunner
首页-MainActitvity
边框绘制-ImageProcess

使用效果

最后github地址
yolov5结合tensflow在移动端的方案