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

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 выбираем 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

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

Для визуализации результата обучения в блокноте есть специальный скрипт. На рисунке ниже показан результат распознавания изображения из тестового набора данных после ~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

Архивируем папку с результатами обучения и заливаем её в 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, в нашем случае это:

1
2
3
4
5
6

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

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

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

• https://tensorflow-object-detection-api-tutorial.readthedocs.io
• https://github.com/EdjeElectronics/TensorFlow-Object-Detection-API-Tutorial-Train-Multiple-Objects-Windows-10
• https://lutzroeder.github.io/netron/