Обучение нейросети распознаванию образов — долгий и ресурсоемкий процесс. Особенно когда под рукой есть только недорогой ноут, а не компьютер с мощной видеокартой. В этом случае на помощь придёт 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 :



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