Распознавание объектов на android с помощью TensorFlow: от подготовки данных до запуска на устройстве

  • Tutorial

Обучение нейросети распознаванию образов — долгий и ресурсоемкий процесс. Особенно когда под рукой есть только недорогой ноут, а не компьютер с мощной видеокартой. В этом случае на помощь придёт Google Colaboratory, которая предлагает совершенно бесплатно воспользоваться GPU уровня Tesla K80 (подробнее).


В этой статье описан процесс подготовки данных, обучения модели tensorflow в Google Colaboratory и её запуск на android устройстве.


Подготовка данных


В качестве примера попробуем обучить нейросеть распознавать белые игральные кости на черном фоне. Соответственно, для начала, надо создать набор данных, достаточный для обучения (пока остановимся на ~100 фото).



Для обучения воспользуемся Tensorflow Object Detection API. Все необходимые для обучения данные мы подготовим на ноутбуке. Нам понадобится менеджер управления окружением и зависимостями conda. Инструкция по установке тут.


Создадим окружение для работы:


conda create -n object_detection_prepare pip python=3.6

И активируем его:


conda activate object_detection_prepare

Установим зависимости которые нам понадобятся:


pip install --ignore-installed --upgrade tensorflow==1.14
pip install --ignore-installed pandas
pip install --ignore-installed Pillow
pip install lxml
conda install pyqt=5

Создадим папку object_detection, и положим все наши фотографии в папку object_detection/images.


В Google Colab есть ограничение на использование памяти, поэтому перед разметкой данных нужно снизить разрешение фотографий, чтобы в процессе обучения не столкнутся с ошибкой "tcmalloc: large alloc....".

Создадим папку object_detection/preprocessing и добавим в неё подготовленные скрипты.


Для изменения размера фото используем скрипт:


python ./object_detection/preprocessing/image_resize.py -i ./object_detection/images --imageWidth=800 --imageHeight=600

Этот скрипт пробежится по папке с указанными фото, изменит их размер до 800x600 и сложит их в object_detection/images/resized. Теперь можно заменить ими оригинальные фотографии в object_detection/images.


Для разметки данных воспользуемся тулзой labelImg.


Клонируем репозиторий labelImg в object_detection


Переходим в папку labelImg


cd [FULL_PATH]/object_detection/labelImg 

и выполняем команду:


pyrcc5 -o libs/resources.py resources.qrc

После этого можно приступить к разметке данных (это самый долгий и скучный этап):


python labelImg.py


В “Open dir” указываем папку object_detection/images и проходим по всем фото, выделяя объекты для распознавания и указывая их класс. В нашем случае это номиналы игральных костей (1, 2, 3, 4, 5, 6). Сохраним метаданные (файлы *.xml) в той же папке.


Создадим папку object_detection/training_demo, которую мы чуть позже зальем в Google Colab для обучения.


Разделим наши фото (с метаданными) на тренировочные и тестовые в соотношение 80/20 и переместим их в соответствующие папки object_detection/training_demo/images/train и object_detection/training_demo/images/test.


Создадим папку object_detection/training_demo/annotations, в которую будем складывать файлы с метаданными необходимыми для обучения. Первым из них будет label_map.pbtxt, в котором укажем отношение класса объекта и целочисленного значения. В нашем случае это:


label_map.pbtxt
item {
    id: 1
    name: '1'
}

item {
    id: 2
    name: '2'
}

item {
    id: 3
    name: '3'
}

item {
    id: 4
    name: '4'
}

item {
    id: 5
    name: '5'
}

item {
    id: 6
    name: '6'
}

Помните метаданные, которые мы получили в процессе разметки данных? Чтобы использовать их для обучения, необходимо конвертировать их в формат TFRecord. Для конвертации воспользуемся скриптами из источника [1].


Проведем конвертацию в два этапа: xml -> csv и csv -> record


Перейдём в папку preprocessing :


cd [FULL_PATH]\object_detection\preprocessing

1.Из xml в csv


Тренировочные данные:


python xml_to_csv.py -i [FULL_PATH]/object_detection/training_demo/images/train -o [FULL_PATH]/object_detection/training_demo/annotations/train_labels.csv

Тестовые данные:


python xml_to_csv.py -i [FULL_PATH]/object_detection/training_demo/images/test -o [FULL_PATH]/object_detection/training_demo/annotations/test_labels.csv

2. Из csv в record


Тренировочные данные:


python generate_tfrecord.py --label_map_path=[FULL_PATH]\object_detection\training_demo\annotations\label_map.pbtxt --csv_input=[FULL_PATH]\object_detection\training_demo\annotations\train_labels.csv --output_path=[FULL_PATH]\object_detection\training_demo\annotations\train.record --img_path=[FULL_PATH]\object_detection\training_demo\images\train

Тестовые данные:


python generate_tfrecord.py --label_map_path=[FULL_PATH]\object_detection\training_demo\annotations\label_map.pbtxt --csv_input=[FULL_PATH]\object_detection\training_demo\annotations\test_labels.csv --output_path=[FULL_PATH]\object_detection\training_demo\annotations\test.record --img_path=[FULL_PATH]\object_detection\training_demo\images\test

На этом мы закончили подготовку данных, теперь надо выбрать модель, которую будем обучать.


Доступные модели для переобучения можно найти тут.


Сейчас мы выберем модель ssdlite_mobilenet_v2_coco, чтобы в дальнейшем запустить обученную модель на android устройстве.


Скачиваем архив с моделью и распаковываем его в object_detection/training_demo/pre-trained-model.


Должно получиться что-то вроде
object_detection/training_demo/pre-trained-model/ssdlite_mobilenet_v2_coco_2018_05_09


Из распакованного архива копируем файл pipeline.config в object_detection/training_demo/training и переименовываем его в ssdlite_mobilenet_v2_coco.config.


Далее нам надо настроить его под свою задачу, для этого:


1. Укажем количество классов


model.ssd.num_classes: 6

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


train_config.batch_size: 18
train_config.num_steps: 20000
train_config.fine_tune_checkpoint:"./training_demo/pre-trained-model/ssdlite_mobilenet_v2_coco_2018_05_09/model.ckpt"

3. Укажем количество фото в тренировочном наборе (object_detection/training_demo/images/train)


eval_config.num_examples: 64

4. Укажем путь к набору данных для тренировки


train_input_reader.label_map_path: "./training_demo/annotations/label_map.pbtxt"
train_input_reader.tf_record_input_reader.input_path:"./training_demo/annotations/train.record"

5. Укажем путь к тестовому набору данных


eval_input_reader.label_map_path: "./training_demo/annotations/label_map.pbtxt"
eval_input_reader.tf_record_input_reader.input_path:"./training_demo/annotations/test.record"

В итоге, у вас должно получиться что то вроде этого.


Далее архивируем папку training_demo и получившийся training_demo.zip заливаем в Google Drive.


Альтернатива работы с архивом
В этой статье рассказывается как примонтировать google drive к виртуальной машине Google Colab, но нужно не забыть поменять все пути в конфигах и скриптах

На этом подготовка данных закончена, перейдем к обучению.


Обучение модели


В Google Drive выбираем training_demo.zip, кликаем на Get shareable link и из полученной ссылки сохраним себе id этого файла:


drive.google.com/open?id=[YOUR_FILE_ID_HERE]


Самый просто способ воспользоваться Google Colab — создать новый блокнот в Google Drive.



По умолчанию обучение будет выполняться на CPU. Чтобы использовать GPU, нужно поменять тип runtime.




Готовый блокнот можно взять тут.


Обучение состоит из следующих этапов:


1. Клонируем репозиторий TensorFlow Models:


!git clone https://github.com/tensorflow/models.git                                                

2. Устанавливаем protobuf и компилируем необходимые файлы в object_detection:


!apt-get -qq install libprotobuf-java protobuf-compiler                                               
%cd ./models/research/
!protoc object_detection/protos/*.proto --python_out=.
%cd ../..

3. Добавляем необходимые пути в переменную окружения PYTHONPATH:


import os
os.environ['PYTHONPATH'] += ":/content/models/research/"
os.environ['PYTHONPATH'] += ":/content/models/research/slim"
os.environ['PYTHONPATH'] += ":/content/models/research/object_detection"
os.environ['PYTHONPATH'] += ":/content/models/research/object_detection/utils"

4. Для получения файла из Google Drive устанавливаем PyDrive и авторизируемся:


!pip install -U -q PyDrive

from pydrive.auth import GoogleAuth
from pydrive.drive import GoogleDrive
from google.colab import auth
from oauth2client.client import GoogleCredentials

auth.authenticate_user()
gauth = GoogleAuth()
gauth.credentials = GoogleCredentials.get_application_default()
drive = GoogleDrive(gauth)

5. Скачиваем архив (нужно указать id вашего файла) и разархивируем его:


drive_file_id="[YOUR_FILE_ID_HERE]"

training_demo_zip = drive.CreateFile({'id': drive_file_id})
training_demo_zip.GetContentFile('training_demo.zip')

!unzip training_demo.zip
!rm training_demo.zip

6. Запускаем процесс обучения:


!python ./models/research/object_detection/legacy/train.py --logtostderr --train_dir=./training_demo/training --pipeline_config_path=./training_demo/training/ssdlite_mobilenet_v2_coco.config

Описание параметров

--train_dir=./training_demo/training — путь к директории, где будут лежать
результаты обучение


--pipeline_config_path=./training_demo/training/ssdlite_mobilenet_v2_coco.config — путь к конфигу


7. Конвертируем результат обучения в frozen graph, который можно будет использовать:


!python /content/models/research/object_detection/export_inference_graph.py --input_type image_tensor --pipeline_config_path /content/training_demo/training/ssdlite_mobilenet_v2_coco.config --trained_checkpoint_prefix /content/training_demo/training/model.ckpt-[CHECKPOINT_NUMBER]
 --output_directory /content/training_demo/training/output_inference_graph_v1.pb

Описание параметров

--pipeline_config_path /content/training_demo/training/ssdlite_mobilenet_v2_coco.config — путь к конфигу


--trained_checkpoint_prefix /content/training_demo/training/model.ckpt-[CHECKPOINT_NUMBER] — путь к чекпоинту, который мы хотим конвертировать.


--output_directory /content/training_demo/training/output_inference_graph_v1.pb — имя сконвертированной модели


Номер чекпоинта [CHECKPOINT_NUMBER], можно посмотреть в папке content/training_demo/training/. ПОСЛЕ обучения там должны появиться файлы типа model.ckpt-1440.index, model.ckpt-1440.meta. 1440 — это и [CHECKPOINT_NUMBER] и номер итерации обучения.



Для визуализации результата обучения в блокноте есть специальный скрипт. На рисунке ниже показан результат распознавания изображения из тестового набора данных после ~20000 итераций обучения.



8. Конвертация обученной модели в tflite.
Для использования tensorflow lite нужно конвертировать модель в формат tflite. Для этого конвертируем результат обучения в frozen graph который поддерживает конвертацию в tflite (параметры такие же как при использовании скрипта export_inference_graph.py):


!python /content/models/research/object_detection/export_tflite_ssd_graph.py --pipeline_config_path /content/training_demo/training/ssdlite_mobilenet_v2_coco.config --trained_checkpoint_prefix /content/training_demo/training/model.ckpt-[CHECKPOINT_NUMBER] --output_directory /content/training_demo/training/output_inference_graph_tf_lite.pb

Для конвертации в tflite нам нужна дополнительная информация о модели, для её получения скачиваем модель из output_inference_graph_tf_lite.pb:



И открываем её в тулзе Netron. Нас интересует названия и размерность входного и выходного узлов модели.




Зная их можно конвертировать pb модель в tflite формат:


!tflite_convert --output_file=/content/training_demo/training/model_q.tflite  --graph_def_file=/content/training_demo/training/output_inference_graph_tf_lite_v1.pb/tflite_graph.pb --input_arrays=normalized_input_image_tensor  --output_arrays='TFLite_Detection_PostProcess','TFLite_Detection_PostProcess:1','TFLite_Detection_PostProcess:2','TFLite_Detection_PostProcess:3' --input_shapes=1,300,300,3 --enable_select_tf_ops --allow_custom_ops  --inference_input_type=QUANTIZED_UINT8 --inference_type=FLOAT --mean_values=128 --std_dev_values=128

Описание параметров

--output_file=/content/training_demo/training/model_q.tflite — путь к результату конвертации


--graph_def_file=/content/training_demo/training/output_inference_graph_tf_lite_v1.pb/tflite_graph.pb — путь к frozen graph, который нужно конвертировать


--input_arrays=normalized_input_image_tensor — название входного узла, которое мы узнали выше


--output_arrays='TFLite_Detection_PostProcess','TFLite_Detection_PostProcess:1','TFLite_Detection_PostProcess:2','TFLite_Detection_PostProcess:3' — названия выходных узлов, которые мы узнали выше


--input_shapes=1,300,300,3 — размерность входных данных, которые мы узнали выше


--enable_select_tf_ops — для использования расширенного runtime TensorFlow Lite


--allow_custom_ops — для использования TensorFlow Lite Optimizing Converter


--inference_type=FLOAT — тип данных для всех массивов в модели кроме входных


--inference_input_type=QUANTIZED_UINT8 — тип данных для всех входных массивов в модели


--mean_values=128 --std_dev_values=128 — среднее значение и стандартное отклонение входных данных, для использования QUANTIZED_UINT8


Архивируем папку с результатами обучения и заливаем её в Google Drive:


!zip -r ./training_demo/training.zip ./training_demo/training/

training_result = drive.CreateFile({'title': 'training_result.zip'})
training_result.SetContentFile('training_demo/training.zip')
training_result.Upload()

Если будет ошибка Invalid client secrets file, то нужно сделать повторную авторизацию в google drive.

Запуск модели на android устройстве


В основе android приложения был использован официальный гайд по object detection, но он был полностью переписан с использованием kotlin и CameraX. Полностью код можно посмотреть тут.


CameraX уже имеет механизм для анализа входящих фреймов с камеры с помощью ImageAnalysis. Логика по распознаванию находится в ObjectDetectorAnalyzer.


Весь процесс распознавания изображения можно разбить на несколько этапов:


1. На вход мы получаем изображение которое имеет YUV формат. Для дальнейшей работы его нужно конвертировать в RGB формат:


val rgbArray = convertYuvToRgb(image)

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


val rgbBitmap = getRgbBitmap(rgbArray, image.width, image.height)
val transformation =  getTransformation(rotationDegrees, image.width, image.height)
Canvas(resizedBitmap).drawBitmap(rgbBitmap, transformation, null)

3. Конвертируем bitmap в массив пикселей, и отдадим его на вход детектора:


ImageUtil.storePixels(resizedBitmap, inputArray)
val objects = detect(inputArray)

4. Для визуализации передадим результат распознавания в RecognitionResultOverlayView и преобразуем координаты с соблюдением соотношения сторон:


val scaleFactorX = measuredWidth / result.imageWidth.toFloat()
val scaleFactorY = measuredHeight / result.imageHeight.toFloat()

result.objects.forEach { obj ->
    val left = obj.location.left * scaleFactorX
    val top = obj.location.top * scaleFactorY
    val right = obj.location.right * scaleFactorX
    val bottom = obj.location.bottom * scaleFactorY

    canvas.drawRect(left, top, right, bottom, boxPaint)
    canvas.drawText(obj.text, left, top - 25f, textPaint)
}

Чтобы запустить нашу модель в приложении, нужно заменить в assets файл с моделью на training_demo/training/model_q.tflite (переименовав её в detect.tflite) и файл с метками labelmap.txt, в нашем случае это:


detect.tflite
1
2
3
4
5
6

Так как в официальном гайде используется SSD Mobilenet V1, в котором индексация меток начинает с 1, а не с 0, нужно поменять labelOffset с 1 на 0 в методе collectDetectionResult класса ObjectDetector.


На этом всё.
На видео результат того, как обученная модель работает на стареньком Xiaomi Redmi 4X :



В процессе работы использовались следующие ресурсы:

Redmadrobot
№1 в разработке цифровых решений для бизнеса

Комментарии 2

    0
    какой максимальное количество элементов (аналогичного оральным костям размеров) можно детектить в реалтайм?
      0

      Не проводил такое исследование. В официальный гайде используют:


      The object detection model we provide can identify and locate up to 10 objects in an image. It is trained to recognize 80 classes of object.

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


      Можно попробовать использовать оптимизацию для распознавания на CPU — Post-training quantization

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

    Самое читаемое