В этой статье я бы хотел рассказать про некоторые приемы работы с данными при обучении модели. В частности, как натянуть сегментацию объектов на ббоксы, а также как обучить модель и получить разметку датасета, разметив всего несколько сэмплов.
Задача
Есть определенный процесс приготовления пиццы и фотографии с его разных стадий (в том числе не только пиццы). Известно, что если рецептура теста испорчена, то на коржике будут белые пупырышки. Также есть бинарная разметка качества теста каждого экземпляра пиццы, сделанная специалистами. Необходимо разработать алгоритм, который будет определять качество теста по фотографии.
Датасет состоит из фотографий, сделанных с разных телефонов, в разных условиях, разных ракурсах. Экземпляров пицц — 17к. Всего фотографий — 60к.
На мой взгляд, задача вполне типичная и хорошо подходит, чтобы показать разные подходы обращения с данными. Для ее решения необходимо:
1. Выбрать фотографии, где есть коржик пиццы;
2. На выбранных фотографиях выделить коржик;
3. Обучить нейросеть на выделенных участках.
Фильтрация фотографий
На первый взгляд кажется, что проще всего было бы отдать эту задачу разметчикам, и потом на чистых данных обучить датасет. Однако я решил, что мне проще самому разметить малую часть, чем объяснять разметчиком, какой ракурс — правильный. Тем более, у меня и не было жесткого критерия правильного ракурса.
Поэтому вот как я поступил:
1. Разметил 100 фоток кромки;
2. Посчитал фичи после глобального пуллинга из сетки resnet-152 с весами от imagenet11k_places365;
3. Взял среднее от фичей для двух классов как анкор;
4. Посчитал расстояние от этого анкора до всех фичей оставшихся 50к фоток;
5. Топ-300 по близости релевантны положительному классу, топ500 самых далеких — отрицательный класс;
6. На этих сэмплах обучил LightGBM на тех же фичах (на картинке указан XGboost, потому что у него есть лого и он более узнаваем, а у LightGBM лого нет);
7. С помощью этой модели получил разметку всего датасета.
Примерно такой же подход я использовал в kaggle соревнованиях в качестве бейзлайна.
Небольшое лирическое отступление
В ODS недавно жаловались, что никто не пишет про свои фейлы. Исправляю ситуацию. Примерно год назад я участвовал в соревновании Kaggle Sea Lions вместе с Евгением Нижибицким. Задача состояла в подсчете морских котиков на снимках с беспилотника. Разметка была дана просто в виде координат тушек, но в какой-то момент Владимир Игловиков разметил их ббоксами и великодушно поделился этим с сообществом. На тот момент я считал себя батей семантической сегментации (после Kaggle Dstl) и решил, что Unet сильно облегчит задачу подсчета, если научиться классно выделять котиков.
В ванильном виде, правда, ее никто уже не использует, а как минимум добавляют Batch Norm. Ну и как правило берут жирный энкодер и раздувают декодер. Еще на смену U-net-like архектурам сейчас пришли новомодные FPN сегментационные сетки, которые показывают хороший перфоманс на некоторых задачах. Однако Unet-like архитектуры и по сей день не потеряли актуальность. Они хорошо работают как бейзлайн, их легко обучать и очень просто варьировать глубину / размер нейрости меняя разные экнкодеры.
Соответственно, я начал обучать сегментацию, имея в качестве таргета на первом этапе только ббоксы котиков. После первой стадии обучения я предсказал трейн и посмотрел, как выглядят предикты. С помощью эвристик можно было выбрать из маски абстрактный confidence (уверенность предсказаний) и условно разделить предсказания на две группы: где все хорошо и где все плохо.
Предсказания, где все хорошо, можно было использовать для обучения следующей итерации модели. Предикты, где все плохо, можно было выбрать с большими площадями без котиков, занулить руками маски и тоже докинуть в трейн. И так итерационно мы с Евгением обучили модель, которая научилась даже сегментировать плавники морских котиков для больших особей.
Но это был лютый фейл: мы потратили кучу времени на то, чтобы научиться круто сегментировать котиков и… Это почти никак не помогло в их подсчете. Предположение о том, что плотность котиков (количество особей на единицу площади маски) постоянна, не работало, потому что беспилотник летал на разной высоте, и фотки имели разный масштаб. И при этом, сегментация все равно не выделяла отдельных особей, если они лежали плотно — что происходило довольно часто. А до инновационного подхода по разделению объектов команды Tocoder’ов на DSB2018 оставался еще год. В итоге мы остались у разбитого корыта и финишировали на 40 месте из 600 команд.
Однако я сделал два вывода: семантическая сегментация — это удобный подход для визуализации и анализа работы алгоритма, а из ббоксов при некотором старании можно сварить маски.
Но вернемся к пицце. Чтобы на отобранных и отфильтрованных фотографиях выделить коржик, самым правильным вариантом было бы отдать задачу разметчикам. На тот момент у нас уже был реализованы ббоксы и алгоритм консенсуса для них. Поэтому я просто накидал пару примеров и отдал на разметку. В итоге я получил 500 сэмплов с точно выделенной областью коржика.
Затем я откопал свой код с котиков и более формально подошел к теперешней процедуре. После первой итерации обучения было так же хорошо видно, где ошибается модель. И уверенность предсказаний можно определить так:
1 — (площадь серой зоны) / (площадь маски) # тут будет формула, обещаю
Теперь, чтобы сделать следующую итерацию натягивания ббоксов на маски, небольшой ансамбль предскажем трейн с ТТА. Это можно считать в некоторой степени WAAAAGH knowledge distillation, но правильней называть Pseudo Labeling.
Дальше нужно глазами выбрать некоторый порог уверенности, начиная с которого мы формируем новый трейн. И опционально можно разметить самые сложные сэмплы, с которыми ансамбль не справился. Я решил, что это будет полезно, и раскрасил где-то 20 картинок, пока переваривал обед.
А теперь финальная часть пайплайна: обучение модели. Чтобы подготовить сэмплы, я экстрактнул область коржика по маске. Также я немного раздул маску диляцией и применил ее к картинке, чтобы удалить фон, поскольку там не должно быть информации о качестве теста. А затем я просто зафайнтюнил несколько моделей из зоопарка Imagenet. Всего у меня получилось собрать порядка 12к уверенных сэмплов. Поэтому я обучал не всю нейросеть, а только последнюю группу сверток, чтобы модель не переобучалась.
Лучшей сингл моделью оказалась Inception-Resnet-v2, и для нее ROC-AUC на одном фолде составил 0.700. Если ничего не выделять и подавать сырые картинки как есть, то ROC-AUC составить 0.58. Пока я разрабатывал решение, у DODO пиццы сварилась следующая партия данных, и можно было протестировать весь пайплайн на честном холдауте. Мы проверили на нем весь пайплайн и получили ROC-AUC 0.83.
Давайте теперь посмотрим на ошибки:
Топ False Negative
Тут видно, что они связаны с ошибкой разметки коржика, поскольку явно присутствуют признаки испорченного теста.
Топ False Positive
Здесь ошибки связаны с тем, что первой моделью выбран не очень удачный ракурс, по которому сложно найти ключевые признаки качества теста.
Заключение
Коллеги иногда подкалывают меня, что я многие задачи решаю сегментацией с использованием Unet. Однако, на мой взгляд, это довольно мощный и удобный подход. Он позволяет визуализировать ошибки модели и уверенность ее предсказаний. Кроме того, весь пайплайн выглядит очень просто и сейчас есть куча репозиториев под любой фреймворк.
Автор: Артур Кузин