Автор: Александр Лекомцев Редактура: Тимур Фатыхов, Андрей Шадриков

Рисунок 1. Новая задача, новая модель, новый датасет — всё в одной статье!

Задача статьи — сделать zero-shot и few-shot сегментацию простой и доступной. Для начала давайте вспомним базовый вариант задачи интерактивной сегментации:

1. Прогоняем картинку через сеть и получаем маску (как при обычной сегментации);
2. Пользователь кликами или выделением bounding box указывает модели на еще не выделенные маской части объекта (или выделенные зря);
3. Подсказка от пользователя (prompt) кодируется и передается в сеть;
4. Сеть прогоняется еще раз уже с промптом от пользователя во входных данных;
5. Снова выполняем действия из 2-го пункта до тех пор, пока маска не окажется достаточно точной. На 3-м шаге добавляем к исходному изображению и подсказке маску сегментации с предыдущего цикла.

В статье же в качестве prompt также используется текст. Сама по себе идея подавать текст на вход для zero-shot сегментации появилась раньше, например, в “Image Segmentation Using Text and Image Prompts” уже использовали эмбеддинги CLIP.

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

А как именно мы кодируем и передаем в сеть три вида промптов: текст, клики и рамки?

Prompt encoding

Рисунок 2. Архитектура Segmentat Anything Model (SAM)

В отличие от работы, которую мы рассмотрели в прошлый раз, здесь клики и боксы кодируются не в виде матрицы, а как наборы точек. Для кликов — это координаты и бинарная метка (объект или фон (x, y, is_foreground)), для боксов — координаты углов (x1, x2, y1, y2).

При этом, к каждому набору координат применяется “positional encoding” (см. Рисунок 3). Элементы матрицы B в формуле сэмплируются из нормального распределения. По сути мы представляем точку распределением точек. Более подробное объяснение можно прочитать в статье.

Рисунок 3. Positional encoding для координат: B — матрица из случайных величин с нормальным распределением и нулевым матожиданием, v — вектор координат

После positional encoding вектор суммируется с выученным эмбеддингом для каждого типа. Для извлечения эмбеддинга из текста используется энкодер из CLIP. Все эти чудеса подаются напрямую в mask decoder (см. Рисунок 4).

Рисунок 4. Архитектура SAM без ViT image encoder (на инференсе после отклика пользователя запускается именно эта часть)

С маской из предыдущего шага все еще проще — прогоняем ее через несколько сверточных слоев и получаем тензор того же размера, что и эмбеддинг изображения. В коде это 256x64x64. Затем эмбеддинги маски и изображения поэлементно складываются. Если это первый прогон — мы данный шаг пропускаем. Идея, кстати, почти без изменений пришла еще из RITM.

Mask Decoder