CRAFT: метод детекции текста на изображениях
Пререквизиты
Эта статья продолжает вводную статью про OCR и дополняет статью про детекторы текста на трансформерах. Рекомендуем сначала ознакомиться с вводной статьёй, затем прочитать эту, а третьей — статью о детекторах текста на трансформерах (если вдруг вы ещё не знакомились с ней 🙂).
Введение
После появления нейросетевых детекторов ресёрчеры пытались адаптировать их для детекции текста. Однако эта задача имеет ряд особенностей, которые не позволяли решать её успешно в общем случае:
- Мы хотим находить текст произвольной формы. Прямоугольные боксы для этого плохо подходят, так как слова могут просто не попасть в такие жёсткие рамки;
- Текст бывает разного размера, а обычные детекторы имеют проблемы с тем, чтобы хорошо детектировать и маленькие, и средние, и большие объекты.
Для решения первой проблемы исследователи из ClovaAI перешли от классической схемы детекции на уровне слов (word-level detection) к детекции на уровне символов (character-level detection). Но здесь возникла следующая проблема: большинство датасетов содержат разметку только на уровне слов. Привлечь разметчиков — плохая идея, поскольку требует существенных временных и материальных затрат. Авторы статьи предложили простое и элегантное решение — получать character-level разметку из синтетических данных (где она бесплатная) и в процессе обучения переносить её на реальные данные. Подробнее об этом будет сказано в блоке «Weakly-supervised обучение» 🙂
Чтобы решить вторую проблему, авторы используют U-net like архитектуру (то есть сегментацию с последующим получением полигональных боксов). Она позволяет лучше работать с фичами разного масштаба за счёт skip connections.
Архитектура CRAFT
Архитектура CRAFT устроена просто. В качестве бэкбона используется VGG-16. В свою очередь промежуточные слои бэкбона вместе со skip connections и UpConv-блоками образуют U-net часть для сбора локальных визуальных фичей. На выходе мы имеем два канала: один для предсказания Region Score, другой — для Affinity Score (далее разберём их подробнее).
Вся архитектура полносвёрточная, поэтому пространственная информация сохраняется. Для нас это критически важно, поскольку мы хотим работать с отдельными символами.
Генерация Ground Truth Labels
Итак, у нас есть два выходных канала — Region Score и Affinity Score. Для каждого изображения в тренировочной выборке мы создаём следующие ground truth labels:
- для region score — маску с вероятностями того, что пиксель содержит центр символа (значения масок максимальны в центре символов и убывают при отдалении от них);
- для affinity score — маску с вероятностями того, что пиксель содержит центр пространства между соседними символами (значения масок максимальны в центрах между символами и убывают при отдалении от них).
Отметим, что маски содержат именно вероятности, полученные с помощью Gaussian Heatmap, а не дискретные значения, как в случае сегментации. Таким образом, мы гибче работаем с рамками, которые не являются жёстко ограниченными.
Применение Gaussian Heatmap:
- мы не хотим работать с каждым пикселем по отдельности внутри бокса, поэтому готовим 2D Gaussian карту;
- так как бокс имеет произвольную форму, мы не можем просто наложить на него эту карту, зато можем провести проективное преобразование (оно переводит прямые в прямые), при котором 2D Gaussian карта примет форму исходного бокса.
Окей, с Region боксами (боксами символа), следовательно, с построением Region Score карт разобрались, а как мы будем считать Affinity Score боксы? Авторы предлагают использовать следующую конструкцию на примере двух соседних символов:
- Возьмём два бокса (просто выпуклые четырёхугольники) соседних символов, проведём внутри каждого диагонали. Точки пересечения диагоналей — центры символов; они вместе с основаниями образуют треугольники;
- Возьмём центры этих треугольников слева и справа. Получим 4 точки, которые и будут образовывать боксы Affinity.
Благодаря использованию таких Region и Affinity карт модель фокусируется как на отдельных символах, так и на их соседях. Это существенно упрощает задачу детекции текста по сравнению с тем, если бы мы пытались искать всё слово целиком.
Weekly-supervised обучение
Описанная выше генерация Ground Truth Labels хорошо подходит для синтетических данных, однако, как уже отмечалось ранее, реальные данные обычно имеют разметку только на уровне слов. Как же тогда перейти от неё к нужной нам разметке?
Ресёрчеры предлагают использовать подход, который можно описать фразой «Fake it till you make it» 😄. Основная идея такого подхода следующая:
- Вырезаем слова по боксам.
- Используем промежуточную модель (сам CRAFT на промежуточных эпохах) для предсказания боксов символов.
- Применяем алгоритм watershed для разделения символов.
- Полученные области символов преобразуем обратно в координаты исходного изображения.
- Генерируем карты регионов символов и их связности.
- Создаём ограничивающие рамки для каждого символа.
- Для улучшения надёжности полученных предсказаний вводим оценку достоверности: соотношение найденных символов к числу символов в Ground Truth. Далее используем её как вес при обучении.
Формула для оценки достоверности:
где:
- \( s_{\text{conf}}(w) \) — оценка уверенности для образца \( w \);
- \( l(w) \) — длина слова;
- \( l_c(w) \) — оценённая длина символов.
Мы используем следующий лосс для обучения:
\( L = \sum_{p} S_c(p) \cdot\left( \|S_r(p) — S^*_r(p)\|^2_2 + \|S_a(p) — S^*_a(p)\|^2_2 \right) \)
где \( S_{r}, S_{a} \) — предсказанные оценки region и affinity, \( S_{r}^{}, S_{a}^{} \) — псевдо-разметка region и affinity.
Таким образом, при обучении точность предсказания символов улучшится с течением времени, оценки уверенности тоже. То есть модель сама себе будет помогать генерировать псевдо-разметку до тех пор, пока у неё не начнёт это хорошо получаться.
Обобщаем вышесказанное и получаем следующий пайплайн:
Инференс
В конце CRAFT мы получаем две хитмапы — для Region и Affinity. Поскольку мы хотим иметь полигоны для текста, нам нужно провести пост-процессинг полученных масок.
- Алгоритм генерации для горизонтального текста:
- Создаём бинарную маску \( M \), где \( M(p) = 1, \space \text{если} \space S_{r}(p) > τ_{r} \space \text{или} \space S_{a}(p) > τ_{a} \), а \( S_{r}(p), S_{a}(p) \) — карты Region и Affinity, \( τ_{r}, τ_{a} \) — пороги для них;
- Применяем маркировку связных компонент (CCL) к \( M \);
- Для каждой компоненты находим минимальный по площади повёрнутый прямоугольник;
- Ограничивающий прямоугольник слова определяем единственным охватывающим прямоугольником, так как имеем цепочку боксов отдельных символов.
- Алгоритм генерации полигонов для изогнутого текста:
- Находим линии локальных максимумов в регионах символов (синие стрелочки на рисунке ниже) — это наиболее вероятные центры символов;
- Выравниваем длины линий локальных максимумов до максимальной в слове;
- Создаём центральную линию, соединяющую центры линий локальных максимумов (жёлтая линия), получаем центральную линию, следующую изгибу текста;
- Поворачиваем линии локальных максимумов перпендикулярно центральной линии — таким образом учитываем наклон символов;
- Определяем конечные точки линий локальных максимумов как кандидатов для контрольных точек полигона;
- Смещаем крайние линии наружу для полного покрытия текстовой области;
- Формируем финальные контрольные точки полигона — точки смещённых крайних линий и остальных линий локальных максимумов становятся финальными контрольными точками полигона.
Обучение
Для обучения использовались следующие датасеты (подробнее о некоторых из них можно почитать здесь):
- ICDAR2013 (IC13);
- ICDAR2015 (IC15);
- ICDAR2017 (IC17);
- MSRA-TD500 (TD500);
- TotalText (TotalText);
- CTW-1500 (CTW);
- SynthText.
Сама тренировка CRAFT делилась на две стадии: сначала его предобучили на SynthText, а затем тюнили на части датасетах выше. Почему только на части? Потому что разметка должна удовлетворять следующим условиям:
- боксы являются четырёхугольниками (боксы четырёхугольного типа);
- возможен подсчёт количества символов.
Подходящими датасетами оказались датасеты из набора ICDAR — IC13, IC15, IC17.
Боксы из датасета MSRA-TD500 тоже являются четырёхугольными, но они не содержат транскрипций, поэтому используются только для тестирования.
Данные из TotalText и CTW-1500 имеют полигональную разметку и используются только для тестирования.
Заключение
Отметим сильные и слабые стороны CRAFT:
- Благодаря использованию Region и Affinity карт CRAFT остаётся более устойчивым к изменению масштаба, чем word-level (на уровне слов) детекторы.
- Если судить по результатам бенчмарков (причём на данных, которых не было в обучении) — CRAFT улавливает общие характеристики текстов, а не подстраивается под конкретные наборы данных.
- CRAFT имеет проблемы в работе с курсивным текстом, в частности с такими языками, как бенгальский и арабский — каждый символ в них написан курсивом, из-за этого их сложнее разделить.
- CRAFT хуже справляется с текстом на сложном фоне при наличии бликов, искажений или нестандартных шрифтов, особенно если такие случаи не были представлены в обучающих данных.
- Медленная обработка — скорость работы модели не подходит для задач реального времени.
Полезные ссылки
- Официальная имплементация на github — Official implementation
- Модель на HuggingFace — HuggingFace
