Разбор SAM (Segment Anything Model) 

SAM — это модель от Meta AI, обученная на новом огромном датасете для сегментации изображений. Этот датасет содержит 11 млн картинок и 1.1 млрд масок сегментации. Meta AI описали, как собирали этот датасет, а также выложили его в открытый доступ вместе с моделью.

Сразу все релевантные ссылки:

• Статья на arxiv;
• Сайт с примерами работы;
• Онлайн-демо;
• Визуализация датасета;
• Скачать датасет.

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

Кроме обычного режима "сегментировать все, что есть на картинке" модель еще умеет сегментировать отдельные объекты по запросу пользователя. Запросы могут быть таких видов:

• точка (или несколько точек) на объекте. Ставите мышкой точку на объект, который хотите сегментировать, и модель сегментирует этот объект;
• bounding box объекта. Мышкой выделяете примерный прямоугольник, в котором находится объект, и модель его сегментирует;
• примерная маска объекта. Ручкой как в Paint закрашиваете примерную область, в которой находится объект. И модель эту область уточняет;
• текст-описание объекта. Например, "коричневый мишка в руках у девочки". Модель найдет этот объект и сегментирует.

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

• edge detection;
• instance segmentation;
• object proposal.

Вообще, авторы считают, что SAM может стать foundation model, т.е. что-то вроде GPT-3 в сфере сегментации картинок. Foundation model — это модель, которая обучилась на какую-то задачу на огромном количестве данных, и это породило у нее способности к решению многих других задач, на которые она напрямую не обучалась, в few-shot и zero-shot режимах.

Правда, прочитав статью, я не уверена, что я с этим согласна. На мой взгляд, те задачи, на которых SAM показывает способности к zero-shot — это та же сегментация, только вид подачи данных на вход модели немного отличается. Т.е. не то чтобы тут у модели возникают какие-то новые неожиданные способности. Но SAM — это точно foundation model в том смысле, что обучена модель была на огромном количестве данных, и из-за этого ее должно быть очень легко дообучить под свою конкретную задачу. Это такой ImageNet для задачи сегментации, короче.

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

Устройство модели SAM

Устройство модели — на картинке выше. Модель состоит из трёх частей:

• Image encoder. Это ViT, предобученный с помощью MAE. Эта часть принимает входящую картинку, которую нужно сегментировать, и превращает ее в вектор (эмбеддинг на картинке)
• Энкодер дополнительной инфы (кондишена), на основе которой нужно сегментировать картинку. Это может быть:
• Mask — заданная юзером маска для объекта, который нужно сегментировать. Маска подаётся на вход сверточному энкодеру, который выдает эмбеддинг того же размера, что image encoder. Этот эмбеддинг потом суммируется с эмбеддингом из image encoder;
• Points — набор точек, которые относятся к объекту, который хотим сегментировать. В этом случае точки переводятся в эмбеддинг с помощью positional encoding'а. Этот эмбеддинг затем подаётся на вход третьей части сети вместе с эмбеддингом картинки, которую нужно сегментировать;
• Box — заданный юзером bounding box объекта, который нужно сегментировать. Box также переводится в эмбеддинг с помощью positional encoding. Конкретнее, с помощью positional encoding кодируются координаты его левого верхнего и правого нижнего углов;
• Text. Вот тут интереснее всего, потому что в датасете, на котором обучается SAM, нет текстовых описаний к изображениям. И если все другие виды инфы (points, box и mask) можно легко получить для любой пары (картинка, карта сегментации) датасета и на этом нейросеть обучать, то с текстом так легко не выйдет. Как же авторы добиваются того, чтобы SAM могла сегментировать объект на картинке на основе его текстового описания, разберём ниже.
• Ну и третья составляющая модели — mask decoder. Она принимает на вход эмбеддинг картинки из image encoder и эмбеддинг допинфы. Архитектура — декодер трансформера с парой добавленных нюансов для лучшей обработки информации (об этом подробнее в разделе А аппендикса статьи).

Важно вот что: mask decoder выдает не одну карту сегментации, а три. Сделано это для того, чтобы модель могла эффективно обрабатывать нечёткие случаи: когда, например, юзер поставил точку, которая может относиться к нескольким объектам на картинке. Например, точка на лице человека: она может относиться как к объекту "лицо", так и к объекту "человек". Если модель выдает только одну карту сегментации, она будет стараться на такие нечёткие запросы выдать карту, которая будет чем-то средним между всеми возможными правильными картами сегментации. А если у нее есть возможность выдавать несколько карт, то каждая карта может соответствовать одному возможному объекту. Количество карт 3 было выбрано эмпирически: авторы говорят, что этого в целом хватает для большинства юзкейсов.

Ниже — пример того, как модель генерирует три маски сегментации для одной точки на изображении:

Для каждой выданной карты сегментации сеть также предсказывает confidence этой маски. При обучении сети используется лосс только из самой confident головы.

Датасет для обучения SAM 

Для обучения SAM авторы статьи собрали самый большой датасет пар (картинка, карта сегментации). В нем 11 млн картинок и 1.1 млрд масок. Масок больше, потому что на одной картинке в среднем находится много объектов, и для каждого из них маска своя. Можно считать, что у нас есть 1.1 млрд пар (объект на картинке, маска сегментации).

Еще раз ссылки на датасет:

• Визуализация датасета;
• Скачать датасет.

Интересно то, как они датасет готовили: это было не полностью с помощью людей-разметчиков (представьте, сколько бы времени заняло такой датасет собрать и сколько бы это стоило). Собирали датасет в три стадии:

1. Manual stage. С помощью людей собрали 4.3 млн масок для 120 тысяч картинок;
2. Semi-automatic stage. На данных с пункта 1 обучили нейронку и ей разметили объекты на 180к картинках. Дальше дали эти размеченные картинки людям и попросили доразметить на них те объекты, которые нейронка пропустила. Чаще всего это были мелкие объекты. В течение этой стадии на новых полученных данных нейронку несколько раз дообучали. Это повысило способность сети сегментировать больше объектов на каждой картинке;
3. Fully-automatic stage. Взяли нейронку, обученную на данных из второго шага. Для каждой картинки генерили плотную сетку точек, и просили модель сегментировать объекты, соответствующие этим точкам. Из полученных таким образом карт сегментации отбирали те, в которых модель наиболее уверена (confidence высокий). И далее их ещё улучшали с помощью классического NMS. Так и набрали 11 млн картинок с 1.1 млрд карт сегментации. В итоговом датасете 99% масок сегментации были получены полностью автоматически.

Такой подход позволяет, при желании, нагенерить ещё данных. Датасет полностью выложен в открытый доступ вместе со статьей.

Замечу ещё раз, что в датасет вообще нет никаких текстовых описаний к никаким объектам. И если из пар (объект на картинке, карта сегментации) можно легко для объекта сгенерить допинфу для обучения модели — кривую маску, bounding box или точку на объекте, то текстовое описание хорошо сгенерить не выйдет. Ниже обсудим, как же авторам удалось научить модель генерить карту сегментации для объекта на основе текстового описания.

Генерация маски сегментации объекта на основе текстового описания объекта

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

А идея очень классная, на мой взгляд. Вот какая:

Вспомним про CLIP: эта модель выучивает общее пространство эмбеддингов для картинок и текста.

Из этого возникает такая идея обучения SAM:

У нас есть датасет вида (объект на картинке, маска сегментации для этого объекта). Давайте для каждого такого объекта вырежем прямоугольный кусок картинки с этим объектом и подадим ее на вход CLIP. На выходе получим эмбеддинг картинки, назовем его Х. Так как картинка в целом содержит только наш объект, то эмбеддинг Х будет содержать информацию о том, что на картинке изображен именно этот объект.

Во время обучения SAM будем в качестве псевдо-эмбеддинга описания объекта подавать на вход модели эмбеддинг Х. То есть, CLIP-эмбеддинг куска картинки, содержащей этот объект:

А после обучения, на этапе инференса, будем уже прогонять текстовое описание желаемого объекта через CLIP, и уже его подавать в качестве кондишена. Так как CLIP выучивает общее пространство эмбеддингов для текста и картинок, такая замена эмбеддинга картинки на эмбеддинга текстового описания должна сработать.

Ну и правда работает, вот примеры результатов:

Заключение

Вот такая идея модели. Подробнее о нюансах архитектуры, обучения и метриках читайте в статье.

Буду рада узнать, что вы думаете насчет того, насколько эта модель важна и как повлияет на прогресс в области сегментации и CV в целом =)