Kotlin Android 第28天，从 0 到 ML - TensorFlow Lite -姿态估计 (Pose estimation)

前言:

   常常看到特效电影幕後花絮，都有请演员在绿幕前在录动作身上都有点点，在萤幕上变成火材人，姿态估计是使用 ML 模型通过估计身体关键关节（关键点）的空间位置，从图像或影片中估计人的姿势的任务。

大纲 :

来看一下的姿态估计codelab的实作吧

范例提供了两个 TensorFlow Lite 姿态估计模型的参考实现：

MoveNet：最先进的姿势估计模型有两种版本：Lighting 和 Thunder。
PoseNet：2017 年发布的上一代姿态估计模型。

MoveNet 有两种版本：
MoveNet.Lightning 比 Thunder 版本更小、更快但准确度较低。它可以在现代智能手机上实时运行。
MoveNet.Thunder 是更准确的版本，但也比 Lightning 更大更慢。它对於需要更高准确性的用例很有用。

MoveNet 在各种数据集上的表现都优於 PoseNet，尤其是在带有健身动作图像的图像中。因此，我们建议在 PoseNet 上使用 MoveNet。

将 TensorFlow Lite 模型添加到assets文件夹

 movenet_lightning.tflite
 movenet_thunder.tflite
 posenet.tflite

build.gradle(app)

dependencies {
   implementation 'org.tensorflow:tensorflow-lite:2.5.0'
   implementation 'org.tensorflow:tensorflow-lite-gpu:2.5.0'
   implementation 'org.tensorflow:tensorflow-lite-support:0.2.0'
}

**姿态估计 (Pose estimation) 是显示在 camrea 的预览画面，就看一下camrea + surfaceView 重点部份吧
**

camrea 的 imageReader

imageReader =
   ImageReader.newInstance(PREVIEW_WIDTH, PREVIEW_HEIGHT,ImageFormat.YUV_420_888, 3)
    imageReader?.setOnImageAvailableListener({ reader ->
        val image = reader.acquireLatestImage()
        if (image != null) {
            if (!::imageBitmap.isInitialized) {
                imageBitmap =
                    Bitmap.createBitmap(
                        PREVIEW_WIDTH,
                        PREVIEW_HEIGHT,
                        Bitmap.Config.ARGB_8888
                    )
            }
            yuvConverter.yuvToRgb(image, imageBitmap)
            // Create rotated version for portrait display
            val rotateMatrix = Matrix()
            rotateMatrix.postRotate(90.0f)

            val rotatedBitmap = Bitmap.createBitmap(
                imageBitmap, 0, 0, PREVIEW_WIDTH, PREVIEW_HEIGHT,
                rotateMatrix, false
            )
            processImage(rotatedBitmap)
            image.close()
        }
    }, imageReaderHandler)


// process image
private fun processImage(bitmap: Bitmap) {
    var person: Person? = null
    var classificationResult: List<Pair<String, Float>>? = null

   //姿态估计模组
    synchronized(lock) {
        detector?.estimateSinglePose(bitmap)?.let {
            person = it
            classifier?.run {
                classificationResult = classify(person)
            }
        }
    }
    frameProcessedInOneSecondInterval++
    if (frameProcessedInOneSecondInterval == 1) {
        // send fps to view
        listener?.onFPSListener(framesPerSecond)
    }

    //姿态估计模组结果
    listener?.onDetectedInfo(person?.score, classificationResult)
    person?.let {
        visualize(it, bitmap)
    }
}

//依姿态估计模组结果画图
private fun visualize(person: Person, bitmap: Bitmap) {
    var outputBitmap = bitmap

    if (person.score > MIN_CONFIDENCE) {
        outputBitmap = VisualizationUtils.drawBodyKeypoints(bitmap, person)
    }

    val holder = surfaceView.holder
    val surfaceCanvas = holder.lockCanvas()
    surfaceCanvas?.let { canvas ->
        val screenWidth: Int
        val screenHeight: Int
        val left: Int
        val top: Int

        if (canvas.height > canvas.width) {
            val ratio = outputBitmap.height.toFloat() / outputBitmap.width
            screenWidth = canvas.width
            left = 0
            screenHeight = (canvas.width * ratio).toInt()
            top = (canvas.height - screenHeight) / 2
        } else {
            val ratio = outputBitmap.width.toFloat() / outputBitmap.height
            screenHeight = canvas.height
            top = 0
            screenWidth = (canvas.height * ratio).toInt()
            left = (canvas.width - screenWidth) / 2
        }
        val right: Int = left + screenWidth
        val bottom: Int = top + screenHeight

        canvas.drawBitmap(
            outputBitmap, Rect(0, 0, outputBitmap.width, outputBitmap.height),
            Rect(left, top, right, bottom), null
        )
        surfaceView.holder.unlockCanvasAndPost(canvas)
    }
}

执行结果:

手机上是用CPU 使用Lightning 加 PoseNet 计算时间比较短，效果比较明显。