Segment Anything

Разбор статьи

Проблема и мотивация

Базовые модели (foundation models), предобученные на веб-масштабных данных, произвели революцию в NLP и согласовании зрения и языка (например, CLIP), однако в общем компьютерном зрении всё ещё не хватало фундаментальной системы, способной к широкому zero-shot обобщению на различные задачи сегментации. Существующие датасеты для сегментации изображений малы по сравнению с массивами пар изображение-текст из сети, а классические архитектуры сегментации, как правило, привязаны к фиксированным семантическим классам или специализированным интерактивным сценариям.

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

1. Задача: Какая целевая функция обеспечивает универсальный zero-shot перенос?
2. Модель: Какая архитектура поддерживает гибкое задание промптов в реальном времени, разрешая при этом внутреннюю визуальную неоднозначность?
3. Данные: Как получить веб-масштабные маски сегментации, если в интернете отсутствуют плотные разметки масок?

Ключевые идеи

• Promptable Segmentation Task (задача сегментации по промптам): Модель должна возвращать валидную маску сегментации по любому входному промпту — будь то точки переднего/заднего плана, bounding box, плотные маски или произвольный текст. При неоднозначности промпта (например, одна точка на рубашке) модель должна возвращать несколько правдоподобных иерархических масок (деталь, часть, целое).
• Асимметричная амортизированная архитектура (SAM): Тяжёлый энкодер изображений запускается один раз для каждого изображения и формирует эмбеддинг, в то время как легковесный энкодер промптов и двухслойный декодер на основе Transformer обрабатывают промпты за $\sim 50\text{ мс}$ на CPU прямо в веб-браузере.
• Flywheel Data Engine: Трёхэтапный пайплайн разметки данных:
  1. С ручным ассистированием (assisted-manual): SAM помогает аннотаторам в браузере, сокращая время разметки одной маски с 34 до 14 секунд (собрано 4,3 млн масок).
  2. Полуавтоматический (semi-automatic): SAM детектирует маски с высокой уверенностью, направляя внимание аннотаторов на неразмеченные, менее заметные объекты (собрано ещё 5,9 млн масок).
  3. Полностью автоматический (fully automatic): SAM получает на вход равномерную сетку точек $32\times 32$ вместе с приближенными кропами изображений, генерируя 1,1 миллиарда масок в составе SA-1B.

Метод

Архитектура модели

SAM состоит из трёх компонентов:

1. Image Encoder: Предобученный с помощью MAE Vision Transformer (по умолчанию ViT-H/16), модифицированный с использованием оконного внимания (windowed attention) и четырёх блоков глобального внимания; обрабатывает входы размера $1024\times 1024$ и выдаёт эмбеддинг $64\times 64$ с 256 каналами.
2. Prompt Encoder:
   • Разреженные промпты: Точки и bounding box проецируются в 256-мерные векторы с помощью позиционных кодирований, складываемых с обучаемыми эмбеддингами типов. Свободный текст кодируется с помощью текстового энкодера CLIP.
   • Плотные промпты: Входные маски (в разрешении в $4\times$ меньше исходного) уменьшаются свертками $2\times 2$ со страйдом 2 и поэлементно складываются с эмбеддингом изображения.
3. Mask Decoder: Двухслойный модифицированный декодер Transformer. Он обновляет токены промптов и эмбеддинги изображения за счёт чередования self-attention для промптов и двунаправленного cross-attention (от токенов к изображению и от изображения к токенам). Признаки изображения увеличиваются в $4\times$ раза с помощью транспонированных свёрток, а MLP проецирует выходной токен для получения динамического линейного классификатора, применяемого поточечно к увеличенным признакам.

Учёт неоднозначности и функция потерь

Один промпт часто описывает несколько вложенных структур. SAM разрешает неоднозначность, предсказывая одновременно 3 маски с соответствующими оценками IoU. Во время обучения обратное распространение ошибки применяется только к минимальному значению потерь среди кандидатов: $${\mathcal{L}}_{\text{mask}}=\underset{i\in \{1,2,3\}}{\min }\left(20\cdot {\mathcal{L}}_{\text{focal}}(m_i,y)+{\mathcal{L}}_{\text{dice}}(m_i,y)\right)$$ Вспомогательная голова минимизирует среднеквадратичную ошибку между предсказанным и реальным IoU: $${\mathcal{L}}_{\text{IoU}}=\Vert {\widehat{\text{IoU}}}_i-\text{IoU}(m_i,y){\Vert }^2$$ При передаче нескольких промптов неоднозначность снимается, и SAM выводит одну маску, используя четвёртый выделенный выходной токен.

Результаты

• Zero-Shot перенос от точек к маске: На 23 разнообразных датасетах сегментации (включая эгоцентрические, подводные, аэро- и микроскопические снимки) SAM превосходит RITM на 16 из 23 датасетов при использовании лишь одной точки в центре объекта (разница достигает $+46.9$ IoU). При устранении неоднозначности оракулом-селектором среди 3 предсказаний SAM обходит RITM на всех 23 датасетах. Экспертные оценки качества пользователями стабильно оценивают маски SAM существенно выше (от 7 до 9 из 10 баллов).
• Zero-Shot детекция границ (BSDS500): При оценке в режиме zero-shot с сеткой точек $16\times 16$, фильтрацией Собеля и edge NMS модель SAM достигает ODS 0,768 и R50 0,928, конкурируя с контролируемыми моделями вроде HED (ODS 0,788, R50 0,923).
• Zero-Shot генерация пропозалов объектов (LVIS v1): SAM достигает 59,3 AR@1000 без надзора по классам, превосходя полностью контролируемый бейзлайн ViTDet-H на средних объектах (81,6 против 80,8), крупных объектах (86,9 против 87,0) и редких категориях (65,8 против 58,3).
• Zero-Shot инстанс-сегментация: При использовании bounding box от ViTDet в качестве промптов SAM достигает 46,5 AP на COCO и 44,7 AP на LVIS. В слепых тестах с участием людей разметчики стабильно называют маски SAM более чёткими и качественными, чем результаты полностью обученного ViTDet.

Ограничения и открытые вопросы

• Мелкие детали и сильный зум: SAM может упускать ультратонкие структуры или создавать мелкие несвязные фрагменты по сравнению с вычислительно тяжёлыми интерактивными методами, которые явно масштабируют фрагменты изображений.
• Задержка полного пайплайна: Хотя декодирование промптов работает в реальном времени ($\sim 50\text{ мс}$), выполнение тяжёлого энкодера изображений ViT-H не укладывается в режим real-time.
• Семантическое понимание: SAM генерирует немарочные (class-agnostic) маски; архитектура не предоставляет семантических или паноптических меток классов «из коробки» без подключения внешних классификаторов.
• Промптинг вида «текст в маску»: Пайплайн запуска по тексту через CLIP пока носит экспериментальный характер и более чувствителен к формулировкам по сравнению с геометрическими промптами.