使用 LiteRT 在 Android 上即時加速圖片區隔

1. 事前準備

輸入程式碼是建立肌肉記憶和加深對教材理解的絕佳方式。複製及貼上程式碼雖然可以節省時間,但長期而言,投資這項做法可提高效率,並培養更強大的程式碼編寫技能。

在本程式碼研究室中,您將瞭解如何建構 Android 應用程式,使用 Google 的全新 TensorFlow Lite 執行階段 LiteRT,對即時攝影機串流執行即時圖像分割。您將使用入門 Android 應用程式,並在其中加入圖像分割功能。我們也會逐步說明預先處理、推論和後續處理步驟。您將學會以下內容:

  • 建構可即時分割圖像的 Android 應用程式。
  • 整合預先訓練的 LiteRT 圖片區隔模型。
  • 預先處理模型的輸入圖片。
  • 使用 LiteRT 執行階段進行 CPU 和 GPU 加速。
  • 瞭解如何處理模型的輸出內容,以顯示區隔遮罩。
  • 瞭解如何調整前置鏡頭。

最後,您會建立類似下圖的內容:

已完成的應用程式

必要條件

本程式碼研究室是為想獲得機器學習經驗的資深行動應用程式開發人員所設計。您必須已經熟悉下列項目:

  • 使用 Kotlin 和 Android Studio 進行 Android 開發
  • 圖像處理的基本概念

課程內容

  • 如何在 Android 應用程式中整合及使用 LiteRT 執行階段。
  • 如何使用預先訓練的 LiteRT 模型執行圖片區隔。
  • 如何預先處理模型的輸入圖片。
  • 如何執行模型推論。
  • 如何處理區隔模型輸出內容,以視覺化呈現結果。
  • 如何使用 CameraX 處理即時攝影機串流。

軟硬體需求

  • 最新版 Android Studio (已在 2025.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 是包含所有必要資源的目錄。在本程式碼研究室中,您只需要 kotlin_cpu_gpu/android_starter 專案。如果遇到困難,不妨查看完成的專案:kotlin_cpu_gpu/android

檔案路徑注意事項

本教學課程會以 Linux/macOS 格式指定檔案路徑。如果您使用 Windows,請視情況調整路徑。

此外,請務必注意 Android Studio 專案檢視畫面與標準檔案系統檢視畫面之間的差異。Android Studio 專案檢視畫面會以結構化方式呈現專案檔案,方便您進行 Android 開發作業。本教學課程中的檔案路徑是指檔案系統路徑,而非 Android Studio 專案檢視畫面中的路徑。

匯入範例應用程式

首先,請將範例應用程式匯入 Android Studio。

  1. 開啟 Android Studio,然後選取「Open」

開啟 Android Studio

  1. 前往並開啟 kotlin_cpu_gpu/android_starter 目錄。

Android Starter

為確保應用程式可使用所有依附元件,匯入程序完成後,請將專案與 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. 瞭解範例應用程式

範例應用程式已具備基本的 UI 和相機處理邏輯。以下快速說明主要檔案:

  • app/src/main/java/com/google/aiedge/examples/image_segmentation/MainActivity.kt:這是應用程式的主要進入點。這個類別會使用 Jetpack Compose 設定 UI,並處理相機權限。
  • app/src/main/java/com/google/aiedge/examples/image_segmentation/MainViewModel.kt:這個 ViewModel 會管理 UI 狀態,並協調圖片分割程序。
  • 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:您可以使用這個建構工具類別設定 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 替換為 (~line 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 替換為 (~line 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 替換為 (~line 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,並提供給 UI。將 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. 測試功能:應用程式啟動後,您應該會看到即時攝影機動態饋給,並疊加彩色區隔。
    • 切換鏡頭:輕觸頂端的鏡頭翻轉圖示,即可切換前置和後置鏡頭。請注意,疊加層會正確調整方向。
    • 變更加速器:輕觸底部的「CPU」或「GPU」按鈕,即可切換硬體加速器。請注意畫面底部顯示的「推論時間」變化。GPU 速度應該會大幅提升。
    • 使用相片庫圖片:輕觸頂端的「相片庫」分頁標籤,從裝置的相片庫選取圖片。應用程式會對所選靜態圖片執行區隔作業。

其他 UI

現在,您已建立功能完整的即時圖像分割應用程式,並採用 LiteRT 技術!

11. 進階 (選用):使用 NPU

這個存放區也包含針對神經處理單元 (NPU) 最佳化的應用程式版本。如果裝置搭載相容的 NPU,這個版本可大幅提升效能。

如要試用 NPU 版本,請在 Android Studio 中開啟 kotlin_npu/android 專案。程式碼與 CPU/GPU 版本非常相似,且已設定為使用 NPU 委派。

如要使用 NPU 委派,請先註冊搶先體驗計畫

12. 恭喜!

您已成功建構 Android 應用程式,使用 LiteRT 執行即時圖像分割。您已瞭解如何:

  • 將 LiteRT 執行階段整合至 Android 應用程式。
  • 載入及執行 TFLite 圖片區隔模型。
  • 預先處理模型的輸入內容。
  • 處理模型輸出內容,建立區隔遮罩。
  • 使用 CameraX 建立即時相機應用程式。

後續步驟

  • 請改用其他圖像分割模型。
  • 使用不同的 LiteRT 委派 (CPU、GPU、NPU) 進行實驗。

瞭解詳情