Обучаем SVTR-Tiny для распознавания текста сцены15.12.2022 20:01

Сегодня мы расскажем вам, как дообучить новую state-of-the-art модель SVTR-Tiny для распознавания текста сцены (текста в реальных уличных условиях) на собственноручно сгенерированных изображениях с помощью API библиотеки PaddleOCR.

В нашей компании компьютерное зрение и, в частности, распознавание текста используется в сервисе проверки контрагентов (СПК) DataNewton для перевода изображений товарных знаков в текстовый формат и для дальнейшего использования этой информации в аналитических целях. Также у нас есть заказчики, для которых мы применяем компьютерное зрение в задачах распознавания текста в документах.

Глобально задача распознавания текста состоит из нескольких частей:  

1. Модель детекции/сегментации определяет область, где находится текст. 

2. Область с текстом из изображения вырезается и модель распознавания считывает этот текст. 

3. Между этими этапами может идти различный препроцессинг, такой как поворот изображения, нормализация и т.д. 

В этой статье мы сконцентрируемся на последнем этапе всего пайплайна OCR (optical character recognition), то есть на распознавании текста, сгенерируем подходящий для этого датасет и обучим SVTR-Tiny.

О задаче и разметке

Задача распознавания текста сцены весьма распространенная. Алгоритмы распознавания текста в реальных уличных условиях все еще разрабатываются и можно сказать, что мы являемся свидетелями эволюции методов детекции и распознавания текста сцены. Существует много размеченных датасетов, которые уже можно использовать для тренировки моделей детекции и распознавания. Если модели детекции вы сможете натренировать на датасетах, состоящих из слов на иностранных языках, то модели распознавания — вряд ли. Поэтому разметка под конкретный язык в данном случае играет важную роль. Идеальным вариантом будет, если вы сгенерируете свой тренировочный датасет. Для генерации мы выбрали SynthText.

От других библиотек генерации изображений с текстом эта библиотека отличается по нескольким параметрам:  

• высокая вариативность генерируемых изображений,  

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

• большое количество фоновых изображений,  

• использование алгоритмов для попадания слов в специальные кластеризованные регионы изображения. 

Для нашей задачи мы модифицировали оригинальный SynthText под русский язык, эту реализацию и будем использовать.

Об архитектуре

Архитектура модели SVTR уже была детально описана в статье на Хабре. Мы можем лишь дополнить, что global mixing и local mixing слои нейронной сети действительно повышают возможности модели в распознавании глобальных и локальных признаков и убедились в этом на собственном примере. На нашей выборке, состоящей из 566 обрезанных картинок с текстом, SVTR-Tiny дает лучший результат, нежели классический CRNN с resnet34 в качестве backbone. Но подробнее о результатах расскажем в конце статьи, а сейчас попытаемся сгенерировать картинки для дообучения SVTR-Tiny и проверить результаты ее работы.

Генерация текстовых данных

Модели распознавания текста обучаются и работают с уже обрезанными изображениями c текстом:

Процесс получения таких изображений производится в несколько этапов, для этого необходимо:

1. Подготовить и сгенерировать текстовый датасет.

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

3. Из сгенерированных изображений с текстом вырезать отдельные изображения со словами.

4. Передать обрезанные изображения со словами в модель для дообучения.

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

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

1. Очистить текстовые данные, удалить лишние символы.

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

Жирным выделены буквы из разных алфавитов, но идентичные или похожие по начертанию

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

2. Преобразовать регистр. Следует заметить, что если в данных, на которых ваша модель будет работать, больше слов в нижнем регистре, то и в генерации нужно ставить соответствующие значения вероятностей. Например, простая реализация функции на Python с равными вероятностями 0.33, 0.33, 0.33 для преобразования слов в нижний регистр, верхний регистр и в слова с заглавной буквы:

def preprocess_word_case(text):
    text = text.lower()
    random_number = random.choices([0,1,2], weights=[0.33, 0.33, 0.33], k=1)[0]
    if random_number == 0:
      text = text.capitalize()
    if random_number == 1:
      text = text.upper()
    if random_number == 2:
      pass
    return text

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

2. Не стоит забывать и о числовых данных и буквах, отдельных от других слов, таких как «и», «в», «к» и т.д. и добавить их к основному датасету. 

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

После всех проделанных манипуляций так выглядят сгенерированные текстовые данные для генерации изображений в формате .txt, состоящие из английских и русских слов.

Подготовка к генерации изображений с текстом

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

├── bg_data
│   └── bg_img
├── data
│   ├── fonts
│   │   └── vn
│   ├── models
│   └── newsgroup
└── prep_scripts

Для того, чтобы иметь возможность генерировать картинки, предварительно необходимо выполнить несколько шагов:

1. Загрузить торрент и скачать файлы, относящиеся к папке bg_data в торрент-файле.

2. Файл bg_img.tar.gz разархивировать в /bg_data/bg_img.

3. Файлы с собственными шрифтами (если планируете их добавлять) с расширениями .ttf .otf .ttc необходимо сложить в папку /data/fonts/vn.

4. Добавить в файл /data/fonts/fontlist.txt дополнительные шрифты.

5. Перед запуском скрипта генерации картинок необходимо обновить модель шрифтов font_px2pt.cp, если ранее вы добавляли новые шрифты. Она находится в папке /data/models. Для этого нужно предварительно запустить скрипт invert_font_size.py, который лежит в корневой директории:

$ cd SynthText-Russian
$ python invert_font_size.py

1. В корневой директории проекта сгенерируется новый файл font_px2pt.cp, его необходимо перенести и заменить старый в /data/models.

2. Файл new_data.txt с нашими сгенерированными текстовыми данными из предыдущего раздела необходимо положить в /data/newsgroup.

Генерация изображений с текстом

Изменяемые параметры для генерации изображений находятся в файле gen.py, который лежит в корневой директории папки с проектом:

NUM_IMG — количество картинок, используемых для генерации. Чтобы использовать все имеющиеся картинки, задайте параметр »-1».

INSTANCE_PER_IMAGE — количество проходов/итераций по одной картинке для генерации. Например, при значении »5» будет сгенерировано 5 различных текстовых вариаций с одной картинкой.

SECS_PER_IMG — время ожидания на генерацию одной картинки.

DATA_PATH — путь к папке с шрифтами, моделью шрифтов, текстовым данным и т.д.

OUT_FILE — результирующий файл с метаданными и сгенерированными картинками.

Теперь можно запускать наш скрипт генерации картинок с текстом:

$ python gen.py

Вот так приблизительно могут выглядеть сгенерированные изображения с текстом:

Сгенерированный файл появится в /results/SynthText.h5.

Преобразование сгенерированных изображений в формат PaddleOCR

На этом этапе из файла SynthText.h5 необходимо извлечь вырезанные слова с текстовой аннотацией для тренировки нейронной сети. В .ipynb ноутбуке генерируется scene_text_recognition.zip файл с тренировочным датасетом, который можно передать в PaddleOCR для тренировки модели распознавания текста. Все тонкости и нюансы преобразования выходного файла SynthText.h5 сгенерированных изображений вы можете посмотреть там.

А вот так выглядят вырезанные части сгенерированных нами изображений, на которых будем дообучать модель распознавания текста:

Дообучение SVTR-Tiny модели на сгенерированных и вырезанных изображениях

Теперь мы можем взять наш scene_text_recognition.zip, который мы получили, используя .ipynb ноутбук из предыдущего раздела, и использовать его для дообучения модели распознавания текста SVTR-Tiny:

1. Установить библиотеки и зависимости:

$ python3 -m pip install paddlepaddle-gpu==2.3.1
$ git clone https://github.com/PaddlePaddle/PaddleOCR
$ cd PaddleOCR
$ pip3 install -r requirements.txt
$ pip install "paddleocr>=2.0.1"

1. Извлечь датасет:

$ unzip -qo '/PaddleOCR/train_data/scene_text_recognition.zip'

1. Загрузить предобученную моделиь «rec_svtr_tiny_none_ctc_en» с этой страницы и переместить ее в /PaddleOCR/pretrain_models.

2. Добавить список символов в файл /ppocr/utils/dict/rus_chars.txt. 

3. Запустить тренировку модели с конфигурацией rec_svtrnet_habr.yml, которую необходимо поместить в /configs/rec по ссылке.

$ python3 tools/train.py -c configs/rec/rec_svtrnet_habr.yml -o Global.pretrained_model=./pretrain_models/rec_svtr_tiny_none_ctc_en_train/best_accuracy

Все шаги и дообучение модели в Google Colab вы можете найти в ноутбуке.

После файнтюнинга в течении 3-х часов на выборке из 278 тысяч сгенерированных изображений, SVTR-Tiny давала на нашей тестовой выборке из 566 изображений товарных знаков среднюю точность по Левенштейну 88,7%, а в случае индифферентности к регистру — 92,96% (в представленном ноутбуке модель тренировалась на 4 тысячах сгенерированных данных).

Наш тестовый датасет из 566 изображений товарных знаков выглядит примерно так:

SVTR-Tiny выигрывает у классической CRNN с resnet34 в качестве backbone как по скорости работы, так и по точности, что действительно заслуживает внимания:

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

© Habrahabr.ru