Автор: Александр Лекомцев
Редактура: Иван Перминов
Давайте для начала тезисно вспомним архитектуру первой версии модели (подробнее о ней мы рассказывали в нашем посте 🙂).
Рисунок 1. Архитектура первой версии SAM
Для любого изображения мы один раз прогоняем image encoder и получаем image embedding. В качестве энкодера используются ViT. Если есть какая-то маска с предыдущих итераций интерактивной сегментации — мы прогоняем её через несколько свёрточных слоёв и складываем с image embedding'ом. Для промпта в виде точек, боксов или текста мы прогоняем свои лёгкие энкодеры и вместе с image embedding'ом передаём в mask decoder.
Рисунок 2. Архитектура mask decoder'а в SAM
В mask decoder мы применяем self-attention для промпта и cross-attention в двух направлениях (от изображения к промпту и от промта к изображению для всех эмбеддингов), чтобы обновить image embedding и prompt embedding. Затем с помощью трёх MLP (как в MaskFormer) и апскейла обновлённого image embedding'а получаем предсказание трёх масок. Ещё один MLP предсказывает confidence score для каждой из масок.
Предсказание трёх масок необходимо для сегментации неоднозначных сцен — в них по одному клику мы не поймём, какой именно объект ищем:
Рисунок 3. Пример проблемы двусмысленности (ambiguity): это неоднозначная сцена, где можно предсказать и одежду, и манекен
Лосс считается по маске с минимальной ошибкой, а confidence scores учатся предсказывать IOU между предсказанной и истинной масками. Подробнее рассказываем о магии вывода маски в колабе и уже упомянутом посте.
Рисунок 4. Архитектура SAM 2
SAM 2 — обобщение архитектуры SAM, которое должно хорошо работать на изображениях и видео. При этом новой модели необходимо получать промпты из точек, масок и боксов на отдельных фреймах и выделять этот объект на протяжении всего видео. Именно поэтому вторая версия схожа с первой — например, prompt encoder и mask decoder значительно не изменились.
Однако в архитектуре второй версии есть несколько ключевых улучшений:
Каждый кадр видео прогоняется через 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'а без применения трюков вроде «вставим как первые слои сетки». Для этого они обращаются к иерархической архитектуре: