Unsupervised Segmentation

Введение

Сегментация — одна из самых востребованных и важных задач компьютерного зрения. Обычно выделяют три основных ее вида:

• семантическая — попиксельная классификация всего изображения;
• инстанс — сегментация и классификация объектов на изображении, которые могут делиться на подклассы;
• паноптик — объединение инстанс и семантической сегментации.

Одна из ключевых проблем при получении качественной модели сегментации — трудоемкость разметки. В отличие, например, от детекции, где разметка — ограничивающий прямоугольник, здесь аннотатору нужно выделять полигональную маску, что занимает гораздо больше времени. Но стоит отметить: для сегментации обычно требуется меньше разметки, чем для детекции. Объясняется это тем, что сегментация — попиксельная классификация. Каждая маска содержит большое количество размеченных пикселей, каждый из которых точно относится к нужному классу. А в детекции мы решаем задачу классификации объекта в ограничивающем прямоугольнике. Не все пиксели внутри него относятся к классу целевого объекта.

Для решения этой проблемы исследователи искали подходы к сокращению времени разметки или автоматизации ее получения. До недавнего времени были популярны методы упрощенной разметки (weakly supervised): разметка ограничивается только классом объекта на изображении (обычная классификация), какой-либо отметкой на объекте (например, линией) или ограничивающим прямоугольником.

Сегодня самые популярные подходы — подходы, которые вообще не используют разметку. Это self-supervised (SSL) методы предобучения визуальных моделей. В статье Emerging Properties in Self-Supervised Vision Transformers авторы акцентировали внимание читателей на семантических свойствах карт внимания трансформерной модели ViT, обученной на данных без разметки с применением контрастивной функции. Да, речь идет о знаменитом подходе DINO. Его мы, кстати, подробно разбирали в нашей статье 😊. Но DINO дает лишь базу в виде предобучения и выявляет потенциал предобученного трансформера в режиме SSL. Чтобы решить задачу сегментации с приемлемым качеством и попасть в лидерборд популярных бенчмарков, нужно дообучать сегментационную голову модели. И здесь уже появляется пространство для творчества и, следовательно, разнообразие подходов.

В этой статье мы рассмотрим самый актуальный и универсальный метод сегментации без разметки, который находится в топах бенчмарков всех видов сегментации — Unsupervised Universal Image Segmentation, или U2Seg.

Unsupervised Universal Image Segmentation (U2Seg)

Самые популярные подходы в задачах сегментации без учителя — STEGO для семантической сегментации и CutLER для инстанс сегментации. Авторы U2Seg предложили архитектуру, которая умеет решать задачи как инстанс, так и семантической сегментации, а следовательно — задачу паноптик сегментации. Таким образом, Unsupervised Universal Segmentation model (U2Seg) — это универсальная модель сегментации.

В U2Seg объединяются архитектуры STEGO и CutLER, поэтому давайте для начала их и рассмотрим 🙂.

STEGO

Со STEGO мы знакомились в нашей статье про DINO.

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

По сути STEGO — способ получения попиксельного эмбеддинга входного изображения, после кластеризации которого возможно добиться качественной семантической маски. Обучение модели сводится к минимизации скалярного произведения двух косинусных мер — F и S. Где F — косинусная мера между эмбеддингами, полученными с помощью замороженного бэкбона, а S — косинусная мера между эмбеддингами на выходе сегментационной головы (см. рисунок 4).

Лосс функция считается для трех вариантов пар изображений: изображение и его ближайший сосед, изображение и его копия, изображение и случайная картинка из выборки. Следовательно, мы имеем контрастивную функцию потерь, где первые две пары положительные, а третья — отрицательная.

CutLER

CutLER также использует богатые фичи DINO для дальнейшего их файнтюна на задаче детекции и инстанс сегментации. Здесь разработчики придумали способ извлечения масок отдельных объектов (инстансов) с помощью MaskCut.

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

U2Seg

Архитектура и обучение U2Seg

Процесс обучения модели U2Seg показан на рисунке 6:

Обучение начинается с выделения масок объектов (things) и масок всего остального, что окружает объекты (stuff — это например, трава, небо, вода, обочина и др.).

Маски объектов выделяются с помощью MaskCut от CutLER (см. рисунок 7):

Недостаток модели CutLER — предсказанные маски объектов не различаются по классам.  В U2Seg эту проблему решили через кластеризацию масок объектов, которые были получены после MaskCut. Для кластеризации использовали обычный KMeans. Таким образом, каждая маска была отнесена к своему классу. В результате получились псевдолейблы для инстанс сегментации.

Чтобы получить семантические маски окружающей среды (stuff), используется подход STEGO, который дистиллирует признаки предобученной модели DINO в дискретные семантические маски.

Макси instance и stuff объединяются по следующему алгоритму:

1. Проверяются пиксели масок объектов \( I_С(i, j) \). Если они ненулевые — текущему пикселю \( (i, j) \) назначается метка instance;
2. Пиксели, для которых \( I_С(i, j)=0 \), проверяются на ненулевые значения маски окружающей среды \( S_S(i, j) \), назначается метка stuff.

Таким образом, задача семантической сегментации сводится к обучению модели на двух классах: “instance” и “stuff”.

Обучение U2Seg сводится к одновременному обучению двух ветвей сети — инстанс и семантической сегментации. Совместная функция потерь выглядит следующим образом:

где \( L_c \) — классификационный лосс, \( L_b \) — лосс детекции, \( L_m \) — сегментационный лосс, \( L_s \) — семантический лосс (попиксельная кросс-энтропия).

В качестве сегментационной модели используется Panoptic Feature Pyramid Networks. У нее две отдельные ветви для инстанс и семантической сегментации. В качестве бэкбона используется ResNet50, инициализированная весами, которые были предобучены DINO.

Эксперименты

Модель обучалась на 1.3 М неразмеченных картинок из ImageNet. Инстанс сегментация тестировалась на датасетах COCO, PASCAL VOC и UVO, паноптик — на COCO и Cityscapes.

Опираясь на данные таблиц, мы можем отметить: U2Seg обгоняет CutLER+ (CutLER + кластеризацию) в задаче zero-shot инстанс сегментации. “Zero-shot” означает, что U2Seg в процессе обучения не видела датасеты COCO и Pascal VOC, а училась только на ImageNet (без разметки) + кластеризации на 800 классов. Хотя для бейзлайна был взят не простой CutLER, а с дополнительной кластеризацией, это не помогло обогнать U2Seg. Она обучалась не только на инстанс, но и на семантической сегментации, что позволило ей получить больше информации из данных и усилить zero-shot прогноз.

На рисунке 10 видно, как U2Seg обгоняет совместный прогноз CutLER + STEGO (CutLER — инстанс, STEGO — семантическая сегментация) в задаче паноптик сегментации. Метрики PQ (panoptic quality), SQ (segmentation quality), RQ (recognition quality) — метрики для оценки качества паноптик сегментации. “Non zero-shot method” предполагает дообучение модели без учителя на данных из домена COCO, на котором и проводилось тестирование.

U2Seg показывает качественную сегментацию с высокой детализацией по сравнению с объединенным прогнозом CutLER и STEGO. Это объясняется эффективностью подхода с дообучением сегментационных ветвей на псевдолейблах.

На рисунке 10 представлен результат сегментации U2Seg в сопоставлении с фактической разметкой и псевдолейблами. По сути псевдолейблы — объединение прогнозов CutLER и STEGO. Финальный прогноз U2Seg отличается от бейзлайна (CutLER + STEGO) более качественным выделением маски инстансов и окружающей среды. Например, в третьем ряду на рисунке мы видим, как бейзлайн отдельно кластеризовал голову и торс мужчины, а U2Seg объединил эти классы в одну категорию — человек.

Чтобы установить соответствия между предсказанным кластером (cluster ID) и фактической меткой (category ID) из тестовой выборки, в задачах сегментации без учителя используется венгерский алгоритм (Hungarain Matching):

1. По заданному порогу уверенности отсекают объекты с наиболее уверенным прогнозом.
2. Далее считают индекс Жаккара (IoU) для каждого предсказанного объекта (для маски / ограничивающего прямоугольника) со всеми объектами из разметки на текущей фотографии.
3. После выбирают объект, чей IoU выше заданного порога, и для его category ID назначают соответствующий category ID.
4. Затем для каждой category ID строят частотную гистограмму category ID, которые назначались текущей category ID на всем тестовом датасете. И на ее основе выбирают для cluster ID наиболее частотную category ID.

На рисунке 11 представлена интересная визуализация активаций полносвязного слоя CutLER и U2Seg с помощью t-SNE. Мы видим, что признаки U2Seg лучше кластеризуются и образуют хорошо считываемые кластера.

Заключение

Модель U2Seg сегодня — возможно, лучшее универсальное решение для задач сегментации без учителя не сегоднящний день.

Она объединяет в себе решение всех трех задач сегментации — инстанс, семантической и паноптик. Модель уверенно держится в топе популярных бенчмарков по сегментации без учителя и является новым state-of-art подходом для паноптик сегментации.

Но, как и любая универсальная модель, U2Seg может проигрывать моделям, специально обученным под отдельные задачи сегментации. Например, на лидерборде бенчмарака семантической сегментации без учителя на датасете COCO-Stuff-27 модель U2Seg занимает лишь третье место, проигрывая модели CAUSE.

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

1. Unsupervised Universal Image Segmentation — статья про U2Seg
2. Unsupervised Semantic Segmentation By Distilling Feature Correspondences — статья про STEGO
3. Cut and Learn for Unsupervised Object Detection and Instance Segmentation — статья про CutLER