Сыграем в квиддич с TensorFlow Object Detection API

Глубокое обучение никогда не перестанет меня удивлять. Оно уже значительно повлияло на разные отрасли и продолжает открывать всё новые возможности. Классификация изображений при помощи сверточных нейронных сетей (CNN) сейчас довольно проста, особенно с появлением фронтенд-решений, например, Keras с бэкендом в виде TensorFlow. Но что если вы хотите идентифицировать больше одного объекта на картинке?

Эта проблема называется “локализация и обнаружение объектов”. Это гораздо сложнее, чем обычная классификация. До 2015 года локализация изображений при помощи CNN была медленной и неэффективной. Если вам интересна история обнаружения объектов, то все вы можете прочитать в этой статье.

Звучит отлично, сложно ли это сделать?

С реализацией нам поможет TensorFlow Object Detection API. Нужно будет только подготовить базу данных и установить несколько настроек. Вы можете обучить свою модель, а можете использовать модели, уже обученные на наборах данных MS COCO, Kitti или Open Images. Их недостаток в том, что эти модели могут предсказывать только классы, определенные наборами данных.

А как обнаружить что-то неопределенное в списке классов? Об этом и пойдет речь в этом тексте. Я покажу вам, как создать собственную программу по обнаружению объектов, используя пример квиддича из вселенной Гарри Поттера (есть ещё пост для любителей “Звездных войн”).

Как изучить Deep Learning за 6 месяцев

Начало

Для начала скопируйте мой репозиторий с 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).

Результат

На производительность модели можно посмотреть в этих видео. Первое видео демонстрирует способность модели распознавать все три объекта, а второе показывает умения модели в качестве ловца.

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