Ускоренная сегментация изображений в реальном времени на Android с помощью LiteRT

1. Прежде чем начать

Набор кода — отличный способ развить мышечную память и углубить понимание материала. Хотя копирование и вставка могут сэкономить время, в долгосрочной перспективе это может привести к повышению эффективности и укреплению навыков программирования.

В этой лабораторной работе вы научитесь создавать приложение для Android, которое выполняет сегментацию изображений в реальном времени на основе трансляции с камеры, используя новую среду выполнения Google для TensorFlow Lite — LiteRT. Вы возьмёте базовое приложение для Android и добавите в него возможности сегментации изображений. Мы также рассмотрим этапы предварительной обработки, вывода и постобработки. Вы:

  • Создайте приложение для Android, которое сегментирует изображения в режиме реального времени.
  • Интегрируйте предварительно обученную модель сегментации изображений LiteRT.
  • Предварительная обработка входного изображения для модели.
  • Используйте среду выполнения LiteRT для ускорения CPU и GPU.
  • Понять, как обрабатывать выходные данные модели для отображения маски сегментации.
  • Узнайте, как настроить фронтальную камеру.

В итоге у вас получится что-то похожее на изображение ниже:

Готовое приложение

Предпосылки

Эта практическая работа разработана для опытных разработчиков мобильных приложений, желающих получить опыт в машинном обучении. Вам необходимо знать:

  • Разработка Android с использованием Kotlin и Android Studio
  • Основные концепции обработки изображений

Чему вы научитесь

  • Как интегрировать и использовать среду выполнения LiteRT в Android-приложении.
  • Как выполнить сегментацию изображения с использованием предварительно обученной модели LiteRT.
  • Как предварительно обработать входное изображение для модели.
  • Как выполнить вывод для модели.
  • Как обрабатывать выходные данные модели сегментации для визуализации результатов.
  • Как использовать CameraX для обработки изображений с камер в реальном времени.

Что вам понадобится

  • Последняя версия Android Studio (протестировано на v2025.1.1).
  • Физическое устройство Android. Лучше всего тестировать на устройствах Galaxy и Pixel.
  • Пример кода (из GitHub).
  • Базовые знания разработки Android на Kotlin.

2. Сегментация изображения

Сегментация изображений — это задача компьютерного зрения, которая включает в себя разбиение изображения на несколько сегментов или областей. В отличие от распознавания объектов, при котором вокруг объекта рисуется ограничивающая рамка, сегментация изображений присваивает каждому пикселю изображения определённый класс или метку. Это обеспечивает гораздо более детальное и детальное понимание содержимого изображения, позволяя определить точную форму и границы каждого объекта.

Например, вместо того, чтобы просто знать, что «человек» находится в рамке, вы можете точно знать, какие пиксели принадлежат этому человеку. В этом руководстве показано, как выполнить сегментацию изображений в реальном времени на устройстве Android с помощью предварительно обученной модели машинного обучения.

Пример сегментации

LiteRT: Расширение возможностей машинного обучения на устройствах

LiteRT — ключевая технология, обеспечивающая высокоточную сегментацию в реальном времени на мобильных устройствах. LiteRT — это высокопроизводительная среда выполнения нового поколения от Google для TensorFlow Lite, разработанная для достижения максимальной производительности базового оборудования.

Это достигается благодаря интеллектуальному и оптимизированному использованию аппаратных ускорителей, таких как графический процессор (GPU) и нейронный процессор (NPU). Перенося интенсивную вычислительную нагрузку модели сегментации с центрального процессора общего назначения на эти специализированные процессоры, 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 и выберите Открыть .

Android Studio Open

  1. Перейдите в каталог kotlin_cpu_gpu/android_starter и откройте его.

Android Starter

Чтобы убедиться, что все зависимости доступны вашему приложению, вам следует синхронизировать свой проект с файлами Gradle после завершения процесса импорта.

  1. Выберите «Синхронизировать проект с файлами Gradle» на панели инструментов Android Studio.

Синхронизация меню

  1. Пожалуйста, не пропускайте этот шаг — если это не сработает, остальная часть урока не будет иметь смысла.

Запустите стартовое приложение

Теперь, когда вы импортировали проект в Android Studio, вы готовы запустить приложение в первый раз.

Подключите Android-устройство к компьютеру через USB и нажмите «Запустить» на панели инструментов Android Studio.

Кнопка «Выполнить»

Приложение должно запуститься на вашем устройстве. Вы увидите трансляцию с камеры в реальном времени, но сегментация пока не будет выполняться. Все файлы, которые вы внесёте в этом руководстве, будут находиться в каталоге LiteRT/litert/samples/image_segmentation/kotlin_cpu_gpu/android_starter/app/src/main/java/com/google/aiedge/examples/image_segmentation (теперь вы знаете, почему Android Studio реструктурирует его 😃).

Директор проекта

Вы также увидите комментарии TODO в файлах ImageSegmentationHelper.kt , MainViewModel.kt и view/SegmentationOverlay.kt . На следующих этапах вы реализуете функцию сегментации изображений, заполнив эти TODO .

4. Разберитесь с начальным приложением

Стартовое приложение уже имеет базовый пользовательский интерфейс и логику управления камерой. Вот краткий обзор ключевых файлов:

  • 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)

После добавления зависимости синхронизируйте свой проект с файлами Gradle, нажав кнопку «Синхронизировать сейчас» , которая появится в правом верхнем углу Android Studio.

Синхронизировать сейчас

2. Изучите ключевые API LiteRT

Откройте ImageSegmentationHelper.kt

Прежде чем писать код реализации, важно понимать основные компоненты API LiteRT, которые вы будете использовать. Убедитесь, что вы импортируете данные из пакета com.google.ai.edge.litert , добавив следующие импорты в начало файла ImageSegmentationHelper.kt :

import com.google.ai.edge.litert.Accelerator
import com.google.ai.edge.litert.CompiledModel
  • CompiledModel : это центральный класс для взаимодействия с моделью TFLite. Он представляет собой модель, предварительно скомпилированную и оптимизированную для конкретного аппаратного ускорителя (например, центрального или графического процессора). Эта предварительная компиляция — ключевая функция LiteRT, обеспечивающая более быстрый и эффективный вывод.
  • CompiledModel.Options : этот класс-конструктор используется для настройки CompiledModel . Наиболее важным параметром является указание аппаратного ускорителя, который вы хотите использовать для запуска модели.
  • Accelerator : это перечисление позволяет выбрать оборудование для вывода. Стартовый проект уже настроен на обработку следующих параметров:
    • Accelerator.CPU : для запуска модели на центральном процессоре устройства. Это наиболее универсальный совместимый вариант.
    • Accelerator.GPU : для запуска модели на графическом процессоре устройства. Для моделей на основе изображений он часто значительно быстрее центрального процессора.
  • Входные и выходные буферы ( TensorBuffer ) : LiteRT использует TensorBuffer для входных и выходных данных модели. Это обеспечивает точный контроль над памятью и позволяет избежать ненужного копирования данных. Вы получаете эти буферы непосредственно из экземпляра CompiledModel с помощью model.createInputBuffers() и model.createOutputBuffers() , а затем записываете в них входные данные и считываете из них результаты.
  • 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 . Этот код настраивает модель с указанным ускорителем (ЦП или ГП) и подготавливает её к выводу. Замените TODO в initSegmenter следующей реализацией (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. Запуск сегментации по кадрам камеры : функции segment в MainViewModel являются точкой входа для нашей задачи сегментации. Они вызываются при получении нового изображения с камеры или выборе изображения из галереи. Затем эти функции вызывают метод segment в нашем ImageSegmentationHelper . Замените TODO в обеих функциях segment следующим кодом (строка ~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 : эта функция служит обёрткой, вызывающей функцию закрытого segment в классе Segmenter . Замените 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. Реализуйте предварительную обработку : в закрытой функции segment внутри класса Segmenter мы выполним необходимые преобразования входного изображения для его подготовки к модели. Это включает масштабирование, поворот и нормализацию изображения. Затем эта функция вызовет другую закрытую функцию segment для выполнения вывода. Замените TODO в функции segment(bitmap: Bitmap, ...) на (~строка 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. Реализация выполнения модели : функция private 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 и передает его в пользовательский интерфейс. Замените TODO в свойстве segmentationUiShareFlow на (~строка 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. Управление ориентацией наложения : найдите TODO в файле SegmentationOverlay.kt и замените его следующим кодом. Этот код проверяет, активна ли фронтальная камера, и, если да, применяет горизонтальное отражение к растровому 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.

Кнопка «Выполнить»

  1. Протестируйте функции : после запуска приложения вы увидите прямую трансляцию с камеры с цветной сегментацией.
    • Переключение камер : нажмите на значок переворота камеры вверху, чтобы переключиться между фронтальной и основной камерами. Обратите внимание, как правильно ориентируется наложение.
    • Изменить ускоритель : нажмите кнопку «ЦП» или «ГП» внизу, чтобы переключить аппаратный ускоритель. Наблюдайте за изменением времени вывода, отображаемым в нижней части экрана. Графический процессор должен работать значительно быстрее.
    • Использование изображения из галереи : нажмите вкладку «Галерея» в верхней части экрана, чтобы выбрать изображение из фотогалереи вашего устройства. Приложение выполнит сегментацию выбранного статического изображения.

Другой пользовательский интерфейс

Теперь у вас есть полнофункциональное приложение для сегментации изображений в реальном времени на базе LiteRT!

11. Дополнительно (необязательно): использование NPU

Этот репозиторий также содержит версию приложения, оптимизированную для нейронных процессоров (NPU). Версия для NPU может обеспечить значительный прирост производительности на устройствах с совместимым NPU.

Чтобы попробовать версию NPU, откройте проект kotlin_npu/android в Android Studio. Код очень похож на версию для CPU/GPU и настроен на использование делегата NPU.

Чтобы воспользоваться делегатом NPU, вам необходимо зарегистрироваться в Программе раннего доступа .

12. Поздравляем!

Вы успешно создали Android-приложение, которое выполняет сегментацию изображений в реальном времени с помощью LiteRT. Вы узнали, как:

  • Интегрируйте среду выполнения LiteRT в приложение для Android.
  • Загрузите и запустите модель сегментации изображений TFLite.
  • Предварительная обработка входных данных модели.
  • Обработайте выходные данные модели для создания маски сегментации.
  • Используйте CameraX для приложения камеры в реальном времени.

Следующие шаги

  • Попробуйте другую модель сегментации изображения.
  • Поэкспериментируйте с различными делегатами LiteRT (CPU, GPU, NPU).

Узнать больше