Как мы распознаем фото документов пользователей. Часть II

В прошлой статье я кратко описал методы и подходы, которые мы используем в inDriver при распознавании фото документов. Во второй части подробно опишу архитектуру CRAFT и CRNN, а также варианты их использования. Прошу под кат!

image-loader.svgСодержание

CRAFT

CRAFT Scene Text Detection — это задача для обнаружения текстовых областей на сложном фоне и маркировки их ограничивающими рамками. Модель CRAFT была предложена в 2019 году, ее основная цель — локализовать отдельные области символов и связать обнаруженные символы с экземпляром текста.

CRAFT использует fully convolutional network архитектуру, основанную на VGG-16 c batch normalization в качестве backbone (базовой сети). Модель VGG16 используется для извлечения признаков и нужна для кодирования входных данных сети в определенное представление признаков. Модель имеет skip connections (пропускные соединения) в части декодирования, которая похожа на UNet в том, что агрегирует низкоуровневые функции через UpConv.

Архитектура CRAFTАрхитектура CRAFT

CRAFT на выходе дает 2 оценки для каждого символа:

  1. character region score (оценка области символа).

  2. affinity score (оценка близости символа). Здесь определяется, можно ли символы объединить в одно слово.

Иными словами, region score используется для локализации отдельных символов на изображении, а affinity score — для группировки каждого символа в одно слово.

Процесс обучения

Чтобы компенсировать отсутствие аннотаций на уровне символов, используется weakly-supervised learning framework (слабо контролируемая система обучения). Она оценивает ground truths (эталонные данные) на уровне символов в существующих наборах данных на уровне реальных слов.

Визуализация работы CRAFT для текстов различной формы. (a) — тепловые карты, (b) — результаты обнаружения текстов различной формыВизуализация работы CRAFT для текстов различной формы. (a) — тепловые карты, (b) — результаты обнаружения текстов различной формы

Главная цель — точно локализовать каждый отдельный символ на естественных изображениях. С этой целью нейронная сеть обучается предсказывать область символов и близость между ними. Так как общедоступного набора данных на уровне символов нет, модель обучается со слабым контролем (weakly-supervised manner).

Ground Truth Label Generation

Для каждого тренировочного изображения мы генерируем ground truth label для region score и affinity score с bounding boxes на уровне символа. Region score представляет собой вероятность того, что данный пиксель является центром символа, affinity score — вероятность пробела между соседними символами.

В отличие от бинарной карты сегментации, которая помечает каждый пиксель дискретно, мы кодируем вероятность центра символа с помощью гауссовой карты интенсивности (Gaussian heatmap). Мы используем представление тепловой карты для оценки как region score, так и affinity score.

Gaussian heatmapGaussian heatmapПроцедура генерации эталонных данных. Создаются ground truth labels на основе синтетического изображения с аннотациями на уровне символовПроцедура генерации эталонных данных. Создаются ground truth labels на основе синтетического изображения с аннотациями на уровне символов

Вычисление значения распределения Гаусса непосредственно для каждого пикселя в bounding box занимает очень много времени. Поскольку bounding box символов на изображении обычно искажаются из-за перспективных проекций, используются следующие шаги для аппроксимации и генерации ground truth как для region score, так и для affinity score:

  1. Подготавливается двумерная изотропная гауссова карта.

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

  3. Происходит деформация (warp) карты Гаусса в область прямоугольника (box area).

Для ground truth affinity score аффинные блоки определяются с использованием смежных блоков символов. Рисуя диагональные линии для соединения противоположных углов каждого блока символов, мы можем сгенерировать 2 треугольника, которые назовем верхним и нижним символьным треугольником. Затем для каждой смежной пары создается поле сходства (affinity box), устанавливая центры верхнего и нижнего треугольников в качестве углов поля.

Предлагаемое определение ground truth позволяет модели обнаруживать большие или длинные текстовые экземпляры, несмотря на использование небольших воспринимающих полей. Обнаружение на уровне символа позволяет сверточным фильтрам фокусироваться только на внутрисимвольных и межсимвольных полях, а не на всем экземпляре текста.

Weakly-Supervised Learning

В отличие от синтетических наборов данных, реальные изображения обычно имеют аннотации на уровне слов. Здесь мы генерируем символьные боксы из каждой аннотации на уровне слов в слабо контролируемой манере (weakly-supervised manner).  

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

Чтобы отразить надежность прогноза промежуточной модели, значение карты достоверности (confidence map) для каждого блока слов вычисляется пропорционально количеству обнаруженных символов, деленному на количество основных символов истинности.

Обучение проводится с использованием как реальных, так и синтетических изображений под слабым контролем (weakly-supervised fashion)Обучение проводится с использованием как реальных, так и синтетических изображений под слабым контролем (weakly-supervised fashion)Процесс разделения символовПроцесс разделения символов

Сначала изображения на уровне слов обрезаются из исходного изображения (word box), потом модель делает оценку region score. А далее алгоритм watershed labeling разделяет области символов, чтобы построить bounding box вокруг каждого символа. После координаты символьных полей преобразуются обратно в исходные координаты изображения с помощью обратного преобразования из croping.

Когда модель обучается с использованием слабого контроля (weak-supervision), мы вынуждены тренироваться с неполными pseudo-GTs (pseudo-ground truths). Если модель обучена с неточными region score, выходные данные могут быть размыты в области символа. Чтобы предотвратить это, мы измеряем качество каждой pseudo-GTs сгенерированной моделью. 

К счастью, в текстовой аннотации есть очень сильная подсказка — длина слова. В большинстве наборов данных предоставляется транскрипция слов, и длина слов может использоваться для оценки достоверности pseudo-GTs.

Для аннотированной выборкиw на уровне слова (word-level) обучающих данных пустьR(w)иl(w)будут bounding box и длиной слова выборкиw соответственно. Посредством процесса разделения символов мы можем получить оценочные ограничивающие прямоугольники символов и соответствующую им длину символовl^c(w). Затем оценка достоверности (confidence score)s_{conf}(w) для выборкиw вычисляется как:  

s_{conf}(w) = \frac{l(w) - min(l(w), |l(w) - l^c(w)|)}{l(w)}

Оценка достоверности по пикселямS_cдля изображения вычисляется как:

\begin{equation*} S_c(p) = \left\{ \begin{array}{lr} s_{conf}(w) &  p \in  R(w),\\ 1 &  \text{otherwise}, \end{array} \right. \end{equation*}

pобозначает пиксель в областиR(w). ОбъектL(Loss Function) определяется как:

L = \sum_{p} S_c(p) \cdot (||S_r(p) - S_{r}^{*}(p)||_{2}^{2} + ||S_a(p) - S_{a}^{*}(p)||_{2}^{2})

S_{r}^{*}(p)и S_{a}^{*}(p) обозначают pseudo-ground truth region score и affinity map, соответственно. АS_r(p) S_a(p)— прогнозируемую region score и affinity score. При обучении с использованием синтетических данных мы можем получить реальную достоверную информацию, поэтому дляS_{c}(p)установлено значение 1. По мере выполнения обучения модель CRAFT может предсказывать символы более точно, а оценки достоверностиs_{conf}(w)постепенно увеличиваются.

Карта region score-символов в процессе обученияКарта region score-символов в процессе обучения

На начальных этапах обучения region scores относительно низкие для незнакомого текста в естественных изображениях. Модель изучает внешний вид новых текстов и синтезированные тексты, которые имеют другое распределение данных по сравнению с набором данных SynthText. 

Если показатель достоверностиs_{conf}(w) ниже 0,5, предполагаемыми bounding box символами следует пренебречь, поскольку они имеют неблагоприятные эффекты при обучении модели. В этом случае мы предполагаем, что ширина отдельного символа постоянна, и вычисляем предсказания на уровне символа, просто разделив область словаR(w)на количество символовl(w). Затемs_{conf}(w)устанавливается на 0,5 для изучения невидимого внешнего вида текстов.

Обучение на синтетических данныхОбучение на синтетических данныхОбучение на реальных данныхОбучение на реальных данных

Inference

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

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

Постобработка для поиска ограничивающих рамок резюмируется следующим образом. Сначала двоичная карта M, покрывающая изображение, инициализируется значением 0.M(p)устанавливается в 1, если S_r(p) > \tau_{r}» src=«https://habrastorage.org/getpro/habr/upload_files/71d/4e9/950/71d4e99506f07e4683d49c8d6a01563e.svg» />или<img alt= \tau_{a}» src=«https://habrastorage.org/getpro/habr/upload_files/53a/a29/bce/53aa29bce1c010b62a0e531cacca6197.svg» />, где \tau_{r} — region threshold,  \tau_{a} — affinity threshold. Затем выполняется Connected Component Labeling (CCL) на M. Наконец, QuadBox получается путем нахождения повернутого прямоугольника с минимальной площадью, охватывающей соединенные компоненты, соответствующие каждой из меток. Для этой цели могут применяться функции connectedComponents и minAreaRect, предоставляемые OpenCV.

Преимущество CRAFT в том, что ему не нужны дополнительные методы постобработки, такие как Non-Maximum Suppression (NMS). Поскольку у нас есть капли изображений областей слов, разделенных CCL, bounding box для слова просто определяется одним охватывающим прямоугольником. С другой стороны, наш процесс связывания символов выполняется на уровне пикселей, которые явно полагаются на поиск отношений между текстовыми компонентами.

Кроме того, мы можем создать многоугольник вокруг всей области символов, чтобы эффективно работать с изогнутыми текстами.

Генерация полигонов для текстов произвольной формыГенерация полигонов для текстов произвольной формы

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

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

Результаты для всего набора текстовых данныхРезультаты для всего набора текстовых данных

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

Training Strategy

Процедура обучения включает в себя 2 шага:  

1. Используем набор данных SynthText для обучения сети для 50 тысяч итераций.

2. Каждый набор тестов принимается для точной настройки модели. 

Некоторые текстовые области «DO NOT CARE» в наборах данных ICDAR 2015 и ICDAR 2017 игнорируются при обучении, если для s_{conf} (w) установлено значение 0. Рекомендую использовать оптимизатор ADAM во всех процессах обучения.

Для обучения без учителя требуются 2 типа данных: четырехугольные аннотации для обрезки изображений слов и транскрипции для расчета длины слова. Этим условиям соответствуют наборы IC13, IC15 и IC17.

В общем доступе можно найти предобученную модель CRAFT только на наборах данных ICDAR и протестированную на других без тонкой настройки. 2 разные модели обучаются с наборами данных ICDAR. Первая модель обучается на IC15 только для оценки IC15. Вторая модель обучается как на IC13, так и на IC17 вместе, что используется для оценки других 5 наборов данных. Никакие дополнительные изображения для обучения не используются. Количество итераций для точной настройки установлено равным 25k.

В случае набора данных CTW-1500 сосуществуют 2 сложные характеристики: аннотации, которые предоставляются на уровне строк и имеют произвольные полигоны. Чтобы исправить это, добавляется еще одна нейронная сеть — LinkRefiner. 

LinkRefiner

Входные данные LinkRefiner представляют собой объединение region score, affinity score и промежуточной карты объектов CRAFT. Выходные данные — уточненную affinity score, скорректированную для длинных текстов. Для объединения символов вместо исходной affinity score используется уточненная affinity score, затем генерация полигонов выполняется так же, как и для TotalText. Только LinkRefiner обучен на наборе данных CTW-1500 все то время, пока CRAFT заморожен.

Архитектура LinkRefinerАрхитектура LinkRefiner

Набор данных CTW-1500 предоставляет аннотации только для полигонов без текстовых транскрипций. Кроме того, аннотации CTW-1500 предоставляются на уровне строк и не рассматривают пробелы в качестве разделительных сигналов. Это далеко от нашего предположения о близости (affinity), которое заключается в том, что оценка близости равна нулю для символов с пробелом между ними.

Чтобы получить один длинный полигон из обнаруженных символов, мы используем мелкую сеть для уточнения ссылок — LinkRefiner. Входные данные LinkRefiner представляют собой объединение оценки региона (region score), оценки сходства (affinity score) и промежуточной карты объектов из сети, которая является результатом Stage4 Block исходной модели CRAFT.

Для ground truth LinkRefiner-линии просто рисуются между центрами парных контрольных точек аннотированных полигонов, что аналогично генерации текстовых строк. Ширина каждой линии пропорциональна расстоянию между парными контрольными точками. Выходные данные модели называются оценкой связи.

Ground truth generation for the LinkRefinerGround truth generation for the LinkRefiner

После обучения у нас есть результаты, полученные моделью, которые представляют собой оценку региона (region score), оценку сходства (affinity score) и оценку связи (link score). Здесь оценка связи (link score) используется вместо исходной оценки сходства (affinity score), и текстовый полигон получается полностью с помощью того же процесса, что и в TotalText. Модель CRAFT локализует отдельные символы, а модель LinkRefiner объединяет символы и слова, разделенные пробелами, которые требуются для оценки CTW-1500.

Результаты по набору данных CTW-1500Результаты по набору данных CTW-1500

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

В конце рассказа о CRAFT поделюсь рядом полезных материалов:

Оригинальная реализация CRAFT. 

PyTorch Re-Implementation of CRAFT с кодом для обучения. 

Библиотека на Python. 

Сейчас SoTA в области детекции текста является TextFuseNet (он точнее CRAFT, но медленнее).

Пример на датасете ICDAR 2015:

recall

precision

TextFuseNet

88.9

91.3

CRAFT

84.3

89.8

Для классификации символов изображений были обучены модели efficientnet и семислойная CNN. На тестах лучше себя показала более простая модель CNN.

Для масок, где не получалось выделить символы в силу потертости и измятости бумаги (особенно актуально для матричных шрифтов), хорошо показала себя модель CRNN. Единственная проблема — нужен большой набор размеченных данных.

CRNN

Модель CRNN представляет собой комбинацию сверточной нейронной сети backbone CNN и рекуррентной нейронной сети RNN. RNN способны улавливать в скрытом состоянии временные и пространственные особенности в последовательности входных данных. Поэтому они подходят для таких задач, как распознавание рукописного ввода и речи.

Архитектура CRNNАрхитектура CRNNПример модели CRNN, где k — kernel size, s — stride, p — padding sizeПример модели CRNN, где k — kernel size, s — stride, p — padding size

В CRNN в конце fully-connected layers (полносвязные слои) не используются, вместо этого выходные данные сверточных слоев (карт признаков) преобразуются в последовательность векторов признаков.

Представление преобразования карт признаков в последовательность векторов признаковПредставление преобразования карт признаков в последовательность векторов признаков

Затем эти векторы передаются в некоторый тип двунаправленной RNN (в нашем случае LSTM). Эта часть модели производит распределение вероятностей для каждого вектора признаков и каждой метки.

В качестве loss-функции используется CTC (Connectionist Temporal Classification) loss. На вход подается одноканальное изображение, а также символы классификации с учетом пробела (обычно он имеет индекс=0). В качестве CTC decoder мы используем greedy_decoder (к сожалению, в pytorch его еще не добавили).

Вариантов CRNN может быть много. Вместо LSTM можно использовать GRU, которая проще и требует меньше ресурсов при обучении. В качестве backbone-сети вместо CNN в примере можно взять другие сети: например, densenet, resnet, mobilenet и т.д. Также CRNN может быть использована для задачи классификации аудио по спектрограммам.

Общий подход к распознаванию документов

  1. Сегментируем нужные области в документе, желательно без лишней информации (рамок, ненужных надписей), чтобы в дальнейшем упростить препроцессинг.

  2. Пробуем с помощью CRAFT получить координаты бокса каждого символа, далее сортируем их по порядку (по оси X). Если удается разделить все символы, переходим к пункту 3, если нет — к пункту 4.

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

  4. Размечаем датасет, добавляем синтетические данные и обучаем модель CRNN.

Также для ускорения работы моделей можно перевести их вTensorRT. Если PyTorch — с помощью библиотеки Torch-TensorRT или через конвертацию весов в формат ONNX.

Спасибо, что дочитали статью до конца. Если у вас остались вопросы, буду рад ответить на них в комментариях.

© Habrahabr.ru