Начальный бизнес-вопрос: можем ли мы сфотографировать кусочек пазла и фотографию его коробки и предсказать, к какому месту в пазле он принадлежит? И если да, то как мы это делаем?

Чтобы ответить на этот вопрос, подход будет разбит на три части:

Часть I. Сегментация экземпляров

Прежде чем мы сможем попытаться предсказать, где в головоломке принадлежит кусок, нам сначала нужно извлечь этот кусок из изображения, чтобы при сравнении его с головоломкой в ​​изображении не было посторонних деталей. Для этого мы рассмотрим шаги по реализации Mask R-CNN, которая может обнаруживать кусочки головоломки с помощью Python 3, Keras и Tensorflow.

Часть II. Извлечение экземпляра объекта

После обучения модели обнаружению частей головоломки, мы захотим создать новое изображение, используя предсказание модели. Мы будем использовать только ограничивающую рамку (ROI, область интереса) изображения и возьмем координаты прогнозируемой маски, чтобы удалить фон и добавить альфа-канал, чтобы фон был прозрачным.

Часть III: кусочек пазла и коробка SIFT Feature Matching

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

Итак, приступим.

Часть I. Сегментация экземпляров

Во-первых, давайте импортируем некоторые из библиотек, которые мы будем использовать, и зададим пути.

Для нашей модели сегментации экземпляра мы будем использовать реализацию Matterport для Mask R-CNN. В корневом каталоге нашего проекта мы хотим клонировать Mask R-CNN и установить его требования.

Поскольку мы пользуемся преимуществом трансферного обучения и начинаем веса с модели, предварительно обученной на наборе данных COCO, нам нужно будет клонировать репозиторий COCO и установить его требования.

Создание набора данных

Чтобы обучить нашу модель, нам нужно будет создать собственные наборы данных для обучения и проверки частей головоломки вместе с аннотациями на уровне пикселей, обозначающими часть головоломки. Для этого мы воспользуемся V GG Image Annotator (VIA), бесплатным инструментом для создания аннотаций на основе HTML, который экспортируется как JSON. Мы будем хранить обучающий набор в каталоге наборов данных / train, а набор для проверки - / val.

Настройка маски R-CNN

Теперь, когда у нас есть подготовленные данные, чтобы заставить работать Matterport Mask R-CNN, нам нужно написать два настраиваемых класса, которые наследуются от классов Config и Dataset, которые он предоставляет.

Класс конфигурации

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

Класс набора данных

Затем подкласс Dataset должен заново реализовать пару методов, используемых для чтения и анализа аннотаций, загрузки изображений и создания масок.

Когда наши конфигурации настроены, последнее, что нам нужно сделать перед окончательным обучением, - это инициализировать модель (с нашим PieceDetectorConfig), загрузить и подготовить набор данных (с нашим VIADataset).

Инициализация модели

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

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

Наконец, мы будем обучать модель 10 эпохам.

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

Проверка модели (режим вывода)

Реализация Matterport Mask R-CNN чрезвычайно хорошо документирована и рекомендуется к прочтению, если вы собираетесь развивать их работу. Ниже мы будем использовать два их метода для оценки модели.

Чтобы визуализировать производительность модели, нам нужно получить доступ к исходному изображению, маске и ограничивающей рамке, которые были созданы с помощью аннотаций:

Затем, чтобы визуализировать предсказание и достоверность, мы передаем эти значения в один из методов, предоставляемых Matterport:

И результат должен выглядеть следующим образом:

Другой способ оценить производительность модели - посмотреть на Intersection over Union (IoU), который измеряет процент перекрытия между базовым истинным bbox / маской и прогнозируемым bbox / маской. Как правило, значения IoU, превышающие 0,50, считаются приемлемыми, но мы работаем в среде с одним классом и одним экземпляром, поэтому значения IoU должны быть намного выше.

И, если мы хотим выполнить прогнозы для нашего тестового (неаннотированного) набора данных, все, что нам нужно сделать, это передать (загруженное) изображение в нашу модель:

Вот и все - мы успешно создали модель, которая может обнаруживать кусочки головоломки. В следующем посте мы продолжим работу над получением прогнозов и созданием нового изображения только ограничивающей рамки с удаленным фоном.

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