Предсказание границ объектов в задаче семантической сегментации

Предсказание границ объектов в задаче семантической сегментации

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

Давайте рассмотрим пример c граундовой сегментацией лошади и сегментацией от BMask R-CNN. По метрике стандартной Intersection over Union (IoU), они совпадают на 92%. Но если посмотреть на глаз, то это не тянет на 92%.

Если мы сравним BMask с PointRend, то второе как раз больше похоже на исходную разметку. Но разница между ними по IoU всего 5%, а на вид так не скажешь.

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

Сразу спойлер: если мы воспользуемся метрикой, предложенной в этой статье, то разница будет более ощутима. Например, последнее предсказание похоже на голд на 91%. Это хороший результат, хотя какие-то отдельные участки не распознались.

Boundary IoU

Чтобы подсчитать эту метрику, мы будем работать с границами. Есть голдовая граница и предсказанная.

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

Дальше мы наложим две границы, посчитаем пересечение и объединение между ними. Так мы и посчитаем наш Boundary для Interception over Union. Это отношение границы пересечения и объединения границ.

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

Как выбирать расстояние d? Для популярных датасетов авторы прикинули и поняли для себя, что d=15 — хорошо. Это позволяет проиллюстрировать разницу между предсказанием границ разных методов. Для своего кастомного датасета это можно подобрать эмпирически.

Вот мы натренировали два метода. Если у них нет разницы, или разница небольшая по простому IoU, то можно измерить ещё и границу, если это важно для задачи. Так мы увидим, насколько широкую границу нужно брать. Если у вас объекты, например, 30 пикселей в ширину, то, наверное, граница в 15 пикселей — многовато. С ней покроется весь объект, и эта метрика ничем не будет отличаться от простого IoU. Исходя из таких соображений, надо выбирать этот параметр.

Обзор статей

Мы нашли несколько статей, которые помогут лучше предсказывать границы, когда делали обзор.

Есть одна группа методов, где мы просто берем лучшую сегментацию, а лучшая сегментация — это чаще всего на трансформерах. Либо это методы, которые уделяют внимание непосредственной границе. К некоторым методам есть код, к некоторым методам в лучших практиках ресёрча есть только инференс-код модели.

Но это не помешало нам использовать и рассмотреть три первых метода, о которых мы и будем говорить.

SegFix

Это алгоритм для пост-процессинга — ему всё равно, какая модель у вас делала семантическую сегментацию. Он старается улучшить предсказания, которые ему подаются на основе предсказанных определенным образом границ.

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

Именно поэтому здесь тренировка отделена от тестирования. Однажды обученный SegFix работает с предсказаниями любой модели на датасете.

Сама модель SegFix довольно простая. Берется бэкбон, из него получается карта признаков, затем используется две ветки: одна, чтобы предсказать границу, вторая — чтобы предсказать направление внутрь объекта. Вторая чуть более хитрая, но авторы здесь используют классификацию на К классов, позже про неё расскажу.

Сами ветки простые: пара сверток и финальный предиктор. Возникает вопрос: как натренировать такую модель сегментации на границу и направления?

Мы берем сегментационную маску Ground Truth, разделяем её на К классов (каждая категория — отдельный канал). Полученные бинарные маски для каждого из К классов мы преобразуем при помощи дистанс трансформа, то есть получаем в каждом ненулевом пикселе расстояние до границы. Чем ближе к границе, тем число меньше, начиная от единицы; ближе к центру оно увеличивается. Затем к полученной карте расстояний мы применяем фильтр Собеля, чтобы получить в каждом пикселе направление внутрь этой маски.

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

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

Поскольку на всех датасетах — на ADE или на Ситскейпе, — у нас размечены все объекты, то мы получаем толщину 10 пикселей для ребер, и они поделены пополам по направлениям.

По поводу карты направлений. При помощи фильтра Собеля мы получаем угол от 0 до 360 градусов. Регрессию не хочется получить, поэтому делим диапазон на 8 или на 4 классов и будем учить обычную сегментацию на К классов. Условно мы берем диапазон от 0 до 45 градусов как движение вверх, от 45 до 90 — как движение вправо вверх и так далее. Получаем финальную маску и просто учим сетку с кросс-энтропией.

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

Авторы предлагают взять направление на карте ребер, и замапить предсказанные классы в сдвиге. Если на карте мы двигаемся вверх, значит, мы берем пиксели, которые находятся сверху от текущего пикселя. И так же аналогично для всех остальных пикселей. Если мы двигаемся по карте направления вниз, то мы мапим это значение в (0; -1). Это значит, что по игреку мы идем вниз, а по иксу никуда не двигаемся, и так для всех классов.

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

Как это выглядит: у нас есть входная карта, лейблов и карта оффсетов. Мы сдвигаемся по карте оффсетов по пикселям и получаем наш рефайнмент. У нас может быть ситуация, когда есть очень толстая граница, поэтому когда мы идем по оффсету, мы не берем пиксель из внутреннего региона изображения, а все ещё берём пиксели с границы.

Если бы здесь была граница три пикселя, то при оффсете в единицу мы будем брать не внутренний регион, а пиксели в границе, которые мы считаем не очень надежными. Поэтому они предлагают значения типа (-1;1) скалировать на некоторый коэффициент, чтобы запрятать пиксели поглубже из внутренней области объекта. В частности, они используют коэффициент, равный двойке.

Что ещё мне не понравилось в этой работе? Они используют Grid_sample для того, чтобы сдвинуть пиксели, а Grid_sample применяется одновременно ко всем пикселям, которые мы его попросим сдвинуть. Из-за этого самые крайние пиксели все-таки меняются не на внутренние, а на пиксели с границы, которые находятся близко к ним. Мне кажется, выгоднее использовать итеративный подход— начинать изнутри, менять на внутренние пиксели и двигаться к границам.

По результатам у них получилось, что SegFix увеличивает метрику во всех сценариях — где-то на 0,5%, где-то на 3% в зависимости от датасета. Практически везде SegFix не ухудшает результат. Поэтому они предлагают использовать его в любой сегментации.

SegFix: результаты

Мы сами тренировали Segfix на СityScapes, и вот результат, который здесь получается. Видно, что некоторые области становятся ровнее, например, машина. Было кривое ребро, после Segfix оно стало более прямым. Также на знаке видно, что Segfix помогает выровнять некоторые ребра на дороге слева и так далее.

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

Вопрос: Можно уточнить про процесс? Мы предсказываем сегментацию, используя две головы — одна предсказывает границы, вторая предсказывает оффсеты. Мы берем те пиксели, которые соответствуют границе, для них смотрим пиксели из фичмапы с оффсетами и заменяем для этих пикселей. Берем тот класс, который на сегментационной маске, если сдвинуть пиксель на величину оффсета. Правильно понимаю?

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

Вопрос: Допустим, у меня есть датасет, где изначально граундтруфная разметка была не точной на границах. Я натренировал сеть, она предсказывает также, как и в граундтруф данных. Можно ли надеяться, что уже натренированный Segfix на ситискейпс может зарефайнить?

Ответ: Я думаю, если изначально была сегментация автомобилей, то можно использовать Segfix, который был натренирован на ситискейпс. У них нет применимости к моделям на которые были обучены на кастомных датасетах. Авторы утверждают, что они пытались обучить Segfix одновременно на ADE и на ситискейпсе, который в принципе довольно общий. Его можно использовать в любых поддоменных областях на этих датасетах.

Вопрос: Правильно ли я понимаю, что без переобучения мы границу не можем резко поменять?

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

Gated-SCNN: Gated Shape CNNs for Semantic Segmentation

Следующая статья это Gated shape CNN. У Gated-SCNN будет 2 бранчи в сетке. Как можно улучшить результат сегментации?

Одна из идей — запускать инференс на картинках большого размера. А лучше — тренировать на картинках большого размера.

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

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

Что авторы здесь предложили? Первый стандартный бэкбон пусть называется обычной веткой. Вторая ветка — это шейп-ветка. Информацию надо как-то передавать из обычной в шейп. Как вы понимаете, в шейп-ветке на исходном решении это все выполняется.

Как делается? Есть самая первая свертка без пулинга и без страйда. Они берут самый начальный фич после первого слоя, и естественно немного сворачивали и после этого запускают свой механизм гейтинга, чтобы подмешать фичи с более глубоких слоев и более низкого разрешения. Чтобы это сделать, они просто увеличивают их: сначала сворачивают, потом делают интерполяцию.

Как они подмешивают? Есть Gated convolutional — сверточный слой, он представляет из себя сигмоида от какой-то свертки предыдущих фич, в общем, это какой-то аттеншн. Но автор утверждает, что это альфа-канал. Мы берем фичи с обычного стрима, с шейпстрима, конкатенируем их, сворачиваем и делаем сигмоид. Получаем коэффициенты, домножаем их на фичи шейп-ветки и все это объединяем в ресидуал-блоке.

Фактически, авторы предлагают поэлементный аттеншн, чтобы подмешать фичи из бэкбона в ветку с большим разрешением. Ширина канала здесь 16, каналов мало, потому что на выходе мы будем карту ребер предсказывать. Там один канал, поэтому можно обойтись небольшим числом каналов, но с большим разрешением.

С помощью этих гейтед-слоев они подмешали все необходимые фичи. Дальше по картинке подсчитали карту градиентов через Кани, конкатенировали с обыкновенными фичами перед ASPP, чтобы предсказать сегментацию. По сути они предложили аттеншн-слой, чтобы видеть фичи на более высокой резолюшн.

Gated-SCNN: Cityscapes results

Ниже визуализации трех альфа-каналов. Представлен сам поэлементный аттеншн: в среднем ряду на более ранних слоях, а это на более глубоких слоях (слева направо).

На более ранних слоях видно, что красное — это большее значение, синее — это меньшее. На них же видно, что аттеншн усиливает отклики на ребрах внутри самого объекта. На более поздних слоях заметно, что аттеншн помогает сетке фокусироваться на ребрах математических, либо на ребрах между объектами.

На самых глубоких слоях надо применить воображение. Авторы приводят результат на Cityscapes, получают 82,8 — это средний Intersection Union. Сейчас это где-то треть страницы Cityscapes на семантической сегментации. Интересно, что авторы используют архитектуру DeepLabV3+, и утверждают, что у неё вот такой перформанс.

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

Gated-SCNN: losses

Какие лоссы тут используются? Сегментационный лосс простой: из интересного тут то, что веса для классов обратно пропорциональны их частоте. Для предсказания ребер может использоваться бинарная кросс-энтропия. В общем, чтобы они вносили пропорциональный контрибьюшн, то есть веса, пропорциональные другому классу.

Gated-SCNN: dual task loss

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

Вторая часть — просто дополнительный сегментационный лосс там, где мы хорошо предсказали карту ребер. Под хорошо предсказали мы понимаем отклик больше 0,8 и лосс там, где у нас сетка уверенно предсказывает ребра. Ну и соответственно, где плохо предсказывает, мы используем игнор лейбл.

Что такое дуал таск лосс? Это модуль разницы между предсказанными картами ребер.

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

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

InverseForm: A Loss Function for Structured Boundary-Aware Segmentation

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

Бывает, что мы предсказали то же самое ребро, только сдвинутое на пару пикселей. Или слегка повернутое, или с масштабом ошиблись, что часто бывает в предикшенах. В этих случаях мы получим большой лосс.

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

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

Тут немного сложный пайплайн. Авторы предлагают тренировать определенным образом сетку, через которую будут пускать лосс. то будет Это и будет InverseForm Loss.

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

Если предикшн будет идеальный, то предсказание будет единичной матрицей. Во время тренировки сети мы будем сравнивать выход сетки, которая предсказывает преобразование с единичной матрицей и брать L2-норму. Это будет итоговый лосс.

Осталось только натренировать сетку на голубом фоне, чтобы она умела выдавать преобразования, когда ей дают на вход две карты ребер. Предлагается тренировать с помощью еще одного трансформера (на желтом фоне), с помощью special-трансформера сети. Она нужна чтобы выдавать более реалистичную трансформацию.

Тренировать её можно как на своем датасете, так и на MNIST stroke, который можно найти на сайте PyTorch. В целом, пайплайн тренировки выглядит вот так — сначала тренируем STN, которые бы выдавал коэффициент преобразования: сдвиг, поворот, масштабирование. Желтую сетку сохраняем. Голубую мы вставляем в датасет для тренировки нашей сети, чтобы для каждой карты ребер по ней она сгенерировала какое-то преобразование. И это преобразование мы применяем к карте ребер. И вот уже эти карты ребер мы передаем на вход нашей сети и просим её натренировать, чтобы она выдавала то преобразование, которое мы получили во вторую карту ребер.

Как только мы натренировали, мы уже можем это вставлять внутрь нашего пайплайна.

Авторы также предложили собственный дизайн сети, но нас интересует просто как выпускался лосс. Мы берем карту ребер с голдовой сегментацией, авторы ещё тайлят (применяют tiling), потому что на полной картинке сети сложно предсказывать коэффициенты. И для каждого тайла считают лосс и этот лосс сравнивают с единичной матрицей, считают L2 и распространяют.

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

Вопрос: Допустим, у нас есть тайлы с Ground Truth и с предикшена. Мы заранее обучили сетку предсказывать параметры преобразования, которые мы можем сделать. Что мы делаем дальше, применяем и уже после этого считаем лосс между Ground truth и предикшеном после преобразования?

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

Результаты этой сетки находятся на первой странице ситискейпса, потому что за бейслайн была тоже выбрана сетка с первой страницы. В результате вышел импрувмент в районе 0,5%. Да, метрику они считают не ту, которую мы рассматривали в начале семинара, потому что статьи выходят с разным шагом, но это что-то очень похожее.

InverseForm: A Loss Function for Structured Boundary-Aware Segmentation

Вот несколько примеров, по которым можно понять, насколько хорошо это работает. Есть голд, есть бейслайн с лоссом, который должен хорошо предсказывать границы. Во-первых, границы сильно прямыми не стали, но где-то сегментация улучшилась. Если смотреть с такой точки зрения рваной границы, то они улучшились, потому что предикшен более правильный, и классификатор меньше ошибается.

Bonus

Был вопрос, что лучше: искать методы, которые улучшают границу или просто взять и инвестировать время в прокачку или поиск самой мощной сегментации?

В относительно недавней статье за декабрь 21-го года, была предложена архитектура Mask2Former. Эта сетка, основанная на трансформере, можно посмотреть насколько точно она дает границы.

Слева — классическая картинка, которую все знают. Можно посмотреть, насколько точно обозначено разделение между человеком и фоном. Выглядит как простая попиксельная разметка. Понятно, что есть какие-то конфьюжн регионы, но все равно понятно, что человек — не жесткий предмет, не самый легкий для сегментации.

Тоже самое можно видеть с котиком. Шерсть просто пиксель в пиксель, кроме макушки.

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

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

Q&A

Вопрос: по какому принципу вы выбирали методы про которые вы рассказывали?

Ответ: По принципу наличия кода. Надо было посмотреть много гипотез, найти те, к которым нет кода, — это гораздо сложнее, чем найти те, к которым код уже есть. UnverseForm затесался, потому что интересная идея.

Первые три метода тестируются на известных датасетах, последний не совсем про улучшение границ. Если вы предсказываете сегментацию и даунсемплите картинку, а хотите хорошую сегментацию, то CascadePSP — то, что нужно.

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

Ответ: Если нет проблем, на чем инферить, и если есть время, я бы попробовал Mask2Former. Достаточно дружелюбный код и интересный результат из коробки. Плюс, куча предтренированных моделей, но не на CityScapes. Можно даже потренировать и самостоятельно зафайнтюнить. А если речь о каком-то edge-device с 200 гб весов, то вряд ли затащишь, тут надо добавлять лоссы.

По материалам открытого семинара Павла Сёмкина, Даниила Осокина «Предсказание границ объектов в задаче семантической сегментации»‎, который прошел в Xperience.AI.