Привет, Хабр! В этой статье мы, команда Sber AI, расскажем о пайплайне для распознавания текста и о нюансах обучения HTR‑моделей, а также поделимся датасетом школьных обезличенных тетрадей. Это почти 2 тысячи страниц с полной разметкой полигонов слов (более 300 тысяч текстов). Если нужно, то датасет есть в открытом доступе на hugging face.
Мы в Sber AI в рамках одного из наших направлений занимаемся распознаванием рукописного текста. В частности наша команда написала пайплайн для более удобного и быстрого проведения экспериментов под разные датасеты. Он состоит из двух модулей — (1) детекция слов и (2) чтение слов. К этому ещё можно добавить этап извлечения связного текста — объединение слов в предложения и страницы. Сложность HTR задачи (handwritten text recognition) в том, что рукопись каждого человека уникальна, на неё влияет множество факторов, включая возраст и настроение. Модель чтения печатного текста можно ускорить добавлением синтетики на основе печати простыми шрифтами на фонах. А вот с HTR‑моделью это не даст такой сильный прирост, так что лучше воспользоваться синтетической рукопиской от GAN.
Отметим, что интересные задачи возникают и в модели для детекции рукописного текста. В таких данных текст, как правило, «прыгает» по странице, каждое слово под своим углом. Некоторые слова накладываются друг на друга, а строка может изгибаться, чтобы она поместилась на одной странице. Есть нюансы и при объединении двух моделей, например, нюансы даунгрейда качества чтения текста при объединении с детекцией (ошибки двух моделей мешают друг другу).
Подробности о датасете
Как уже говорилось выше, он состоит из набора фотографий школьных тетрадей. На одной картинке может быть одна или две страницы с рукописным текстом ученика (и пометками учителя).
Разметка в формате COCO (Common Objects in Context), ссылка на hugging face (там же в ридми более детальное описание структуры json с аннотацией). В разметке слова выделены полигонами, для каждого есть транскрипция текста. Есть несколько классов для текста: написанного учеником, учителем, и отдельный класс для междустрочных комментариев ученика (выполнены чаще всего карандашом). Классы можно объединить, но в некоторых случаях может быть полезно предсказывать их отдельно (например, если важно извлечь основной текст из тетради, без пометок и правок учителя).
Часто встречаются тетради, лежащие под наклоном, в некоторых случаях страницы могут быть приподняты и «волниться» — всё это будет сказываться на общем качестве распознавания. В разметке есть полигоны для границ листов (обводка страницы), что может помочь для восстановления наклона и перспективы, а также отсеивать ложные детекции слов вне листа на постпроцессинге.
Также есть разметка, которая поможет понять, на одной ли строке находятся слова. Это особенно важно при извлечении связного текста из предсказаний. Здесь сложно написать эвристику, т.к. всегда найдётся настолько нестандартный случай таких поворотов и изгибов строки, что он сломает любые if»ы. Но в разметке есть ключ «group_id»: если у двух слов одинаковый индекс group_id, значит они лежат на одной строке. Из такой разметки можно создать линии, соединяющие слова, и обучить её на сегментации, или же обучить отдельную модель кластеризации.
Сегментация и слипание масок
Для детекции слов мы используем простую сегментационную Unet‑like архитектуру (linknet) — другие semantic‑архитектуры можно взять из smp‑либы. Модель предсказывает несколько масок под каждый класс из датасета. Сегментация — простое и быстрое решение, но цена — необходимость писать постпроцессинг для предсказанных масок: трешхолдинг, извлечение контуров, исключение ложных предсказаний в несколько пикселей, апскейлинг контуров до размера исходной картинки и др.
Как было видно из картинок тетрадей выше, слова в рукописных датасетах могут располагаться очень близко друг к другу. Маски таких рядом стоящих слов будут «слипаться», на этапе постпроцессинга они объединятся в один полигон, чего мы избегаем. Это мешает чтению текста в OCR модели, а также нарушает логику определения строк на странице.
Для борьбы с этой проблемой на препроцессинге мы уменьшаем полигоны слов. Таким образом, полигон не обводит слово, а расположен как бы внутри него. Это позволяет «разлепить» слова и улучшает итоговые метрики. Впрочем, иногда такой процессинг может создавать разрывы внутри разметки полигона, так что нужно быть аккуратным. Ещё важно не забыть на этапе постпроцессинга увеличить предсказанные полигоны для корректного вырезания текста по ним.
На этапе препроцессинга добавляем маски границ для слов, они предсказываются отдельным классом сегментации. Простое добавление таких границ улучшает разделяемость слов на предсказании и повышает итоговые метрики. На инференсе сами границы никак не используются, но можно провести эксперименты с вычитанием маски границ из маски слов, что в вашем датасете может дополнительно помочь с разделяемостью объектов. Такой препроцессинг подсмотрели из DBnet.
Есть и другие варианты разделения объектов одного класса в сегментации:
Использовать instance segmentation (как mask rcnn). Можно взять архитектуры из mmdetection или detectron2. Но такие архитектуры тяжеловесны и переусложнены, их труднее допилить под свою задачу.
Использовать deep watershed transofrm. Учим дополнительно вектора, обозначающие границы слов + значения энергии/интенсивности для определения «глубины» масок. По этим параметрам сможем восстановить границы объектов и отделить их друг от друга.
OCR и beam search
После этапа сегментации предсказанные полигоны для слов нужно вырезать и отправить на предсказание в модель чтения текста. Например, в CRNN (довольно старенькая простая архитектура, отличное объяснение которой можно посмотреть здесь). Как альтернатива — seq2seq‑модели, которые позволят читать кропы с несколькими строчками текста, расположенными одна над другой (чего не может сделать CRNN в силу своей структуры), или более свежие архитектуры на базе трансформеров.
Важный этап OCR — декодирование предсказанного текста из предсказанной матрицы размерностью NxM (N — алфавит, M — time‑step, соответствующие ширине картинки).
Самый простой способ вытащить текст из матрицы — greedy search подход. Мы просто берём букву с максимальной вероятностью на каждом шаге. И это уже неплохо будет работать. Иногда модель сама будет «додумывать» скрытые места на картинке (рекуррентные сети в моделях запоминают какие‑то лингвистические особенности текста во время обучения). Но, конечно же, часто мы будет получать и совсем наивные ошибки — например, «домаьняя работа». Сеть может перепутать буквы и вместо «ш» предсказать «ь».
Здесь нам поможет beam search и добавление лингвистической модели. Отличие beam от greedy search в том, что вместо того, чтобы выбрать только один лучший путь, состоящий из букв с максимальными вероятностями, мы на каждом этапе декодирования держим несколько возможных путей — beam/луч («домаьняя», «домааняя», «домашняя» и т. д., размером beam_width). Осталось только выбрать какой‑то один вариант из выбранных — здесь поможет language model.
Немного про beam search в ctc loss
Вообще beam search для CTC loss (здесь можете посмотреть подробнее, как он работает, или почитать здесь) полезен ещё тем, что суммирует вероятности тех «сырых» путей, которые приводят к одному слову. Например, в ctc los пути «коооот» и «коооттт» и «кккот» приводят к одному слову «кот», а значит, мы можем объединить эти пути на этапе декодирования и суммировать их вероятности.
Дело в том, что мы можем перевзвешивать наши пути с учетом вероятностей из LM, и таким образом корректировать предсказания OCR. Скажем, мы знаем, что последовательность букв «шняя» встречается гораздо чаще, чем «ьняя» и, следовательно, при декодировании, beam search будет перевзвешивать вероятности путей («домашняя» и «домаьняя» и др.) с учетом частоты n‑gram внутри них. Таким образом, beam search и character LM могут существенно улучшить качество чтения модели (+5 п.п. accuracy на тетрадном датасете), декодируя её выход с учётом языковых особенностей домена.
В качестве beam search мы используем parlance ctcdecode в связке с kenLM. При подготовке лексикона для kenLM важно между буквами поставить пробелы, чтобы либа правильно собралась для character level метода, link.
Но beam search можно использовать не только на уровне символов (character level), но и на уровне слов (word level). В этом случае мы можем добавить словарь и таким образом ограничить предсказания только этим словарём. Сама OCR останется такой же, изменится декодирование выхода: beam search в процессе декодирования матрицы OCR будет оставлять только те пути, которые могут привести к словам из нашего лексикона (если путь отсутствует в словаре, то он исключается из списка доступных путей в beam).
На выходе у нас получится несколько слов, и выбрать из них одно снова поможет language model. Только в этот раз LM обучена не на последовательностях символов, а на сочетаниях слов. Если мы декодировали два слова «работа» и «забота», и знаем, что предыдущее слово было «домашняя», то благодаря LM мы можем скорректировать вероятности слов и выбрать «домашняя работа», потому что такое сочетание слов чаще встречается в нашем домене. Подробнее об этом можно почитать здесь и здесь. Но word level декодинг будет лучше подходить для задач ASR (распознавание аудиопотока), так как речь человека действительно более или менее ограничена словарём, а текст на картинках включает в себя сокращения, знаки препинания и цифры. Библиотеки для world level beam search: torchaudio и pyctcdecode, CTCWordBeamSearch.
beam search в seq2seq
Для seq2seq-архитектур beam search будет работать несколько иначе. CRNN на выходе дает итоговую матрицу со всеми вероятностями по алфавиту по всей ширине картинки, так что мы сразу получаем возможность проводить декодинг. Seq2seq же проводит предсказания в цикле: смотрит на картинку и предсказывает первую букву, далее снова смотрит на картинку и, принимая во внимание первую букву, предсказывает вторую, и так далее.
Таким образом, seq2seq-модель делает следующее предсказание с учетом ранее предсказанной буквы. А значит, если мы хотим получить от seq2seq несколько вариантов текста, то и forward pass нужно будет запускать неоднократно, увеличивая время работы сети.
Извлечение связного текста
Когда мы объединим детекцию текста и чтение в один пайплайн (вырезаем найденные слова и подаем их в OCR‑модель), то новой задачей станет извлечение связного текста из таких предсказаний (анализ layout страницы — макета/расположения текста). Ведь модель детекции не гарантирует, что полигоны располагаются на странице строго слева‑направо и сверху‑вниз.
Человек не пишет слова ровно по горизонтали: каждое слово может иметь свой наклон, а сама строка сильно «крениться» в концу листа. Страница может иметь рельеф и лежать неровно на поверхности. Мы можем написать эвристику на постпроцессинге: например, конец текущего слова находится рядом с началом следующего, и значит, эти слова на одной строке. Но всегда найдется настолько смещённое слово, которое сломает if»ы, от чего собьются все последующие номера линий на странице.
В датасете тетрадей у полигонов слов есть параметр group_id, отвечающий на номер строки. Для слов с одинаковым group_id мы можем нарисовать линию, которая будет как бы перечёркивать эти слова. Затем эту линию обучим в сегментации отдельным классом и на постпроцессинге сможем понять, что полигоны слов находятся на одной строке, если они пересекаются одной линией. Альтернативное решение — написать классификатор/кластеризатор по этим group_id.
Пример
В других датасетах может появиться необходимость писать модели по парсингу таблиц (строки/колонки), или нахождению абзацев текста, оглавления и т.д.
Разметка по строкам
Восстановление угла поворота картинки
Текст под наклоном ухудшает качество чтения. Особенно это актуально для CRNN‑архитектуры, которая читает текст как бы слева‑направо, предсказывая букву для каждого «столбца» OX на изображении. В итоге мы можем упростить работу OCR, если заранее приведём текст на странице в горизонтальный вид.
Модель сегментации для каждого слова на картинке предсказывает полигон, по координатам которого мы можем вытащить угол поворота (используя cv2.fitLine). Но из‑за неоднородности контуров самих полигонов угол поворота часто не будет соответствовать реальному наклону самого слова, и мы только увеличим долю ошибок на этапе чтения.
Лучше будет, если за угол поворота слова мы будет брать угол линии, на которой находится данное слово. Напомню, что линию нам предсказывает сегментация. Или же мы можем усреднить углы всех линий/слов на странице, повернув сразу всю картинку. В этом случае мы восстанавливаем наклон страницы. Наша команда остановилась на последнем варианте, но для каждой задачи можно пробовать и сравнивать разные варианты.
Еще один способ — сегментировать всю страницу на фотографии (в датасете тетрадей есть разметка полигона страниц). Далее — перспективным преобразованием по точкам страницы восстанавливать ей прямоугольное положение.
У такого метода есть и минусы. Например, когда углы страницы скрыты за рамками изображения, а также большая погрешность преобразования из‑за ошибки сегментации в нахождении координат одного из углов. Их всего 4, соответственно, и вклад каждого угла в точность восстановления большой.
Снижение качества при объединении моделей
При объединении двух моделей важно учитывать, чтобы не падала общая метрика, о которой мы поговорим чуть позднее. Дело в том, что OCR мы обучаем и тестируем на кропах, вырезанных по разметке, сделанной людьми. Но на инференсе кропы с текстом вырезаются уже по предсказаниям сегментации, которые не смогут в точности повторить разметку человека. В итоге процент ошибок чтения текста увеличивается просто за счет того, что на вход OCR принимает данные, которых она не видела во время обучения. Далее рассмотрим несколько способов уменьшить такой даунгрейд качества пайплайна.
Отличие масок разметки от предсказаний
Первый вариант — во время обучения OCR вырезать кропы слов не ровно по разметке, а добавляя немного фона вокруг слов. Так мы увеличим размер картинок, что позволит смелее применять аугментации вида random crop во время обучения. На картинках, вырезанных ровно по разметке, текст располагается слишком близко к краям изображения, и random_crop будет иногда отрезать буквы (аннотация при этом будет оставаться для полного текста).
Можно написать и такую аугментацию: на трейне добавлять случайное значения к координатам разметки (x, y), смещая их, и вырезать каждый раз немного другую область из исходной большой картинки. Особенно хорошо такой трансформ сработает, если OCR обучаем по маскам слов — так мы пытаемся имитировать погрешности сегментации. Вырезать текст по маскам вместо bbox»ов может быть удобнее в датасетах с «кучным» текстом, когда при их кропе по bbox»ам будет захватываться лишний текст с других строк.
Также можно обучить OCR совместно с сегментацией. Сегментация на этом этапе должна быть обучена, а её веса заморожены. Сегментация будет искать слова на странице, затем вырезать их и далее подавать на вход OCR, которая в свою очередь уже предсказывает текст. Backprop при этом далее делаем только для OCR. Здесь самая хитрая часть — для полигонов, предсказанных сегментацией, найти таргет текст из разметки. Можно сматчить по IoU предсказанные полигоны с ground truth полигонами и таким образом взять таргет‑текст. Понятно, что в OCR не нужно подавать кропы, для которых не был найдет таргет‑текст (false positive сегментации).
Метрика пайплайна
В целевой метрике важно отразить точность чтения текста с учётом ошибок в детекции: false positive (модель задетектила текст там, где его нет) и false negative (модель не задетектила текст, хотя он на самом деле есть). Но для оценки пайплайна мы не можем усреднить отдельные метрии двух сеток — сегментации и OCR — как раз по причине влияния качества сегментации на OCR. Соответственно, метрику нужно сделать «сквозную», учитывающую ошибки двух моделей.
Выглядеть такая метрика может следующим образом:
Для всех предсказанных полигонов матчим таргет текст из разметки (по IoU).
Если не нашлось ни одного полигона из разметки для данного предсказанного полигона, то для него в качестве gt‑текста оставляем пустоту «« (это будут FP сегментации).
Если у нас остались gt‑полигоны, не сматченные ни с одним предсказанным полигоном, то мы добавляем gt‑текст к отсутствующему полигону (это FN сегментации, для таких случаев OCR предсказывать ничего не будет).
Далее из всех предсказанных полигонов мы кропаем слова и читаем текст OCR‑моделью. Таким образом:
Предсказанный полигон | Ground truth полигон | Предсказанный текст | Ground truth текст | Комментарий |
pred_polygon_0 | gt_polygon_0 | «Комашняя» | «Домашняя» | |
pred_polygon_1 | gt_polygon_1 | «работа» | «работа» | |
pred_polygon_2 | - | «аваа» | «» | Это FP сегментации, детекция текста там, где его нет. OCR здесь может прочитать случайные буквы. |
- | gt_polygon_3 | «» | «Классная» | Здесь нет предсказанного полигона, а значит нет и запуска OCR-модели. |
В итоге мы имеем списки из gt‑слов и список из предсказанных слов и можем подсчитать метрики — accuracy, Левенштейна, CER и др., сравнивая два этих списка. Благодаря такой метрике мы будем штрафовать пайплайн за ложные срабатывания в детекции слов, а также за незадетектированные слова. Всё это отражается в одной итоговой цифре, по которой удобно сравнивать модели из разных экспериментов. Также это позволяет увидеть ситуации, когда увеличение метрики по отдельной модели не приводит к росту качества на всём пайплайне, что бывает часто.
Заключение
Подведем небольшой итог. Мы рассмотрели особенности рукописных датасетов на примере школьных тетрадей — такие как неразборчивость почерка, «кучность» текста, разные наклоны у строк и др. Посмотрели, как можно разделять объекты одного класса при семантической сегментации (как, например, предсказание масок границ для объектов), и также изучили метод beam search для декодирования предсказаний текста.
Важный этап задачи распознавания текста — извлечение связного текста из предсказаний. Модель должна уметь извлекать строки и знать последовательность слов внутри строки. Хорошо учитывать это на этапе планирования разметки, чтобы впоследствии не пришлось доразмечать датасет или менять архитектуру моделей. Также мы изучили причины даунгрейда качества чтения текста при объединении моделей и написали метрику, которая помогает численно оценивать качество пайплайна.
Вы можете взять за основу наш код для чтения текста: модели сегментации и OCR — эти репы, кстати, можно использовать и по отдельности в других задачах (сами модели можно конвертировать в ONNX и ускорить рантайм на CPU с использованием openvino). И третья репа — для объединения первых двух моделей в единый пайплайн. В конфиге пайплайна блоками прописывается основной цикл инференса на изображении: запуск сегментации, восстановление наклона страницы, кроп текста и запуск OCR, постпроцессинг и т. д. В репах есть примеры запуска на Google Colab. А пример деплоя можно найти на hugging face spaces (инференс занимает время, т.к. в базовых серверах на hugging face выделяют всего 2 cpu)
Ну а на этом все. Очень надеюсь, что статья была полезной, и вы узнали что‑то новое. Если у вас есть собственные кейсы в этом направлении — поделитесь, пожалуйста, в комментариях.
Ссылки из статьи:
Датасет тетрадей на русском языке и на английском
Наш код по распознаванию: OCR, сегментация, пайплайн
Пример деплоя распознавания.
Коллектив авторов Sber AI и Sber IDP: @ddimitrov @kuznetsoff87 @JuliaAgafonova @gazizovmarat
