Распознавание текста вместе с CLIP4STR
Введение
Задача распознавания текста на фотографии — важная задача во многих технологических и бизнес-процессах: от анализа сканов документов до автономного вождения.
Глобально можно выделить два основных этапа ее решения:
- Детекция текста
- Его распознавание внутри задетектированных ограничивающих прямоугольников
В этой статье мы поговорим о втором этапе, а именно — о Scene text recognition (STR).
Scene text recognition (STR) — распознавание текста на фотографиях реального мира.
На рисунке 1 представлены примеры одного из популярных датасетов STR — это ICDAR2013 и ICDAR2015.
На фотографиях выше мы видим, что текст бывает сильно искажен (вытянут в длину или ширину, имеет непрямоугольную форму и др.), по-разному ориентирован в пространстве, размещен на различных сложных фонах. Он даже может быть рукописным. Это все усложняет задачу его распознавания.
Самые популярные модели распознавания текста преимущественно базируются на визуальных энкодерах: CRNN, VitSTR, TrOCR и др. Но задача распознавания текста мультимодальна — у нас есть изображение и текст. Поэтому применение мультимодальной модели, Visual-language model (VLM), кажется вполне разумным 🙂
Также подумали и авторы модели CLIP4STR. Они решили сделать новый бейзлайн для STR-задач на основе самой популярной VLM-модели — CLIP.
💡 CLIP [подробный разбор архитектуры], или Contrastive Language-Image Pre-training — модель от OpenAI, обученная на огромном датасете пар изображений и текста. Модель с двумя энкодерами для текста и картинок, которые можно использовать для генерации соответствующих качественных фичей. Их затем можно применять для своих даунстрим задач.
Идея CLIP4STR довольно проста: почему бы не использовать для распознавания текста в реальном мире модель, обученную на парах «текст — картинка», следовательно, имеющую «сильные» энкодеры для представления обеих модальностей. Действительно, можно применять картиночный энкодер CLIP для получения эмбеддингов изображения с текстом и из них генерировать какой-то текст (как, например, это делается в ViTSTR), а затем использовать текстовый энкодер CLIP для исправления ошибок в тексте с предыдущего этапа.
Если проанализировать слой внимания визуального энкодера CLIP (см. рисунок 2), можно увидеть: на семплах из датасета STR CLIP в тех местах, где на фото есть текст, веса внимания больше. То есть CLIP может выступать как хороший бэкбон для задачи распознавания текста.
Виды SOTA для STR-задач
Перед более подробным рассмотрении CLIP4STR давайте вспомним, какие SOTA-модели существуют для задачи распознавания текста.
В основном модели STR делятся на два вида: контекстно-независимые и контекстно-зависимые.
Контекстно-независимые модели STR
К контекстно-независимым моделям можно отнести, например, CRNN-CTC модель.
CRNN — это комбинация сверточного бэкбона и рекуррентной нейронной сети (см. рисунок 3). Свёрточная часть кодирует входное изображение во внутреннее представление, которое затем разбивается на последовательность токенов для рекуррентной нейронной сети. Проблема в том, что длина последовательности, которая подается на вход RNN (то есть разбитая в последовательность карта признаков CNN), не приравнивается длине распознаваемого слова (она обычно в несколько раз длиннее). Поэтому здесь не получится применить обычную кросс-энтропийную функцию, нужно предварительно сопоставить фрагмент изображения с конкретной буквой, что позволяет реализовать CTC (Connectionist Temporal Classification) loss-функцию.
Также можно упомянуть модели на базе сегментационных сетей, например, CA-FCN. Здесь разработчики решили уйти от предсказания 1D-последовательности и предложили учитывать положение символов в 2D-пространстве с помощью задачи сегментации изображения (см. рисунок 4). Сегментационная сеть предсказывает многоклассовую маску с символами. В сети также реализован механизм внимания, который позволяет лучше отделять символы от фона. На финальной стадии распознавания полученные маски символов сортируются слева направо.
Но все методы не могут справиться с ситуацией, когда часть слова перекрыта, или какие-то символы в слове искажены. То есть эти модели не учитывают семантику слов. Следующий логичный шаг развития STR-моделей — применение трансформерных моделей. Они очень хорошо себя показывают в задачах восстановления или продолжения последовательности, будь то слов или символов. Такие модели называются контекстно-зависимыми.
Контекстно-зависимые модели STR
К контекстно-зависимым моделям относится, например, TrOCR.
Трансформерная модель TrOCR строится на базе ViT (vision transformer) в качестве энкодера и трансформерного декодера. Энкодер ViT бьет входное изображение на непересекающиеся патчи и кодирует слово в последовательность эмбеддингов, а декодер принимает на вход эмбеддинг патчей картинок и формирует на выходе последовательность букв.
Преимущества подхода:
- TrOCR позволяет использовать предобученные на большом объеме данных трансформерный картиночный энкодер и текстовый декодер, что помогает избежать ресурсоемкого и длительного этапа предобучения и использовать все преимущества, которые дает современный self-supervised learning
- ViT позволяет избавиться от индуктивного байеса, характерного для сверточных сетей, за счет механизма внимания, поскольку в трансформерах учитываются корреляции между различными патчами изображения, а не только между компактно расположеными пикселями, как в сверточных сетях. Это особенно актуально для текстовых изображений, где больший смысл имеет семантика слова, а не изображения
- TrOCR упрощает как этап препроцессинга, так и постпроцессинга
- а еще легко расширяется до мультиязыковой модели: достаточно расширить словарь и использовать мультиязыковой текстовый декодер
Помимо TrOCR, есть много моделей с использованием трансформеров: ViTSTR (предлагают отказаться от текстового декодера), PARSeq (используют текстовый декодер, но дополнительно применяют обобщенный авторегрессионный подход Permuted Sequence Modeling (PSM), о нем поговорим ниже).
Модель CLIP также пытались применять неоднократно: CLIPTER, CLIP-OCR. Мы же рассмотрим CLIP4STR, которая занимает лидирующую позицию в самом представительном бенчмарке задач STR — ICDAR2013.
Архитектура CLIP4STR
Энкодер
Как видно из названия модели, в основе ее архитектуры лежит CLIP, а именно — энкодер изображения и текста (см. рисунок 6). Причём, если в оригинальном CLIP на выходе энкодера картинок используется только class token, то здесь используются все токены последовательности. Логика такая: мы хотим восстановить последовательность символов на фото, поэтому нам нужно больше подробной информации об изображении, которую способен дать энкодер. Полученная последовательность символов на выходе картиночного декодера подается на вход кросс-модальной ветки, состоящей из текстового энкодера и декодера (сиреневые блоки на рисунке 6). Эта ветвь улучшает прогноз визуальной ветки (зеленые блоки на рисунке 6). Часть текстового энкодера CLIP заморожена, что существенно снижает затраты на обучение. В то же время энкодер изображений полностью разморожен, поскольку домен STR-картинок сильно отличается от домена, где обучался CLIP. Важно отметить: градиент не проходит от зелёной ветви к сиреневой, то есть обе ветви обучаются независимо.
Выходы энкодера изображения
Декодер
Текстовый (кросс-модальный) и картиночной декодеры схожи по архитектуре (см. рисунок 7) и отличаются только входами.
Оба декодера используют технику Permuted Sequence Modeling (PSM), предложенную в модели PARSeq.
Обычно в моделях по распознаванию символов применяют авторегрессионный (AR) подход к восстановлению последовательности букв на выходе: слева направо, либо справа налево. Поскольку трансформерная архитектура позволяет работать параллельно со всеми токенами последовательности, мы можем реализовать всевозможные сценарии восстановления элементов последовательности за счет маскирования внимания. На рисунке 8 приведено несколько вариантов масок внимания для последовательности длиной 3: a, b — маски, реализующие последовательность слева направо и справа налево; c, d — рандомные маски внимания. Таким образом, PSM позволяет реализовать различные сценарии восстановления последовательности, что улучшает качество STR-модели.
Для лучшего понимания метода PSM представим его как способ ансамблирования AR и NAR (неавторегрессионных) моделей с общими весами.
Итак, декодер принимает на вход либо эмбеддинг картиночного энкодера
Внутри декодера — два слоя многоголового внимания (Multi-Head Attention, или MHA). Первый MHA отвечает за вычисления внимания между контекстом и обучаемым позиционным эмбедингом, где D — размерность эмбеддинга:
Второй MHA — за вычисления внимания между
\( m_2=softmax(\frac{m_1F^T}{\sqrt{D}})F+m_1 \)
Обучение CLIP4STR
При обучении входной контекст декодера формируется на основании набора масок внимания (см. рисунок 8). То есть модель учится воспроизводить последовательность как в манере AR (слева направо или справа налево), так и в случайной форме — NAR.
В качестве функции потерь считается сумма кросс-энтропий между фактом \( \hat{y} \) и выходами картиночной ветви —
\( L=CE(y^i,\hat{y})+CE(y,\hat{y}) \)
Инференс CLIP4STR
Давайте рассмотрим, как в CLIP4STR реализуется инференс — применение обученной модели.
На рисунке 9 приведен алгоритм работы декодера при инференсе изображения, который можно разделить на три этапа:
- Авторегрессионная генерация последовательности картиночной ветвью.
- Авторегрессионная генерация последовательности кросс-модельной ветвью.
- Неавторегрессионное уточнение последовательности.
На первом этапе картиночный декодер итеративно генерирует выходное слово, начиная с токена [B].
На втором этапе первое приближение выходного слова, полученное картиночным декодером, подается в качестве контекста на вход кросс-модальной ветви. Кросс-модальный декодер также в форме AR генерирует свой вариант выходного слова.
На этих двух этапах используются только авторегрессионные маски внимания: слева направо и справа налево.
На третьем этапе при работе декодера применяются случайные маски внимания. Здесь уточняется слово, предсказанное ранее.
Результаты
В эксперименте участвовало три варианта модели с предобученными энкодерами из OpenCLIP: CLIP-ViT-B/16, CLIP-ViTL/14, и CLIP-ViT-H/14 (см. рисунок 10):
Как видно из результатов сравнения CLIP4STR c моделями SOTA для задач STR (см. рисунок 11), CLIP4STR существенно обгоняет другие модели на базе CLIP на сложном реальном датасете ICDAR. При этом масштаб модели не сильно влияет на метрику (см. на датасеты ICDAR, SVT, IIIT5K).
В то же время масштаб модели помогает справиться с датасетом с большим количеством пропущенных символов. Например, если сравнить CLIP4STR-L и H на бенчмарке HOST.
Важно отметить: CLIP4STR сопоставим с современными SOTA на основных STR-бенчмарках или даже обгоняет их.
Эксперименты
Расскажу немного о своем опыте запуска CLIP4STR на синтетике в сравнении с CRNN.
В качестве тренировочных и тестовых данных я использовал синтетику, сгенерированную с помощью фреймворка synthtiger от ClovaAI. На рисунке 12 представлены примеры обучающих данных для CLIP4STR и CRRN.
Для CRNN все картинки приводятся к одному размеру без искажений пропорций за счет паддинга. А для CLIP4STR — к стандартному размеру трансформерного энкодера (224х224).
На основе обучающего датасета можно сделать предположение: трансформерная модель должна быть более робастной ко входным искажениям текста. Давайте посмотрим, что получилось в результате обучения.
Для эксперимента использовался CRNN с предобученным Resnet18 на Imagenet 1k и CLIP4STR c весами CLIP-ViT-B/16 от OpenAI. Метрика accuracy на тесте у CRNN составила 0.74, у CLIP4STR — 0.9. Это вполне объяснимо, если учитывать сложность трансформерной модели и объем выборки на этапе предобучения CLIP. Посмотрим на реальные примеры сравнения прогнозов двух моделей:
Модель на базе CLIP отлично справляется, когда буквы видно частично, как на рисунке 13 (a), или текст сильно искажен, как на рисунке 13 (b).
Рисунок 13 (с) — хороший пример того, что модель может выступать в роли автокорректора и исправлять опечатки. Иногда это полезно, но такое “придумывание” модели может приводить к нежелательным результатам. Как, например, на рисунке 13 (d), где CLIP4STR, с одной стороны, справилась лучше и распознала все символы на картинке, но в то же время решила “додумать” окончание, так как последняя буква на фото была обрезана.
Заключение
CLIP4STR, трансформерная модель на базе CLIP, — хорошая SOTA в задаче STR сегодня. Она достигает state-of-the-art результатов на 11 базовых STR-бенчмарках и хорошо проявляет себя в рамках проведенного эксперимента.
Конечно, такая модель более требовательна к вычислительным ресурсам и объему обучающей выборки, но это является особенностью большинства трансформерных моделей.
Также стоит учитывать склонность LLM к галлюцинациям и “додумываниям”, например, в тех ситуациях, когда часть букв не видна или имеется символ, которого не было в словаре.
Полезные ссылки
