В современном мире машинное обучение применяется в самых разных областях: классификация изображений, прогнозы покупательского спроса, рекомендации фильмов и музыки для конкретного человека, кластеризация. При этом для достаточно больших моделей вычисление результата (и в гораздо большей степени обучение модели) может быть ресурсоёмкой операцией.
Для того, чтобы использовать обученные модели на не самых мощных устройствах, компания 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 килобайт.
Статья была опубликована на сайте МКА