Всем привет! В прошлой статье мы рассказывали, как можно ускорить процесс разметки данных с помощью интерактивной сегментации, и уже упоминали state-of-the-art-решение в этой области — модель Segment Anything. Сегодня остановимся на том, как можно улучшить качество и производительность SAM: научить модель генерировать более детализированные и гранулярные маски, а также ускорить её работу в 50 раз и адаптировать для мобильных устройств. А в следующей части на примерах покажем, что способности SAM не ограничиваются одной только сегментацией: модель может применяться для решения самых разных задач компьютерного зрения.

Что такое Segment Anything (SAM)? 

SAM (Segment Anything Model) — это сегментационная модель, которая была выпущена Meta AI*  весной 2023 года и быстро стала одной из самых популярных AI-моделей. SAM называют первой фундаментальной моделью в компьютерном зрении и сравнивают с ChatGPT в NLP из-за рекордно большого количества разнообразных данных, которые видела модель (SAM обучалась на датасете SA-1B, содержащем более одного миллиарда масок); а также из-за её способности к zero-shot transfer, то есть способности легко обобщаться для решения смежных задач.

SAM сегментирует объекты на картинке в соответствии с промптом: им может быть точка на изображении, приблизительный прямоугольник или произвольный текст. Авторы отмечают, что при заранее вычисленном эмбеддинге изображения модель способна работать на CPU в браузере в реальном времени, а также может применяться без дообучения для задач, отличных от сегментации (например, для задач детекции границ и положения объектов — edge detection и object proposal).

Высокое качество сегментации и хорошая обобщаемость принесли популярность оригинальной SAM; но можем ли мы ещё улучшить качество модели?

Semantic-SAM

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

Чтобы добиться многоуровневого понимания семантики объектов, модель обучалась на семи датасетах, включая SA-1B. Датасеты содержат сегментационные маски с разными уровнями детализации: некоторые из них содержат аннотации только на уровне объекта (например, COCO), а другие — на уровне объектов и частей объектов (как Pascal Part).

Semantic-SAM способна выполнять как обычную сегментацию, так и сегментацию на основе промпта. В качестве промпта, как и в оригинальном SAM, может выступать текст, клики и ограничивающие прямоугольники (bounding box’ы). Так как промпты могут относиться к разным уровням одного и того же объекта (например, одна и та же точка может подразумевать выделение одной части тела или всего человека на фото целиком), то каждый прямоугольник или клик кодируется одним позиционным эмбеддингом и 6 эмбеддингами контента, относящимся к 6 уровням детализации объекта. 

Авторы пришли к выводу, что оригинальная SAM не может генерировать качественные сегментационные маски на разных уровнях детализации из-за того, что во время обучения в SAM используется сопоставление масок «многие-к-одному». То есть несколько предсказанных масок будут сопоставлены с одной ground-truth маской, как на картинке выше. Поэтому, чтобы модель умела устанавливать соответствие между несколькими предсказанными масками и ground-truth масками объекта, во время обучения в Semantic-SAM используется схема сопоставления «многие-ко-многим» с помощью венгерского алгоритма. Такой подход гарантирует, что одна точка может генерировать предсказания на нескольких уровнях детализации. 

HQ-SAM

У оригинального SAM есть существенный недостаток: для объектов, имеющих сложную структуру, сгенерированные маски получаются неточными. Если в структуре объекта присутствует много тонких линий, SAM часто выделяет их как единое целое, не выделяет вообще или путает объект с фоном, как на примерах ниже. Улучшением качества сегментации сложных объектов решили заняться авторы модели HQ-SAM.

Так как дообучение существующего декодера SAM или добавление в модель нового декодера приводили к ухудшению zero-shot-перформанса, авторы решили добавить лишь небольшое количество новых параметров (<0,5% от всех параметров модели). Всего в HQ-SAM представлены два новых ключевых компонента: токен высокого качества (HQ-Output Token) и блок слияния глобальных и локальных признаков (Global-Local Feature Fusion). Во время обучения все оригинальные веса SAM фиксируются и обучаются только параметры новых слоёв: HQ-Output Token, связанный с ним трёхслойный MLP-блок, а также Fusion-блок.

Рассмотрим подробнее два основных изменения в структуре SAM:

• High-Quality Output Token (HQ-Output Token). Новый токен соединяется с выходными и промпт-токенами и подаётся на вход в декодер масок. Аналогично другим выходным токенам, HQ-токен сначала выполняет процедуру self-attention с другими токенами, а затем token-to-image и обратный image-to-token attention.

• Global-local Feature Fusion. Как понятно из названия, в этом блоке происходит слияние глобальных и локальных признаков. Признаки с разных слоёв SAM объединяются для создания новых HQ-признаков, которые будут использоваться для предсказания масок высокого качества. Вследствие объединения сгенерированные HQ-признаки обладают как глобальным семантическим контекстом, так и детализированной информацией о локальных границах объектов.

Для обучения HQ-SAM был собран новый сегментационный датасет HQSeg-44k, содержащий 44 тысячи сегментационных масок высокой точности. Для создания нового датасета были размечены изображения из шести существующих датасетов, суммарно охватывающих около 100 различных семантических классов.

Так как все веса оригинальной модели зафиксированы, дообучать необходимо только добавленный HQ-токен и Fusion блок. Как следствие, обучение HQ-SAM занимает всего 4 часа на 8 GPU. С улучшением качества предсказаний разобрались, но как улучшить производительность?

Как ускорить SAM (FAST SAM)

Авторы статьи Fast Segment Anything подсветили существенный минус оригинального SAM: большая вычислительная сложность из-за использования трансформерного энкодера ограничивает широкое применение модели. По этой причине было решено заменить тяжёлый энкодер ViT в SAM на более легковесный свёрточный детектор. Благодаря этому удалось добиться ускорения в 50 раз, сохранив качество предсказаний на сравнимом уровне.

Метод FastSAM состоит из двух этапов: сегментация всех объектов на изображении и выбор объекта согласно промпту.

• На этапе сегментации всех объектов применяется архитектура YOLOv8-seg. Изображение проходит через свёрточный backbone для извлечения признаков, после чего подаётся на вход в Feature Pyramid Network (FPN), чтобы объединить признаки разных размерностей. Полученные карты признаков далее подаются в две параллельные ветки: для детекции и для сегментации объектов. На выходе с ветки детекции мы получаем предсказания категорий объектов и соответствующие им bounding box’ы, а на выходе с ветки сегментации — информацию о сегментационных масках объектов. Объединение этих результатов позволяет методу YOLOv8-Seg одновременно обнаруживать и выделять все объекты на изображении.

• На втором этапе метода FastSAM промпт от пользователя применяется для выделения интересующих нас объектов. Как и в оригинальном SAM, промптом может быть точка на изображении, прямоугольник или текст.

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

Если в качестве промпта подаётся прямоугольник, используется метрика Intersection over Union (IoU) для сопоставления выбранного прямоугольника с bounding box’ами из различных масок.

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

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

SAM на мобилках (FASTER SAM)

Что если мы захотим запустить SAM на мобильном устройстве? Для этого существует Faster SAM — ещё более легковесная версия SAM, созданная специально для работы на низкопроизводительных девайсах типа мобильных телефонов.

В оригинальном SAM основную вычислительную сложность представляет энкодер: в нем 600М параметров против 4М в декодере, поэтому, чтобы сделать модель более легковесной, менять следует именно энкодер. В MobileSAM авторы заменили ViT-H, облегченным энкодером ViT-S, уменьшенной версией ViT. ViT-S имеет меньше параметров и более эффективен в вычислительной скорости, что делает его более подходящим для мобильных приложений.

Однако трансформерный энкодер и декодер масок в оригинальном SAM связаны — они обучаются вместе. По этой причине простая замена ViT-H на ViT-S приводит к неудовлетворительной производительности. Трудности возникают ещё и из-за того, что оптимизация энкодера изображений зависит от качества декодера и наоборот. Когда оба модуля учатся с нуля, их обучение до хорошего состояния становится более сложной задачей. Чтобы решить эту проблему, авторы предложили метод раздельной дистилляции для обучения нового SAM. Весь процесс дистилляции при этом делится на два этапа: дистилляция энкодера изображений и дообучение декодера масок.

Дистилляция — это техника, используемая для обучения новой модели путём передачи знаний из предварительно обученной модели. В контексте MobileSAM раздельная дистилляция используется для обучения легковесного энкодера изображений, который совместим с декодером масок в исходной SAM. При этом происходит обучение отдельной модели, имитирующей поведение исходного энкодера изображений в SAM.

MobileSAM более чем в 60 раз меньше по размеру и при этом имеет качество, сравнимое с оригинальным SAM. Кроме того, MobileSAM в 4 раза быстрее и в 7 раз меньше, чем упомянутый ранее FastSAM. Эти усовершенствования делают MobileSAM подходящим для мобильных устройств с ограниченными вычислительными ресурсами и объёмом памяти: MobileSAM можно обучить на одном GPU менее чем за день, а инференс одной картинки занимает 10 мс.

Заключение

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

В следующей части статьи мы расскажем о способностях SAM адаптироваться под самые разные задачи компьютерного зрения с минимальными изменениями (или вообще без них).

Подписывайтесь на канал команды RnD CV, чтобы оставаться в курсе последних новостей компьютерного зрения и отслеживать успехи нашей команды.

Авторы

• Александр Капитанов @hukenovs

• Нагаев Александр @nagadit

• Елизавета Петрова @kleinsbotle

Ссылки

1. Segment Anything

2. Semantic-SAM: Segment and Recognize Anything at Any Granularity

3. Segment Anything in High Quality

4. Fast Segment Anything

5. Faster Segment Anything: Towards Lightweight SAM for Mobile Applications

*Решением суда запрещена «деятельность компании Meta Platforms Inc. по реализации продуктов — социальных сетей Facebook и Instagram на территории Российской Федерации по основаниям осуществления экстремистской деятельности».