使用 LiteRT 在 Android 上实现实时加速图像分割

1. 准备工作

输入代码是培养肌肉记忆和加深对材料理解的好方法。虽然复制粘贴可以节省时间,但从长远来看,投资于这种实践可以提高效率并增强编码技能。

在此 Codelab 中,您将学习如何使用 Google 的新 TensorFlow Lite 运行时 LiteRT 构建一个 Android 应用,该应用可对实时摄像头画面执行实时图像分割。您将使用一个初始 Android 应用,并为其添加图像分割功能。我们还将介绍预处理、推理和后处理步骤。您将学习以下内容:

  • 构建一个可实时分割图片的 Android 应用。
  • 集成预训练的 LiteRT 图片分割模型。
  • 预处理模型的输入图片。
  • 使用 LiteRT 运行时进行 CPU 和 GPU 加速。
  • 了解如何处理模型的输出以显示分割掩码。
  • 了解如何针对前置摄像头进行调整。

最后,您将创建与下图类似的内容:

完成后的应用

前提条件

本 Codelab 专为想要获得机器学习经验的有经验移动开发者而设计。您应熟悉以下内容:

  • 使用 Kotlin 和 Android Studio 进行 Android 开发
  • 图像处理的基本概念

学习内容

  • 如何在 Android 应用中集成和使用 LiteRT 运行时。
  • 如何使用预训练的 LiteRT 模型执行图像分割。
  • 如何为模型预处理输入图片。
  • 如何针对模型运行推理。
  • 如何处理分割模型的输出以直观呈现结果。
  • 如何使用 CameraX 进行实时摄像头画面处理。

所需条件

  • 最新版本的 Android Studio(已在 v2025.1.1 上测试)。
  • 实体 Android 设备。最好在 Galaxy 和 Pixel 设备上进行测试。
  • 示例代码(来自 GitHub)。
  • 具备使用 Kotlin 进行 Android 开发的基础知识。

2. 图像分割

图像分割是一项计算机视觉任务,涉及将图像划分为多个段或区域。与在对象周围绘制边界框的对象检测不同,图片分割会为图片中的每个像素分配特定的类别或标签。这样一来,您就可以更详细、更精细地了解图片的内容,从而确切知道每个对象的形状和边界。

例如,您不仅可以知道某个方框中有人,还可以准确知道哪些像素属于该人。本教程演示了如何使用预训练的机器学习模型在 Android 设备上执行实时图像分割。

细分示例

LiteRT:推动设备端机器学习的边缘化

LiteRT 是一项关键技术,可在移动设备上实现实时、高保真分割。LiteRT 是 Google 为 TensorFlow Lite 打造的下一代高性能运行时,旨在充分发挥底层硬件的性能。

它通过智能且优化地使用 GPU(图形处理单元)和 NPU(神经处理单元)等硬件加速器来实现这一目标。通过将分割模型的高强度计算工作负载从通用 CPU 分流到这些专用处理器,LiteRT 可显著缩短推理时间。这种加速功能使我们能够在实时摄像头画面上流畅运行复杂的模型,从而拓展了直接在手机上使用机器学习技术所能实现的范围。如果没有这种性能,实时分割功能就会太慢且不稳定,无法提供良好的用户体验。

3. 进行设置

克隆代码库

首先,克隆 LiteRT 的代码库:

git clone https://github.com/google-ai-edge/LiteRT.git

LiteRT/litert/samples/image_segmentation 是包含您需要的所有资源的目录。在此 Codelab 中,您只需要 kotlin_cpu_gpu/android_starter 项目。如果您遇到问题,不妨查看完成的项目:kotlin_cpu_gpu/android

关于文件路径的说明

本教程以 Linux/macOS 格式指定文件路径。如果您使用的是 Windows,则需要相应地调整路径。

另请务必注意 Android Studio 项目视图与标准文件系统视图之间的区别。Android Studio 项目视图是项目文件的结构化表示形式,专为 Android 开发而整理。本教程中的文件路径是指文件系统路径,而不是 Android Studio 项目视图中的路径。

导入 starter 应用

首先,将 starter 应用导入 Android Studio。

  1. 打开 Android Studio,然后选择 Open

Android Studio Open

  1. 前往 kotlin_cpu_gpu/android_starter 目录并将其打开。

Android 启动器

为确保您的应用能使用所有的依赖项,在导入过程完成后,应将您的项目与 Gradle 文件同步。

  1. 在 Android Studio 工具栏中选择 Sync Project with Gradle Files

菜单同步

  1. 请勿跳过此步骤,否则您将无法理解本教程的其余内容。

运行起始应用

现在,您已将项目导入 Android Studio,接下来就可以首次运行该应用了。

通过 USB 线将 Android 设备连接到计算机,然后点击 Android Studio 工具栏中的 Run

“运行”按钮

应用应在设备上启动。您会看到实时摄像头画面,但尚未进行分割。您在本教程中进行的所有文件编辑都将在 LiteRT/litert/samples/image_segmentation/kotlin_cpu_gpu/android_starter/app/src/main/java/com/google/aiedge/examples/image_segmentation 目录下进行(现在您知道 Android Studio 为何要重构此目录了 😃)。

项目目录

您还会在 ImageSegmentationHelper.ktMainViewModel.ktview/SegmentationOverlay.kt 文件中看到 TODO 注释。在以下步骤中,您将通过填充这些 TODO 来实现图片分割功能。

4. 了解 starter 应用

起始应用已具有基本的界面和相机处理逻辑。以下是关键文件的简要概述:

  • app/src/main/java/com/google/aiedge/examples/image_segmentation/MainActivity.kt:这是应用的主要入口点。它使用 Jetpack Compose 设置界面并处理相机权限。
  • app/src/main/java/com/google/aiedge/examples/image_segmentation/MainViewModel.kt:此 ViewModel 用于管理界面状态并协调图像分割流程。
  • app/src/main/java/com/google/aiedge/examples/image_segmentation/ImageSegmentationHelper.kt:我们将在其中添加图片分割的核心逻辑。它将负责加载模型、处理相机帧和运行推理。
  • app/src/main/java/com/google/aiedge/examples/image_segmentation/view/CameraScreen.kt:此可组合函数用于显示相机预览和分割叠加层。
  • app/src/main/assets/selfie_multiclass.tflite:这是我们将使用的预训练 TensorFlow Lite 图像分割模型。

5. 了解 LiteRT 并添加依赖项

现在,我们向初始应用添加图片分割功能。

1. 添加 LiteRT 依赖项

首先,您必须将 LiteRT 库添加到项目中。这是关键的第一步,可让您通过 Google 的优化运行时在设备上启用机器学习。

打开 app/build.gradle.kts 文件,并将以下行添加到 dependencies 代码块中:

// LiteRT for on-device ML
implementation(libs.litert)

添加依赖项后,点击 Android Studio 右上角显示的 Sync Now 按钮,将项目与 Gradle 文件同步。

立即同步

2. 了解 Key LiteRT API

打开ImageSegmentationHelper.kt

在编写实现代码之前,请务必了解您将使用的 LiteRT API 的核心组件。确保您是从 com.google.ai.edge.litert 软件包导入,并将以下导入项添加到 ImageSegmentationHelper.kt 的顶部:

import com.google.ai.edge.litert.Accelerator
import com.google.ai.edge.litert.CompiledModel
  • CompiledModel:这是与 TFLite 模型交互的核心类。它表示已针对特定硬件加速器(如 CPU 或 GPU)预编译和优化的模型。这种预编译是 LiteRT 的一项关键功能,可实现更快、更高效的推理。
  • CompiledModel.Options:您可以使用此 build 类来配置 CompiledModel。最重要的设置是指定要用于运行模型的硬件加速器。
  • Accelerator:此枚举可让您选择用于推理的硬件。新手入门项目已配置为处理以下选项:
    • Accelerator.CPU:用于在设备的 CPU 上运行模型。这是最通用的选项。
    • Accelerator.GPU:用于在设备的 GPU 上运行模型。对于基于图像的模型,这通常比 CPU 快得多。
  • 输入和输出缓冲区 (TensorBuffer):LiteRT 使用 TensorBuffer 作为模型输入和输出。这样,您就可以对内存进行精细控制,并避免不必要的数据复制。您将使用 model.createInputBuffers()model.createOutputBuffers() 直接从 CompiledModel 实例获取这些缓冲区,然后将输入数据写入这些缓冲区,并从中读取结果。
  • model.run():此函数用于执行推理。您只需将输入和输出缓冲区传递给它,LiteRT 即可处理在所选硬件加速器上运行模型的复杂任务。

6. 完成 ImageSegmentationHelper 的初始实现

现在,该编写一些代码了。您将完成 ImageSegmentationHelper.kt 的初始实现。这包括设置 Segmenter 私有类来保存 LiteRT 模型,并实现 cleanup() 函数以正确释放该模型。

  1. 完成 Segmenter 类和 cleanup() 函数:在 ImageSegmentationHelper.kt 文件中,您会找到一个名为 Segmenter 的私有类和一个名为 cleanup() 的函数的框架。首先,完成 Segmenter 类,具体做法是定义其构造函数以保存模型,为输入/输出缓冲区创建属性,并添加 close() 方法以释放模型。然后,实现 cleanup() 函数以调用此新的 close() 方法。将现有的 Segmenter 类和 cleanup() 函数替换为以下内容:(~第 83 行)
    private class Segmenter(
        // Add this argument
        private val model: CompiledModel,
        private val coloredLabels: List<ColoredLabel>,
    ) {
        // Add these private vals
        private val inputBuffers: = model.createInputBuffers()
        private val outputBuffers: = model.createOutputBuffers()
    
        fun cleanup() {
          // cleanup buffers
          inputBuffers.forEach { it.close() }
          outputBuffers.forEach { it.close() }
          // cleanup model
          model.close()
        }
    }
    
  2. 定义 toAccelerator 方法:此方法将加速器菜单中定义的加速器枚举映射到导入的 LiteRT 模块特有的加速器枚举(大约在第 225 行):
    fun toAccelerator(acceleratorEnum: AcceleratorEnum): Accelerator {
      return when (acceleratorEnum) {
        AcceleratorEnum.CPU -> Accelerator.CPU
        AcceleratorEnum.GPU -> Accelerator.GPU
      }
    }
    
  3. 初始化 CompiledModel:现在,找到 initSegmenter 函数。您将在此处创建 CompiledModel 实例,并使用该实例来实例化您现在定义的 Segmenter 类。此代码使用指定的加速器(CPU 或 GPU)设置模型,并准备好进行推理。将 initSegmenter 中的 TODO 替换为以下实现(Cmd/Ctrl+f“initSegmenter”或第 62 行附近):
    cleanup()
    try {
      withContext(singleThreadDispatcher) {
        val model =
          CompiledModel.create(
            context.assets,
            "selfie_multiclass.tflite",
            CompiledModel.Options(toAccelerator(acceleratorEnum)),
            null,
          )
        segmenter = Segmenter(model, coloredLabels)
        Log.d(TAG, "Created an image segmenter")
      }
    } catch (e: Exception) {
      Log.i(TAG, "Create LiteRT from selfie_multiclass is failed: ${e.message}")
      _error.emit(e)
    }
    

7. 开始分段和预处理

现在我们已经有了模型,接下来需要触发细分流程并为模型准备输入数据。

触发细分

分割过程在 MainViewModel.kt 中开始,该进程会从摄像头接收帧。

打开MainViewModel.kt

  1. 从相机帧触发分割MainViewModel 中的 segment 函数是分割任务的入口点。每当相机提供新图片或从媒体库中选择新图片时,系统都会调用这些方法。然后,这些函数会调用 ImageSegmentationHelper 中的 segment 方法。将两个 segment 函数中的 TODO 替换为以下内容(第 107 行左右):
    // For ImageProxy (from CameraX)
    fun segment(imageProxy: ImageProxy) {
        segmentJob =
            viewModelScope.launch {
                imageSegmentationHelper.segment(imageProxy.toBitmap(), imageProxy.imageInfo.rotationDegrees)
                imageProxy.close()
            }
    }
    
    // For Bitmaps (from gallery)
    fun segment(bitmap: Bitmap, rotationDegrees: Int) {
        segmentJob =
            viewModelScope.launch {
                val argbBitmap = bitmap.copy(Bitmap.Config.ARGB_8888, true)
                imageSegmentationHelper.segment(argbBitmap, rotationDegrees)
            }
    }
    

预处理图片

现在,我们回到 ImageSegmentationHelper.kt 来处理图片预处理。

打开ImageSegmentationHelper.kt

  1. 实现公共 segment 函数:此函数充当封装容器,用于调用 Segmenter 类中的私有 segment 函数。将 TODO 替换为(大约第 95 行):
    try {
      withContext(singleThreadDispatcher) {
        segmenter?.segment(bitmap, rotationDegrees)?.let { if (isActive) _segmentation.emit(it) }
      }
    } catch (e: Exception) {
      Log.i(TAG, "Image segment error occurred: ${e.message}")
      _error.emit(e)
    }
    
  2. 实现预处理Segmenter 类中的私有 segment 函数用于对输入图片执行必要的转换,以便为模型做好准备。这包括缩放、旋转和归一化图片。然后,此函数将调用另一个私有 segment 函数来执行推理。将 segment(bitmap: Bitmap, ...) 函数中的 TODO 替换为(大约在第 121 行):
    val totalStartTime = SystemClock.uptimeMillis()
    val rotation = -rotationDegrees / 90
    val (h, w) = Pair(256, 256)
    
    // Preprocessing
    val preprocessStartTime = SystemClock.uptimeMillis()
    var image = bitmap.scale(w, h, true)
    image = rot90Clockwise(image, rotation)
    val inputFloatArray = normalize(image, 127.5f, 127.5f)
    Log.d(TAG, "Preprocessing time: ${SystemClock.uptimeMillis() - preprocessStartTime} ms")
    
    // Inference
    val inferenceStartTime = SystemClock.uptimeMillis()
    val segmentResult = segment(inputFloatArray)
    Log.d(TAG, "Inference time: ${SystemClock.uptimeMillis() - inferenceStartTime} ms")
    
    Log.d(TAG, "Total segmentation time: ${SystemClock.uptimeMillis() - totalStartTime} ms")
    return SegmentationResult(segmentResult, SystemClock.uptimeMillis() - inferenceStartTime)
    

8. 使用 LiteRT 进行主推断

对输入数据进行预处理后,我们现在可以使用 LiteRT 运行核心推理。

打开ImageSegmentationHelper.kt

  1. 实现模型执行:私有 segment(inputFloatArray: FloatArray) 函数是我们直接与 LiteRT run() 方法交互的地方。我们将预处理后的数据写入输入缓冲区,运行模型,然后从输出缓冲区读取结果。将此函数中的 TODO 替换为(大约在第 188 行):
    val (h, w, c) = Triple(256, 256, 6)
    
    // MODEL EXECUTION PHASE
    val modelExecStartTime = SystemClock.uptimeMillis()
    
    // Write input data - measure time
    val bufferWriteStartTime = SystemClock.uptimeMillis()
    inputBuffers[0].writeFloat(inputFloatArray)
    val bufferWriteTime = SystemClock.uptimeMillis() - bufferWriteStartTime
    Log.d(TAG, "Buffer write time: $bufferWriteTime ms")
    
    // Optional tensor inspection
    logTensorStats("Input tensor", inputFloatArray)
    
    // Run model inference - measure time
    val modelRunStartTime = SystemClock.uptimeMillis()
    model.run(inputBuffers, outputBuffers)
    val modelRunTime = SystemClock.uptimeMillis() - modelRunStartTime
    Log.d(TAG, "Model.run() time: $modelRunTime ms")
    
    // Read output data - measure time
    val bufferReadStartTime = SystemClock.uptimeMillis()
    val outputFloatArray = outputBuffers[0].readFloat()
    val outputBuffer = FloatBuffer.wrap(outputFloatArray)
    val bufferReadTime = SystemClock.uptimeMillis() - bufferReadStartTime
    Log.d(TAG, "Buffer read time: $bufferReadTime ms")
    
    val modelExecTime = SystemClock.uptimeMillis() - modelExecStartTime
    Log.d(TAG, "Total model execution time: $modelExecTime ms")
    
    // Optional tensor inspection
    logTensorStats("Output tensor", outputFloatArray)
    
    // POSTPROCESSING PHASE
    val postprocessStartTime = SystemClock.uptimeMillis()
    
    // Process mask from model output
    val inferenceData = InferenceData(width = w, height = h, channels = c, buffer = outputBuffer)
    val mask = processImage(inferenceData)
    
    val postprocessTime = SystemClock.uptimeMillis() - postprocessStartTime
    Log.d(TAG, "Postprocessing time (mask creation): $postprocessTime ms")
    
    return Segmentation(
      listOf(Mask(mask, inferenceData.width, inferenceData.height)),
      coloredLabels,
    )
    

9. 后期处理和显示叠加层

运行推理后,我们会从模型中获得原始输出。我们需要处理此输出,以创建视觉分割掩码,然后将其显示在屏幕上。

打开ImageSegmentationHelper.kt

  1. 实现输出处理processImage 函数将模型的原始浮点输出转换为表示分割掩码的 ByteBuffer。为此,它会找到每个像素的最高概率类别。将其 TODO 替换为(大约在第 238 行):
    val mask = ByteBuffer.allocateDirect(inferenceData.width * inferenceData.height)
    for (i in 0 until inferenceData.height) {
        for (j in 0 until inferenceData.width) {
            val offset = inferenceData.channels * (i * inferenceData.width + j)
    
            var maxIndex = 0
            var maxValue = inferenceData.buffer.get(offset)
    
            for (index in 1 until inferenceData.channels) {
                if (inferenceData.buffer.get(offset + index) > maxValue) {
                    maxValue = inferenceData.buffer.get(offset + index)
                    maxIndex = index
                }
            }
            mask.put(i * inferenceData.width + j, maxIndex.toByte())
        }
    }
    return mask
    

打开MainViewModel.kt

  1. 收集和处理细分结果:现在,我们回到 MainViewModel 来处理来自 ImageSegmentationHelper 的细分结果。segmentationUiShareFlow 会收集 SegmentationResult,将遮罩转换为彩色 Bitmap,并将其提供给界面。将 segmentationUiShareFlow 属性中的 TODO 替换为(大约第 63 行)- 不要替换已有的代码,只需填充正文:
    viewModelScope.launch {
      imageSegmentationHelper.segmentation
        .filter { it.segmentation.masks.isNotEmpty() }
        .map {
          val segmentation = it.segmentation
          val mask = segmentation.masks[0]
          val maskArray = mask.data
          val width = mask.width
          val height = mask.height
          val pixelSize = width * height
          val pixels = IntArray(pixelSize)
    
          val colorLabels =
            segmentation.coloredLabels.mapIndexed { index, coloredLabel ->
              ColorLabel(index, coloredLabel.label, coloredLabel.argb)
            }
          // Set color for pixels
          for (i in 0 until pixelSize) {
            val colorLabel = colorLabels[maskArray[i].toInt()]
            val color = colorLabel.getColor()
            pixels[i] = color
          }
          // Get image info
          val overlayInfo = OverlayInfo(pixels = pixels, width = width, height = height)
    
          val inferenceTime = it.inferenceTime
          Pair(overlayInfo, inferenceTime)
        }
        .collect { flow.emit(it) }
    }
    

打开view/SegmentationOverlay.kt

最后一步是,当用户切换到前置摄像头时,正确调整分割叠加层的方向。前置摄像头的画面自然会进行镜像处理,因此我们需要对叠加层 Bitmap 应用相同的水平翻转,以确保它与摄像头预览正确对齐。

  1. 处理叠加层方向:在 SegmentationOverlay.kt 文件中找到 TODO,并将其替换为以下代码。此代码会检查前置摄像头是否处于活动状态,如果处于活动状态,则在叠加层 Bitmap 绘制到 Canvas 上之前,对其应用水平翻转。(第 42 行左右):
    val orientedBitmap =
      if (lensFacing == CameraSelector.LENS_FACING_FRONT) {
        // Create a matrix for horizontal flipping
        val matrix = Matrix().apply { preScale(-1f, 1f) }
        Bitmap.createBitmap(image, 0, 0, image.width, image.height, matrix, false).also {
          image.recycle()
        }
      } else {
        image
      }
    

10. 运行并使用最终应用

您现在已完成所有必要的代码更改。现在,您可以运行应用,看看自己的工作成果了!

  1. 运行应用:连接 Android 设备,然后点击 Android Studio 工具栏中的 Run

“运行”按钮

  1. 测试功能:应用启动后,您应该会看到带有彩色分割叠加层的实时摄像头 Feed。
    • 切换摄像头:点按顶部的摄像头翻转图标,即可在前置摄像头和后置摄像头之间切换。请注意叠加层如何正确调整方向。
    • 更改加速器:点按底部的“CPU”或“GPU”按钮可切换硬件加速器。观察屏幕底部显示的推理时间的变化。GPU 应该会快得多。
    • 使用图库图片:点按顶部的“图库”标签页,从设备的照片库中选择图片。应用将对所选静态图片运行分割。

其他界面

现在,您已拥有一个由 LiteRT 提供支持且功能完善的实时图像分割应用!

11. 高级(可选):使用 NPU

此代码库还包含一个针对神经处理单元 (NPU) 优化的应用版本。NPU 版本可在具有兼容 NPU 的设备上显著提升性能。

如需试用 NPU 版本,请在 Android Studio 中打开 kotlin_npu/android 项目。该代码与 CPU/GPU 版本非常相似,并且配置为使用 NPU 委托。

如需使用 NPU 委托,您需要加入抢先体验计划

12. 恭喜!

您已成功构建了一个使用 LiteRT 执行实时图像分割的 Android 应用。您已了解如何:

  • 将 LiteRT 运行时集成到 Android 应用中。
  • 加载并运行 TFLite 图像分割模型。
  • 预处理模型的输入。
  • 处理模型输出以创建分割掩码。
  • 使用 CameraX 开发实时相机应用。

后续步骤

  • 尝试使用其他图片分割模型。
  • 尝试不同的 LiteRT 委托(CPU、GPU、NPU)。

了解详情