Machine learning

В современном мире машинное обучение применяется в самых разных областях: классификация изображений, прогнозы покупательского спроса, рекомендации фильмов и музыки для конкретного человека, кластеризация. При этом для достаточно больших моделей вычисление результата (и в гораздо большей степени обучение модели) может быть ресурсоёмкой операцией.

Для того, чтобы использовать обученные модели на не самых мощных устройствах, компания Google представила свой фреймворк TensorFlow Lite. Для работы с ним необходимо обучить модель, построенную с помощью фреймворка TensorFlow (не Lite!), затем перенести её в формат TensorFlow Lite. После этого модель может быть легко использована на встраиваемых или мобильных устройствах.

В этой статье мы рассмотрим все шаги для запуска модели на Android.

Обучение и перенос модели

Для примера возьмём одну из стандартных моделей MobleNet. Обучение (training) будет проводиться на наборе картинок ILSVRC-2012-CLS.

Для обучения скачаем на достаточно мощную машину с ОС Linux набор моделей:

git clone (ссылка)

Установим систему сборки bazel по инструкции на сайте.

Запустим обучение модели:

bazel build -c opt mobilenet_v1_{eval,train}
./bazel-bin/mobilenet_v1_train --dataset_dir путь_к_набору_картинок
—checkpoint_dir путь_к_чекпоинтам

Данные команды выполнят обучение модели и создадут файлы с расширением *.tflite, необходимые для использования интерпретатором TensorFlow Lite.

В более общем случае, если необходимо использовать модель, описанную с применением фреймворка TensorFlow, то после обучения её необходимо сохранить и перенести в формат TensorFlow Lite с помощью конвертора. В нашем случае этот шаг не требуется.

Приложение Android

Приложение будет содержать следующие функции:

  • Захват изображений с камеры в реальном времени
  • Классификация изображения с использованием модели TensorFlow Lite
  • Отображение на экран результата классификации

Как заготовку для данного приложения можно использовать исходные тексты примера «image_classification» с сайта.

Захват изображений в реальном времени

Будет использован фреймворк «android.hardware.camera2», доступный начиная с Android 5. С помощью системного сервиса CameraManager получаем доступ к камере и, реализуя интерфейс OnImageAvailableListener, будем получать кадры в методе onImageAvailable.

Кадры будут приходить несколько раз в секунду с частотой, зависящей от аппаратной реализации встроенной камеры. Для наглядности также будем выводить изображение, получаемое с камеры на экране, разместив компонент TextureView в макете (layout) с размером, совпадающим с размером экрана (за исключением небольшого участка снизу TextureView, где мы разместим информацию о результатах классификации объекта). Для этого свяжем данный компонент с выводом с камеры.

В методе onImageAvailable мы сначала получим последний доступный кадр, вызвав acquireLatestImage(). Может возникнуть ситуация, когда классификация  кадра  займёт много времени, и камера за это время успеет выдать более одного кадра. Именно поэтому мы берём последний кадр, пропуская необработанные. Кадр приходит в формате YUV420. Преобразуем его в формат ARGB8888 с помощью вызова convertYUV420 к библиотеке ImageUtils. Так как нужный нам формат — это массив чисел с плавающей точкой одинарной точности (float) в диапазоне от -1 до 1, то выполним и это преобразование.

Классификация изображения

Перед выполнением классификации, в начале выполнения приложения, необходимо загрузить модель из файла, зашитого в самом приложении в каталоге assets. В этот каталог необходимо скопировать файл модели mobilenet_v1_1.0_224.tflite, полученный на этапе «Обучение и перенос модели», и описание объектов классификации label.txt.

После этого создаётся объект интерпретатора TensorFlow Lite. В качестве параметров конструктора передаётся модель.

Классификация выполняется простым вызовом метода run интерпретатора. Первый параметр метода — это распознаваемый кадр в формате массива чисел с плавающей точкой, полученный нами с камеры. Второй параметр — это массив массивов чисел с плавающей точкой. Эти массивы будут заполнены результатами работы модели. Результатами в нашем случае является массив вероятностей. Каждый элемент массива c индексом i — это вероятность, что в кадре находится именно i-тый объект.

Отображение на экран результата классификации

Отобразим на экране название самого вероятного объекта, используя предварительно загруженные названия объектов из файла названий, хранящегося в assets. В используемой нами модели названия на английском, но можно предварительно перевести их на русский язык, изменив файл labels.txt.

Как вариант, можно отобразить не просто название самого вероятного объекта, а список из 3-5 наиболее вероятных объектов вместе с соответствующими вероятностями.

Заключение

Мы рассмотрели, как можно использовать предварительно обученную модель классификации изображений в ОС Android. При этом использование моделей возможно и на встраиваемых устройствах, так как интерпретатор TensorFlow Lite имеет также интерфейс C++ и занимает размер около 300 килобайт.

Статья была опубликована на сайте МКА