Глубокое обучение никогда не перестанет меня удивлять. Оно уже значительно повлияло на разные отрасли и продолжает открывать всё новые возможности. Классификация изображений при помощи сверточных нейронных сетей (CNN) сейчас довольно проста, особенно с появлением фронтенд-решений, например, Keras с бэкендом в виде TensorFlow. Но что если вы хотите идентифицировать больше одного объекта на картинке?
Эта проблема называется “локализация и обнаружение объектов”. Это гораздо сложнее, чем обычная классификация. До 2015 года локализация изображений при помощи CNN была медленной и неэффективной. Если вам интересна история обнаружения объектов, то все вы можете прочитать в этой статье.
Звучит отлично, сложно ли это сделать?
С реализацией нам поможет TensorFlow Object Detection API. Нужно будет только подготовить базу данных и установить несколько настроек. Вы можете обучить свою модель, а можете использовать модели, уже обученные на наборах данных MS COCO, Kitti или Open Images. Их недостаток в том, что эти модели могут предсказывать только классы, определенные наборами данных.
А как обнаружить что-то неопределенное в списке классов? Об этом и пойдет речь в этом тексте. Я покажу вам, как создать собственную программу по обнаружению объектов, используя пример квиддича из вселенной Гарри Поттера (есть ещё пост для любителей “Звездных войн”).
Начало
Для начала скопируйте мой репозиторий с GitHub. Это будет вашей основной директорией. Также вы можете скопировать репозиторий с моделями TensorFlow. В таком случае вам нужны будут только папки slim и object_detection, а все остальное вы можете удалить. Не переименовывайте ничего внутри папок.
Зависимости
Вам нужно установить несколько зависимостей в TensorFlow, выполнив следующие команды в базовой директории:
pip install -r requirements.txt
API использует Protobufs, чтобы конфигурировать и обучать модели. Перед использованием нам нужно скомпилировать библиотеки Protobuf. Сначала установите компилятор при помощи следующей команды:
sudo apt-get install protobuf-compiler
Теперь вы можете компилировать библиотеки:
protoc object_detection/protos/*.proto --python_out=.
Вам необходимо добавить путь к вашему базовому каталогу, а также к вашей slim-директории к пути к Python. Обратите внимание, что вы должны выполнять этот шаг каждый раз, когда вы открываете новый терминал. Вы можете сделать это, выполнив команду ниже. Или вы можете добавить её в файл ~ / .bashrc для автоматизации процесса.
export PYTHONPATH=$PYTHONPATH:`pwd`:`pwd`/slim
Подготовка данных
Сначала я хотел создать программу, которая будет находить снитч. Но затем я решил повысить ставку. Что если идентифицировать все мячи в игре?
В квиддиче используют три уникальных объекта. Два бладжера, один квоффл и один снитч.
Мы начнем с подготовки файла label_map.pbtxt, который будет содержать все объекты с метками, а также ID-номер каждой метки. Обратите внимание, что ID метки должен начинаться с 1. Вот содержимое файла, которое я использовал для проекта.
item {
id: 1
name: ‘snitch’
}
item {
id: 2
name: ‘quaffle’
}
item {
id: 3
name: ‘bludger’
}
Время собирать набор данных. Я собрал кадры из отрывка фильма о Гарри Поттере, используя небольшой код и фреймворк OpenCV. Затем я использовал другой код, чтобы извлечь 300 изображений из набора данных. Все это доступно в файле utils.py в моем репозитории GitHub.
Да, только 300 изображений, мой набор данных не был огромным. Но это потому, что я не смог бы подписать много изображений. При необходимости вы можете использовать платные сервисы вроде Amazon Mechanical Turk.
Аннотации
Каждая задача по локализации изображений требует аннотаций для сравнения. Здесь используются аннотации в формате XML-файлов с 4 координатами, представляющими локацию рамки вокруг объекта, и его меткой. Мы используем формат Pascal VOC. Аннотация будет выглядеть подобным образом:
<annotation>
<filename>182.jpg</filename>
<size>
<width>1280</width>
<height>586</height>
<depth>3</depth>
</size>
<segmented>0</segmented>
<object>
<name>bludger</name>
<bndbox>
<xmin>581</xmin>
<ymin>106</ymin>
<xmax>618</xmax>
<ymax>142</ymax>
</bndbox>
</object>
<object>
<name>quaffle</name>
<bndbox>
<xmin>127</xmin>
<ymin>406</ymin>
<xmax>239</xmax>
<ymax>526</ymax>
</bndbox>
</object>
</annotation>
Неужели вам придется вручную вносить аннотации в XML-файлы? Конечно, нет. Существуют инструменты для рисования рамки вокруг объекта и создания метки. Для пользователей Linux/Windows — это инструмент LabelImg, а для пользователей Mac — RectLabel.
Несколько заметок перед сбором данных:
- Не переименовывайте файлы с изображениями после аннотации. Код будет искать изображение внутри XML-файла, который автоматически заполняет название изображения. Убедитесь, что у вашего изображения и XML-файла одно и то же название.
- Измените размер изображения до желаемого перед аннотацией. Если вы сделаете это позже, то вам придется изменять данные внутри XML-файлов.
- LabelImg может добавить дополнительные элементы в файлы, не удаляйте их, если они не мешают вашему коду.
Если у вас что-то не получится, в файле utils.py есть несколько утилит, которые могут вам помочь. Или вы можете скачать мой аннотированный набор данных.
Создайте текстовый файл под названием trainval. В нем должны содержаться названия всех ваших изображений и XML-файлов. Например, если у вас есть img1.jpg, img2.jpg и img1.xml, img2.xml в наборе данных, то ваш файл будет выглядеть так:
img1 img2
Разделите свой набор данных на две папки: images и annotations. Поместите label_map.pbtxt и trainval.txt в папку annotations. Создайте папку под названием xmls внутри annotations и поместите туда все XML-файлы. Ваша иерархия директорий будет выглядеть так:
-base_directory |-images |-annotations ||-xmls ||-label_map.pbtxt ||-trainval.txt
API принимает данные в формате файла TFRecords. Ваш набор данных можно будет конвертировать при помощи небольшой утилиты. Для этого используйте файл create_tf_record.py из моего репозитория. Вам нужно будет ввести следующую команду:
python create_tf_record.py \
--data_dir=`pwd` \
--output_dir=`pwd`
После этого вы найдете два файла: train.record и val.record. Набор данных делится на 70% для обучения и 30% для оценки. Вы можете изменить это соотношение в функции main() файла.
Обучение модели
Теперь нам нужно выбрать модель локализации для обучения. Но проблема в том, что вариантов слишком много. Каждый отличается от другого по скорости и точности. Вам нужно выбрать правильную модель для правильной работы. Можете почитать эту статью для большего понимания.
SSD работают быстро, но могут ошибиться в распознавании маленьких объектов с должной точностью, а Faster RCNN работают медленно, но с большей точностью.
TensorFlow Object Detection API предоставляет ряд предобученных моделей. Я рекомендую производить обучение предварительно тренированной модели. Это сильно сократит время обучения.
Скачайте одну из этих моделей и извлеките содержимое в базовую директорию. Так как мне была больше интересна точность, но также я хотел получить разумное время обработки, я выбрал версию ResNet-50 модели Faster RCNN. После извлечения вы получите контрольные точки модели, график вывода и файл pipeline.config.
Остается одна последняя вещь. Вам нужно указать «работу обучения» в файле pipeline.config. Поместите его в базовую директорию. В последних нескольких строчках вам нужно будет установить значения для указания местоположения ваших файлов.
gradient_clipping_by_norm: 10.0
fine_tune_checkpoint: "model.ckpt"
from_detection_checkpoint: true
num_steps: 200000
}
train_input_reader {
label_map_path: "annotations/label_map.pbtxt"
tf_record_input_reader {
input_path: "train.record"
}
}
eval_config {
num_examples: 8000
max_evals: 10
use_moving_averages: false
}
eval_input_reader {
label_map_path: "annotations/label_map.pbtxt"
shuffle: false
num_epochs: 1
num_readers: 1
tf_record_input_reader {
input_path: "val.record"
}
}
Если у вас есть опыт в установке гиперпараметров модели, вы можете это сделать. Здесь есть несколько кратких советов от создателей.
Теперь вы готовы обучить модель. Введите следующую команду, чтобы начать обучение.
python object_detection/train.py \ --logtostderr \ --pipeline_config_path=pipeline.config \ --train_dir=train
Видеокарта моего ноутбука не справилась бы с размером модели, поэтому я запустил её на процессоре. Это заняло от 7 до 13 секунд на шаг. После 10 тысяч шагов модель добилась достаточной точности. Я прекратил обучение после 20 тысяч шагов, потому что это и так уже заняло два дня.
Вы можете продолжить с контрольной точки, изменив атрибут fine_tune_checkpoint в model.ckpt на model.ckpt-xxxx, где xxxx представляет номер шага из сохраненной контрольной точки.
Экспорт модели
Зачем обучать модель, если не использовать её для обнаружения объектов? Нужно снова использовать API. Но вот в чем проблема. Модуль вывода требует в качестве входных данных граф модели. Но при помощи следующей команды вы можете экспортировать модель в граф.
python object_detection/export_inference_graph.py \ --input_type=image_tensor \ --pipeline_config_path=pipeline.config \ --trained_checkpoint_prefix=train/model.ckpt-xxxxx \ --output_directory=output
Вы можете получить файл frozen_inference_graph.pb вместе с группой файлов контрольных точек.
В моем репозитории вы найдете файл inference.py. Вы можете его использовать, чтобы запустить или протестировать модуль обнаружения объектов. Код довольно понятен, и он похож на демо от создателей TensorFlow. Вы можете запустить его через следующую команду:
python object_detection/inference.py \
--input_dir={PATH} \
--output_dir={PATH} \
--label_map={PATH} \
--frozen_graph={PATH} \
--num_output_classes={NUM}
Замените {PATH} на имя файла или путь к нужному файлу, также замените {NUM} на количество объектов, которое нужно определить вашей модели (в моем случае, 3).
Результат
На производительность модели можно посмотреть в этих видео. Первое видео демонстрирует способность модели распознавать все три объекта, а второе показывает умения модели в качестве ловца.
Довольно впечатляет, хотя и у модели возникают сложности с тем, чтобы отличать головы от мячей. Но это хорошая производительность для такого размера набора данных.