Segment Anything Model 2

SAM

Давайте для начала тезисно вспомним архитектуру первой версии модели (подробнее о ней мы рассказывали в нашем посте 🙂).

Для любого изображения мы один раз прогоняем image encoder и получаем image embedding. В качестве энкодера используются ViT. Если есть какая-то маска с предыдущих итераций интерактивной сегментации — мы прогоняем её через несколько свёрточных слоёв и складываем с image embedding’ом. Для промпта в виде точек, боксов или текста мы прогоняем свои лёгкие энкодеры и вместе с image embedding’ом передаём в mask decoder.

В mask decoder мы применяем self-attention для промпта и cross-attention в двух направлениях (от изображения к промпту и от промта к изображению для всех эмбеддингов), чтобы обновить image embedding и prompt embedding. Затем с помощью трёх MLP (как в MaskFormer) и апскейла обновлённого image embedding’а получаем предсказание трёх масок. Ещё один MLP предсказывает confidence score для каждой из масок.

Предсказание трёх масок необходимо для сегментации неоднозначных сцен — в них по одному клику мы не поймём, какой именно объект ищем:

Лосс считается по маске с минимальной ошибкой, а confidence scores учатся предсказывать IOU между предсказанной и истинной масками. Подробнее рассказываем о магии вывода маски в колабе и уже упомянутом посте.

SAM 2

SAM 2 — обобщение архитектуры SAM, которое должно хорошо работать на изображениях и видео. При этом новой модели необходимо получать промпты из точек, масок и боксов на отдельных фреймах и выделять этот объект на протяжении всего видео. Именно поэтому вторая версия схожа с первой — например, prompt encoder и mask decoder значительно не изменились.

Однако в архитектуре второй версии есть несколько ключевых улучшений:

• замена ViT-энкодера на Hiera и изменения в mask decoder’е;
• адаптация архитектуры к работе с последовательностями кадров за счёт трёх новинок — Memory Encoder, Memory Bank и Memory Attention Module.

Замена ViT-энкодера

Каждый кадр видео прогоняется через image encoder один раз. Hiera обрабатывает кадры видео быстрее ViT в 2 раза, если сравнивать похожие по метрикам модели. Ещё один интересный факт — самая большая модель Hiera в SAM 2 весит 850MB, тогда как аналогичный ViT encoder в SAM весит 2.4.GB (хотя по бенчмаркам Hiera получше).

В статье “Hiera: A Hierarchical Vision Transformer without the Bells-and-Whistles” авторы показывают, как сделать быструю архитектуру на основе vision transformer’а без применения трюков вроде «вставим как первые слои сетки». Для этого они обращаются к иерархической архитектуре:

• на первых слоях сети не извлекаем сложные признаки, следовательно, уменьшаем число каналов, но сохраняем высокое пространственное разрешение (ширину и высоту);
• на последних слоях увеличиваем число каналов и уменьшаем разрешение, потому что нам нужны более сложные фичи.

Всего Hiera encoder делится на 4 stages, каждая из которых представляет собой блок, похожий на блок ViT, но с несколькими модификациями и особым соотношением каналов и разрешения. Выходы из stages 3-4 передаются как image embedding в memory attention, потому что имеют больше высокоуровневых фичей, и затем отправляются в mask decoder.

При этом у нас имеется skip connection для фичей со stages 1-2 до transposed convolutions. Они используются для апскейла перед получением маски. Объяснить это можно тем, что у фичей с первых stages более высокое пространственное разрешение. Так мы можем получать более точные границы масок. Похожий трюк использовался в разных кейсах, например, в HQ-SAM, где повышали точность сегментации границ и мелких деталей.

Ещё одно важное улучшение связано с видео. Поскольку объекты в нём могут появляться и пропадать, в mask decoder добавили отдельную occlusion head. Она выдаёт скор отсутствия объекта в кадре. Если скор выше заранее установленного порога — мы не предсказываем объект на этом кадре и отправляем occlusion embedding в memory bank. Эта голова учится с cross-entropy loss. Её задача — предсказать положение объекта в кадре / его отсутствие.

Кроме того, мы сохраняем список из object pointers («obj ptr» на рисунке выше) — небольших векторов, которые вместе с memory bank используются в memory attention для обновления image embedding’а перед попаданием в mask decoder.

Memory Encoder, Memory bank и Memory Attention Module

Чтобы применять работающие с двумерными данными модели к трёхмерным простраственно-временным данным, нужно как-то кодировать, сохранять и передавать информацию между кадрами.

Для этого используется стек из L (в итоговой версии L = 4) трансформерных блоков, где первый принимает на вход image embedding текущего кадра. Каждый из них выполняет:

1. Self-attention для обработки текущего embedding’а кадра.
2. Cross-attention к памяти прошлых кадров, их промптам (если их вводили для этих кадров), object pointers из Memory Bank.
3. Прогон MLP и получение обновлённой версии image embedding’а, которая затем идёт в Mask Encoder.

Что происходит внутри:

1. Маска на выходе из mask decoder’а проходит через несколько свёрточных слоёв и уменьшается.
2. Уменьшенная версия маски складывается с image embedding’ом (уже готовым из Hiera).
3. Результат ещё несколько раз прогоняется через свёртки, и применяется positional encoding.

Полученный эмбеддинг добавляется в Memory Bank и используются для формирования будущих прогнозов.

Авторы также применили GRU перед добавлением в Memory Bank (хороший помощник в трекинге), но в текущей постановке эффект оказался слабым (заметное улучшение только на LVOSv2). Потому в дальнейшем они решили его не использовать.

Object pointers — небольшие векторы, в которых хранится высокоуровневая информация об объектах для сегментации. Они получены из выходных токенов декодера масок для каждого кадра. Их применение как раз показало себя хорошо 🙂

Он работает по принципу двух FIFO-очередей из пар эмбеддингов и object pointers: одна очередь хранит память последних N кадров без промптов, другая — память M кадров с промптами. Например, если при сегментации видео стартовая маска — единственный промпт, Memory Bank содержит «память» о первом кадре (и хранит в течение обработки всего видео), а также о N последних кадрах без промпта.

Для эмбеддингов последних N кадров в очереди без промптов изпользуется RoPE positional encoding для добавления информации об удалённости во времени. Для очереди M кадров с промптами это не используется, поскольку между двумя кадрами с промптом может пройти очень много времени — будет тяжело собрать датасе, покрывающий все возможные случаи. Поэтому positional encoding будет только путать модель.

Memory Bank — последний компонент архитектуры для обработки видео. Именно он обеспечивает сохранение информации между кадрами.

Датасет

Теперь перейдём к датасету, на котором это всё обучалось 🙂

Когда мы хотим получить модель для сегментации любых объектов на видео, с большинством датасетов для video object segmentation нас ждёт одна из двух проблем:

• недостаточно данных (например, в очень хорошо размеченном DAVIS всего 50 видео);
• недостаточно аннотаций для объектов на фоне (например, проходящих мимо людей) в датасетах с множеством видео (например в YouTube-VOS с его ~6 часами).

Кроме того, части объекта обычно не размечаются как отдельные объекты, а если нам нужно «segment anything» — эти части тоже придётся разметить. А ещё акцент в разметке датасетов на людях, животных и транспорте тоже не играет нам на руку.

Что делать? Конечно же, собирать свой датасет на 200+ часов видео!

Итак, по опыту SAM мы уже знаем, что авторы очень полюбили итеративное улучшение датасета. Как же они собирали датасет на этот раз?

Основные этапы создания датасета

1. Phase 1: SAM. Аннотаторы размечали маску целевого объекта на каждом кадре видео с помощью оригинального SAM и вручную. Каждый кадр обрабатывался как отдельная картинка. На этом этапе получили точную разметку, но он оказался трудозатратным (37.8s на кадр). Таким же способом размечали данные для тестирования, чтобы исключить bias от SAM 2 (который использовался на следующих этапах).
2. Phase 2: SAM 2 Mask. Чтобы перейти ко второй фазе, нужно было обучить SAM 2 Mask — версию модели, способной принимать на вход только промпты-маски (не точки и не боксы, как финальная). Для этого авторы использовали данные из датасета SA-1B (на котором учился первый SAM), opensource датасеты из списка выше и данные с Phase 1. Во второй фазе SAM 2 принимал маски как промпты и применял их на соседние кадры. Аннотаторы корректировали предсказания SAM 2 Mask вручную, после чего модель переобучали с новыми данными, а цикл повторяли дважды. В итоге получилось уменьшить среднее время аннотации кадра в 5 раз.
3. Phase 3: SAM 2. Финальная версия SAM 2 была дообучена для работы с промптами-точками и промптами-боксами (которые можно легко получить из масок). Аннотаторам требовались лишь уточняющие клики для корректировки масок. В сравнении с начальной скоростью разметки ускорение произошло в 8.4 раз, а цикл повторяли 5 раз.

Дополнительные этапы подготовки датасета

1. Верификация данных: отдельная команда разметчиков определяла, стоит принимать данные без изменений или нужно отправлять их на доразметку объекта. Она могла также отметить объект как «not well defined» — в таком случае его убирали из разметки насовсем.
2. Автоматическая генерация масок: SAM 2 получал промпты в виде равномерной сетки точек (так делали и в demo первого SAM, чтобы он обнаружил все объекты). Затем маски отправлялись на верификацию и при прохождении добавлялись в основной датасет, в противном случае — на ручную доразметку. Это повысило количество данных и «плотность» покрытия кадров масками, а заодно помогло найти случаи, где было нужно дообучить модель через ручную доразметку.

В итоге получился самый большой (как по длительности, так и по количеству масок) датасет для задачи сегментации объектов на видео! 😎

Результаты

На рисунке ниже приведено сопоставление SOTA-решений для video object segmentation. Они сравниваются по стандартным для задачи метрикам \( \textit{J}\&\textit{F} \) ( \(J \) — Jaccard index, он же IOU, \( F \) — F-score для пикселей границы, метрика — их усреднение, подробнее можно прочитать тут) и \( \textit{G} \) (то же самое, но с усреднением по категориям, подробнее тут). Неудивительно, что SAM 2 побеждает во всех случаях 😄

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

Отметим ещё один интересный момент — в работе над картинками результаты SAM 2, всё-таки, не всегда превосходят результаты первой версии. Например, в статье “Is SAM 2 Better than SAM in Medical Image Segmentation?” авторы приходят к выводу, что SAM качественнее справляется с сегментацией объектов на низкоконтрастных изображениях (КТ и ультразвук), а на более контрастных изображениях SAM и SAM 2 примерно равны. Хотя SAM обрабатывает каждый срез исследования как отдельную картинку, а SAM 2 — как видео.

Возможно, это связано с весомым отличием доменов видео с камер и медицинских снимков — как минимум голова окклюзии в медицинских данных не особо нужна. На то, что эта проблема не самой архитектуры указывает ещё и то, что при внесении небольших изменений в выбор векторов в Memory Bank и дообучение на новом домене результаты второй версии оказались лучше. Об этом подробнее можно прочитать в статье “Medical SAM 2: Segment medical images as video via Segment Anything Model 2”.

Заключение

Мы рассмотрели, как Segment Anything эволюционировала во второй версии, чтобы преодолеть ограничения первой в задаче сегментации объектов на видео. Перечислим основные изменения:

• Memory Encoder, Memory bank и Memory Attention Module позволяют передавать информацию между кадрами;
• замена ViT encoder на более быструю Hiera помогла быстрее обрабатывать видео, при этом доступ к фичам со stages 1-2 в mask decoder повысил точность сегментации;
• новый датасет под задачу video object segmentation, для сбора которого применили итеративную (и интерактивную :)) разметку с помощью обучаемой модели.

В итоге получили универсальную модель, которая эффективно работает как с картинками, так и с видео в задачах интерактивной сегментации, упрощая процесс разметки для ещё большего числа задач в CV.

Полезные ссылки

• SAM 2: Segment Anything in Images and Videos — оригинальная статья с большими подробностями про сбор датасета, архитектуру и обучение (вплоть до указания для каждого этапа количества масок на изображении, использования scheduler с learning rate, публичных датасетов для zero-shot валидации);
• SAM 2 Demo — демонстрация работы новой модели;
• Hiera: A Hierarchical Vision Transformer without the Bells-and-Whistles — статья про новый энкодер, который работает быстрее ViT.