Всем привет! Прошло 2 года с момента моей первой публикации, которая была посвящена детектору объектов RetinaNet. Всё это время во мне часто просыпалось желание снова поделиться с вами каким-либо полезным и интересным материалом. И вот, в поисках лучшей модели для решения задачи OCR, я наткнулся на модель SVTR в репозитории PaddleOCR и статью "SVTR: Scene Text Recognition with a Single Visual Model". По словам авторов статьи модель является SOTA в задаче scene text recognition. Ознакомившись с устройством сети и протестировав модель на своих данных, я убедился, что модель действительно заслуживает внимания, так как она достаточно быстро сходится, выдавая высокий процент точности и малое время инференса.

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

Введение

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

Обычно модели нейронных сетей состоят из нескольких частей. Одна часть выделяет признаки из входного изображения, а другая обрабатывает эти признаки как последовательность, предсказывая символ за символом. Ярким представителем таких сетей является CRNN, который был предложен ещё в далёком 2015м, но до сих пор является популярным, так как прост в обучении и очень эффективен. Состоит эта модель из CNN блока, который выделяет признаки, и BiLSTM блока, который обрабатывает полученные признаки. Выход из последнего блока блока декодируется в текст с помощью CTC. Позже были предложены ряд энкодер-декодер архитектур, а так же методы обучения с передачей знаний о языке моделям. Но первые, как правило, основаны на авто-регрессионных методах, что увеличивает время инференса, а вторые сложны в обучении и менее точны по сравнению с SOTA моделями.

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

В ходе исследований было предложено несколько конфигураций архитектуры SVTR. SVTR-L достигла уровня SOTA в эффективности распознавания текста сцены на английском и китайском языках. В то время как SVTR-T имеет всего 6.03M параметров и имеет скорость инференса 4.5ms на NVIDIA 1080Ti GPU.

Теперь подробно рассмотрим архитектуру модели!

Архитектура

На входе нейронной сети располагается блок STN (Spatial Transformer Network, который был предложен в статье “ASTER: An Attentional Scene Text Recognizer with Flexible Rectification”). Этот блок выполняет выпрямление текста, находящегося на исходном изображении. Сперва изображение с входным размером размером (32, 100) уменьшается до размера (32, 64) и подаётся в простую свёрточную сеть с полносвязным слоем в конце. Выход сети - координаты ключевых точек, определяющих нижнюю и верхнюю границы текста. Затем с помощью Thin-Plate Spline (который подробно описан в статье ASTER) выполняются преобразования исходного изображению, располагая текст на нём горизонтально. На рисунке ниже приведены примеры с результатом применения таких преобразований. Главном преимуществом использования STN по сравнению с другими методами выпрямления текста (например, аффинными преобразованиями) является отсутствие необходимости в разметке, так как сеть обучается засчёт градиентов, распространяющихся от сети распознавания.

Далее, новое изображение с выпрямленным текстом подаётся на вход блоку Patch Embedding. Суть этого блока заключается в разбитии входного изображения на патчи (малые части одного размера) и представлении каждого патча в виде вектора. Такой подход был впервые предложен в статье "An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale" и предполагал разбитие изображения на непересекающиеся патчи размером (16, 16) и применение линейного преобразования к каждому из них. Но в SVTR для кодирования патчей размером (4, 4) было предложено использовать два свёрточных слоя размером (3, 3) и шагом 2. Такой подход немного увеличивает вычислительные затраты, но даёт более обширное представление о каждой части изображения. На рисунке ниже изображена схема работы рассмотренных подходов.

К выходу Patch Embedding добавляется position embedding, для передачи информации о позиции каждого патча на изображении. В оригинальной статье каждый полученный эмбеддинг обозначается как character component, который мы для упрощения будем называть компонентом. Полученный набор компонент представляет собой тензор размером (H/4, W/4, D0), где D0 - размер компонента (эмбеддинга), H и W - высота и ширина исходного изображения соответственно.

После блока Patch Embedding следует ключевая часть архитектуры SVTR. Поэтому, прежде чем подробно разобрать её устройство, вспомним какие проблемы наиболее распространены при решении задачи распознавания текста. Одной из таких проблем является неправильная классификация похожих символов (например цифру “0” и буква “О”). Таким образом, правильное определение класса символа значительно зависит от качества выделенных признаков из соответствующей ему области на изображении. Как правило в существующих исследованиях выделенные признаки представляют собой последовательность, в которой закодирован текст, а каждый признак соответствует небольшой области на изображении. Часто признаки содержат шум, особенно когда текст расположен неровно. Тонкий момент заключается в том, что обработка выполняется с упором на получение последовательности, в которую уже закодирован текст. Такой подход оказывается недостаточным для описания каждого символа отдельно.

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

Для извлечения глобальных и локальных признаков авторами были разработаны global mixing и local mixing слои. Global mixing представляет собой multi-head self-attention слой, который вычисляет зависимости между каждым компонентом. Local mixing отличается тем, что для каждого компонента вычисляется взаимосвязь только с соседними компонентами (после перемножения key и query всех компонентов к полученной матрице внимания применяется заранее составленная маска). Как описали авторы статьи, local mixing, аналогично операции свёртки, проходит скользящим окном размером (7, 11) по всему набору компонентов с шагом (3, 5), таким образом на каждом шаге вычисляя зависимость между компонентами в текущем окне.

Итак, вернёмся к тензору, полученному от Patch Embedding размером (H/4, W/4, D0). Этот тензор проходит 3 стадии обработки. На каждой стадии сперва меняется размер входного тензора из (h, w, d) в (h*w, d), где h - количество компонентов по высоте, w - количество компонентов по ширине, d - размер компонента. Затем идёт последовательность mixing blocks, которые состоят из layer normalization, local или global mixing, layer normalization и mlp с добавлением shortcut connections. После применения нескольких mixing блоков в конце каждой стадии полученный тензор с размером (h*w, d) снова преобразуется в тензор с размером (h, w, d). 

Так как обработка с постоянным пространственным разрешением вычислительно затратна и приводит к избыточному представлению, то полученный тензор проходит блок merging, состоящий из слоя свёртки и layer normalization, уменьшая размер тензора по высоте в 2 раза, увеличивая количество каналов, но оставляя ширину неизменной. Эта операция не только снижает вычислительные затраты, но и способствует изучению изображения в разном масштабе, не оказывая влияния на пространственную информацию по ширине (что уменьшает шанс кодирования соседних символов в один).  

Пройдя две стадии стадии обработки через блоки mixing и merging, на третей (заключительной) стадии merging блок заменяется на combing. Всё дело в том, что при входном изображении с размером 32 по высоте на выходе 3-го mixing блока образуется тензор с размером 2 по высоте, поэтому операция свёртки заменяется на pooling слой, за которым идёт полносвязный слой, нелинейная функция активации и dropout. Таким образом, в конце 3-й стадии обработки формируется тензор с размером (1, W/4, D3).

После объединения признаков блоком combing следует последний блок - полносвязный слой, на выходе которого получаем последовательность размером W/4, состоящей из компонентов размером N, где N - количество символов для предсказания. Предсказанная последовательность и является закодированным текстом, где каждый компонент одного символа предсказывает один определенный класс, а компоненты не текста предсказывают класс blank.

В итоге, авторами было предложено 4 варианта архитектуры SVTR разного размера.

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

Что ещё интересно отметить, так это сравнение различных конфигураций блоков и их влияние на результат в рамках архитектуры SVTR-T. Было экспериментально подтверждено положительное влияние блока Patch Embedding (Ours) по сравнению со свёрточным слоем с ядром (7, 7) с шагом 4 (Overlap) и линейными преобразованиями предложенными в статье ViT (Linear), а также влияние уменьшения размерности (Prog) блоком Merging по сравнению с обработкой без данного блока (None).

Также были проведены эксперименты, определяющие наилучшие комбинации Local (L)  и Global (G) mixing слоёв. Самый высокий показатель точности получила архитектура с 6 Local mixing слоями в начале (в 3-х блоках 1-й стадии и первых 3-х блоках 2-й стадии) и 6 Global mixing слоями в конце (в последних 3-х блоках 2-й стадии и 3-х блоках 3-й стадии.) Помещая в начало local mixing слои, нейросеть извлекает важные локальные детали из каждого символа отдельно, заставляя global mixing фокусироваться на изучении закодированных компонент текста, не вдаваясь в локальные детали. Это подтверждается ухудшением результата при перестановке global и local mixing слоёв.

Визуализация карт внимания подтверждает, что нейронная сеть SVTR извлекает признаки разного уровня. В первом ряду (на изображении ниже) можно увидеть, что области активации находятся на разных частях символа “B”, выделяя его локальные признаки. Во втором ряду области активации уже покрывают каждый отдельный символ целиком (“В”, “L” и “S”), находя зависимости между компонентами, принадлежащими одному символу, и обрабатывая этот символ целиком. В третьем ряду эти области распространяются уже на соседние символы, что говорит о способности извлекать более глобальные признаки и находить зависимости между компонентами соседних символов.

Заключение

В итоге была получена новая state-of-the-art модель, которая не только превосходит по точности ранее разработанные архитектуры, но и, по сравнению с некоторыми методами, имеет в разы меньшее количество параметров и время инференса.

Спасибо за внимание! Надеюсь эта статья была вам интересна!