TensorFlow Lite + GPUImage 实现AI背景虚化（一）

发表于 | 分类于 | | 阅读次数:
如何用TensorFlow Lite结合GPUImage在Android中实现AI背景虚化

最近项目中有研究背景虚化功能，需求是通过写一个GPUImage的滤镜，结合TensorFLow Lite来实现对图片中指定物体的背景虚化功能。这部分内容基本都是通过看官方文档和自己摸索学习，这里总结并整理一份笔记，内容主要包括Android接入TensorFlow Lite，通过运行AI模型来识别图片中物体，并对其做背景虚化。一共分为三部分，本文是第一部分，TensorFlow Lite的接入及模型的使用。

什么是Tensorflow Lite

Tensorflow Lite是一个Google官方的机器学习库，主要功能就帮助我们在移动端运行Tensorflow Lite模型。TensorFlow Lite主要由两个组件构成，详见官网介绍：

TensorFlow Lite consists of two main components:

The TensorFlow Lite interpreter, which runs specially optimized models on many different hardware types, including mobile phones, embedded Linux devices, and microcontrollers.
The TensorFlow Lite converter, which converts TensorFlow models into an efficient form for use by the interpreter, and can introduce optimizations to improve binary size and performance.

这两个组件，一个是转换器，一个是解释器。

转换器顾名思义，就是要把常规的Tensorflow模型转换为Android端适用的模型文件。而解释器则用来在Android端具体去导入并运行模型。我们开发时要做的就是在项目中接入Tensorflow Lite库，然后用库中提供的解释器的一些API来导入AI模型文件，运行模型，解析得到的数据。

Tensorflite的接入

Google为我们提供了一个demo项目，里面包括了图像识别、图像分割等不同模型的使用示例，可以直接参考：Tensorflow Lite Android使用demo

首先，要准备好Tensorflow lite模型文件，在Android端是以一种后缀为 .tflite 的文件。可以直接用demo中的示例，或者自己用转换器生成的tflite格式的模型文件。

接下来在项目module下的的 build.gradle 文件中添加如下依赖:

1
2
3
4
//tensorflow lite核心库
implementation "org.tensorflow:tensorflow-lite:2.1.0"
//tensorflow lite辅助库（可以不用添加）
implementation "org.tensorflow:tensorflow-lite-support:0.0.0-nightly"

这样已经可用了，不过Google推荐我们再对依赖的ndk做一些精简，继续在 build.gradle 文件中加入以下配置：

1
2
3
4
5
6
7
android {
    defaultConfig {
        ndk {
            abiFilters 'armeabi-v7a', 'arm64-v8a'
        }
    }
}

最后，我们将tflite文件（这里命名为model.tflite），将这个文件移动到项目model下的 src/main/assets 目录下，如图：

ok，接下来就是在代码中娶导入并使用模型文件了。

具体使用

导入模型后，具体的使用方法大致有以下几步：

  1. 将输入数据转换为正确格式的InputBuffer
  2. 创建好准备接收运行结果的OutputBuffer
  3. 创建模型对应的解释器并运行模型
  4. 处理运行模型得到的输出数据

Tensorflow Lite使用的输入输出数据以Buffer承载，这里可以简单了解下：Java NIO 之 Buffer(缓冲区) 。接下来以图像分割模型举例。

输入数据的处理

在图像分割模型中，我们的输入数据显然是一张图片。图片的大小是不固定的，而一般来说模型需要输入的图片都是给的大小的，因此第一步是要把不同尺寸大小的图片经过缩放或者补充像素等方式转换为给的固定大小的图片，这里其实就是一个Bitmap的处理，这里就不细说了。

处理好Bitmap后，接下来要把Bitmap先转换为一个三维数组。这个三维数组是有图像中每一个像素点的RGB值构成。比如图像的大小是500500，那么这个三维数组大小就是500 500 * 3。代码如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21

/**
 * 将Bitmap转换为一个RGB数组
 * @return 一个三维数组，第三维是像素点的RGB值
 */
fun Bitmap.rgbArray(): Array<Array<FloatArray>> {
    val pixelsValues = IntArray(this.width * this.height)
    this.getPixels(pixelsValues, 0, this.width, 0, 0, this.width, this.height)
    val result = Array(this.height) { Array(this.width) { FloatArray(3) } }

    var pixel = 0
    for (y in 0 until this.height) {
        for (x in 0 until this.width) {
            val value = pixelsValues[pixel++]   
            result[y][x][0] = ((value shr 16 and 0xFF).toFloat())//r
            result[y][x][1] = ((value shr 8 and 0xFF).toFloat())//g
            result[y][x][2] = ((value and 0xFF).toFloat())//b
        }
    }
    return result
}

拿到这个三维数组后，要将其转换为一个ByteBuffer。因为前面三维数组的大小是 500 500 3，所以对应的FloatBuffer大小是 500 500 3,因为一个Java中 1 Float = 4 Byte，所以对应ByteBuffer，大小就是 500 500 3 * 4。另外一般来说模型的输入数据都需要一个归一化的预处理，其实也就是对数值做一个简单的加减乘除，代码如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

/**
     * 将rgb数组转换为模型运行所需要的byteBuffer
     * @param array rgb数组
     * @param normalizeOp 归一化操作
     */
    private fun rgbToByteBuffer(array: Array<Array<FloatArray>>, width: Int, height: Int, normalizeOp: ((Float) -> Float) = { it }): ByteBuffer {
        val inputImage = ByteBuffer.allocateDirect(1 * width * height * 3 * 4)
        inputImage.order(ByteOrder.nativeOrder()).rewind()
        for (y in 0 until height) {
            for (x in 0 until width) {
                for (z in listOf(0, 1, 2)) {
                    val value = array[y][x][z]
                    inputImage.putFloat(normalizeOp(value))
                }
            }
        }
        inputImage.rewind()
        return inputImage
    }
    

    val input = rgbToByteBuffer(rgbArray, 500, 500) { (it * 2 / 255F) - 1 }

转换为ByteBuffer后就OK了，这就是我们需要的输入数据。

准备输出数据

因为输入数据和输出数据同样是ByteBuffer，只是大小不一样。输出数据的Buffer大小根据模型来定。比如在这个图像分割模型中，输出数据应该是一个Long类型的二维数组，大小是 500 500。那么对应的LongBuffer大小就是 500 500。又因为 1 Long = 8 Byte，所以最终的输出Buffer应该是一个大小为 500 500 8 的ByteBuffer，代码如下：

1
2

val output = ByteBuffer.allocateDirect(TF_DPI_SEG * TF_DPI_SEG * 8).order(ByteOrder.nativeOrder()).asLongBuffer()

创建解释器并运行

Tensorflow Lite解释器即 Interpreter ，创建方式很简单，代码如下：

1
2
3
4

val byteBuffer = FileUtil.loadMappedFile(context, "model.tflite")
val options = Interpreter.Options()
val interpreter = Interpreter(byteBuffer,options)

直接new一个对象就好了，构造器中需要两个参数，第一个将模型文件转为ByteBuffer，第二个则是当前需要的一些配置Options。模型文件向ByteBuffer的转换可以直接用support库中提供的工具类来处理。

创建好解释器对象后，直接调用他的 run() 方法并传入前面准备好的输入、输出数据即可：

1
2
3
//...
interpreter.run(input, output)
//...

输出数据的处理

最后就是输出数据的处理了。前面说过，我们的目标是把找出图片中的目标物体对其做背景虚化，所以需要用原图，通过模型运算来生成一张标记出或者描出目标物体的图，这张图中目标物体应该是一种颜色，非目标物体是另一种颜色，这样方便后面的处理。因此在得到输出数据ByteBuffer后，需要先转换为一个常规的一维或者二维数组，再根据数组生成一张Bitmap。代码如下：

1
2
3
4
//...
var array = IntArray(output.limit()) {
    output[it].toInt()
}

上面的代码把输出数据转换为了一个Int数组，数组中的每一个值都表示对应像素点的物体类型，比如1是人，2是狗，3是猫等等。接下来就是用数组生成Bitmap了，代码如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

/**
 * 用一个RGB数组生成bitmap
 * @param pixelsOp 数组中每一个值转换为具体像素值的过程  */
fun <T> createMask(array: Array<T>, width: Int, height: Int, pixelsOp: (T) -> Int): Bitmap {
    val conf = Bitmap.Config.ARGB_8888
    val maskBitmap = Bitmap.createBitmap(width, height, conf)

    for (y in 0 until height) {
        for (x in 0 until width) {
            val value = array[x + width * y]
            val color = pixelsOp(value)
            maskBitmap.setPixel(x, y, color)
        }
    }

    return maskBitmap
}

val mask = createMask(array.toTypedArray(), TF_DPI_SEG, TF_DPI_SEG) { value ->
            Color.argb(
                255,
                //根据实际情况计算
                , 0, 0
            )
        }

根据实际情况，比如目标物体是猫，那么属于猫的部分就是黑色，非猫的部分是红色，那么这样最终输出的Bitmap就是：

原图：

结果图：