U2Net 模型在 android 中的使用

Question

我按照这些instructructions把原来的u2net模型权重文件u2net.pth转成了tensorflow lite，转成功了

但是我在 tensrflow lite 的 android 中使用它时遇到了问题，我无法使用 tflite-support script 将图像分割器元数据添加到这个模型中，所以我更改了模型并仅返回 1 个输出 d0（它是所有的组合，即 d1、d2、...、d7）。然后元数据添加成功，我可以使用该模型，但它没有提供任何输出并返回相同的图像。

所以任何帮助将不胜感激，让我知道我在哪里搞砸了，以及我如何使用它在带有 android 的 tensorflow lite 中正确使用这个 u2net 模型，提前致谢..

Answer 1

我会在这里写一个很长的答案。与 U2Net 的 github 存储库取得联系会让您努力检查预处理和 post 处理步骤，以便您可以在 android 项目中应用相同的方法。

首先预处理：在 u2net_test.py 文件中，您可以在 this line 处看到所有图像都使用函数 ToTensorLab(flag=0) 进行了预处理。导航到这里，您会看到 flag=0 的预处理是这样的：

else: # with rgb color (flag = 0)
            tmpImg = np.zeros((image.shape[0],image.shape[1],3))
            image = image/np.max(image)
            if image.shape[2]==1:
                tmpImg[:,:,0] = (image[:,:,0]-0.485)/0.229
                tmpImg[:,:,1] = (image[:,:,0]-0.485)/0.229
                tmpImg[:,:,2] = (image[:,:,0]-0.485)/0.229
            else:
                tmpImg[:,:,0] = (image[:,:,0]-0.485)/0.229
                tmpImg[:,:,1] = (image[:,:,1]-0.456)/0.224
                tmpImg[:,:,2] = (image[:,:,2]-0.406)/0.225

注意2个步骤。

首先将每个颜色像素值除以所有颜色像素值的最大值：

image = image/np.max(image)

和

其次在每个颜色像素值上应用均值和标准：

tmpImg[:,:,0] = (image[:,:,0]-0.485)/0.229
tmpImg[:,:,1] = (image[:,:,1]-0.456)/0.224
tmpImg[:,:,2] = (image[:,:,2]-0.406)/0.225

所以基本上在 Kotlin 中，如果你有一个位图，你必须做类似的事情：

fun bitmapToFloatArray(bitmap: Bitmap):
                Array<Array<Array<FloatArray>>> {
            
            val width: Int = bitmap.width
            val height: Int = bitmap.height
            val intValues = IntArray(width * height)
            bitmap.getPixels(intValues, 0, width, 0, 0, width, height)

            // Create aa array to find the maximum value
            val fourDimensionalArray = Array(1) {
                Array(320) {
                    Array(320) {
                        FloatArray(3)
                    }
                }
            }
            // https://github.com/xuebinqin/U-2-Net/blob/f2b8e4ac1c4fbe90daba8707bca051a0ec830bf6/data_loader.py#L204
            for (i in 0 until width - 1) {
                for (j in 0 until height - 1) {
                    val pixelValue: Int = intValues[i * width + j]
                    fourDimensionalArray[0][i][j][0] =
                        Color.red(pixelValue)
                            .toFloat()
                    fourDimensionalArray[0][i][j][1] =
                        Color.green(pixelValue)
                            .toFloat()
                    fourDimensionalArray[0][i][j][2] =
                        Color.blue(pixelValue).toFloat()
                }

            }
            // Convert multidimensional array to 1D
            val oneDFloatArray = ArrayList<Float>()

            for (m in fourDimensionalArray[0].indices) {
                for (x in fourDimensionalArray[0][0].indices) {
                    for (y in fourDimensionalArray[0][0][0].indices) {
                        oneDFloatArray.add(fourDimensionalArray[0][m][x][y])
                    }
                }
            }

            val maxValue: Float = oneDFloatArray.maxOrNull() ?: 0f
            //val minValue: Float = oneDFloatArray.minOrNull() ?: 0f

            // Final array that is going to be used with interpreter
            val finalFourDimensionalArray = Array(1) {
                Array(320) {
                    Array(320) {
                        FloatArray(3)
                    }
                }
            }
            for (i in 0 until width - 1) {
                for (j in 0 until height - 1) {
                    val pixelValue: Int = intValues[i * width + j]
                    finalFourDimensionalArray[0][i][j][0] =
                        ((Color.red(pixelValue).toFloat() / maxValue) - 0.485f) / 0.229f
                    finalFourDimensionalArray[0][i][j][1] =
                        ((Color.green(pixelValue).toFloat() / maxValue) - 0.456f) / 0.224f
                    finalFourDimensionalArray[0][i][j][2] =
                        ((Color.blue(pixelValue).toFloat() / maxValue) - 0.406f) / 0.225f
                }

            }

            return finalFourDimensionalArray
        }

然后这个数组被送入解释器，因为你的模型有多个输出，我们使用 runForMultipleInputsOutputs:

// Convert Bitmap to Float array
             val inputStyle = ImageUtils.bitmapToFloatArray(loadedBitmap)

            // Create arrays with size 1,320,320,1
            val output1 =  Array(1) { Array(CONTENT_IMAGE_SIZE) { Array(CONTENT_IMAGE_SIZE) { FloatArray(1)}}}
            val output2 =  Array(1) { Array(CONTENT_IMAGE_SIZE) { Array(CONTENT_IMAGE_SIZE) { FloatArray(1)}}}
            val output3 =  Array(1) { Array(CONTENT_IMAGE_SIZE) { Array(CONTENT_IMAGE_SIZE) { FloatArray(1)}}}
            val output4 =  Array(1) { Array(CONTENT_IMAGE_SIZE) { Array(CONTENT_IMAGE_SIZE) { FloatArray(1)}}}
            val output5 =  Array(1) { Array(CONTENT_IMAGE_SIZE) { Array(CONTENT_IMAGE_SIZE) { FloatArray(1)}}}
            val output6 =  Array(1) { Array(CONTENT_IMAGE_SIZE) { Array(CONTENT_IMAGE_SIZE) { FloatArray(1)}}}

            val outputs: MutableMap<Int,
                    Any> = HashMap()
            outputs[0] = output1
            outputs[1] = output2
            outputs[2] = output3
            outputs[3] = output4
            outputs[4] = output5
            outputs[5] = output6
          
            // Runs model inference and gets result.
            val array = arrayOf(inputStyle)
            interpreterDepth.runForMultipleInputsOutputs(array, outputs)

然后我们使用解释器的第一个输出作为you can see at u2net_test.py file. (I have also printed results of line 112但是好像没有效果。您可以自由尝试使用颜色像素值的最小值和最大值）。所以我们有 post proseccing 就像你在 save_output function:

看到的那样

// Convert output array to Bitmap
val (finalBitmapGrey, finalBitmapBlack) = ImageUtils.convertArrayToBitmapTensorFlow(
                output1, CONTENT_IMAGE_SIZE,
                CONTENT_IMAGE_SIZE
            )

上面的函数是这样的：

fun convertArrayToBitmapTensorFlow(
            imageArray: Array<Array<Array<FloatArray>>>,
            imageWidth: Int,
            imageHeight: Int
        ): Bitmap {
            val conf = Bitmap.Config.ARGB_8888 // see other conf types
            val grayToneImage = Bitmap.createBitmap(imageWidth, imageHeight, conf)

            for (x in imageArray[0].indices) {
                for (y in imageArray[0][0].indices) {
                    val color = Color.rgb(
                        //
                        (((imageArray[0][x][y][0]) * 255f).toInt()),
                        (((imageArray[0][x][y][0]) * 255f).toInt()),
                        (((imageArray[0][x][y][0]) * 255f).toInt())
                    )

                    // this y, x is in the correct order!!!
                    grayToneImage.setPixel(y, x, color)
                }
            }
            return grayToneImage
        }

然后这张灰度图你可以随意使用。

由于预处理的多个步骤，我直接使用了解释器，没有额外的库。如果您可以通过所有步骤插入元数据，我将在本周晚些时候尝试，但我对此表示怀疑。

如果您需要一些说明，请随时问我。

Colab 笔记本 link

编码愉快

U2Net 模型在 android 中的使用

Usage of U2Net Model in android

android

pytorch

image-segmentation

tensorflow-lite