Генератор синтетических документов для обучения моделей ИИ

Где взять пару миллионов документов с разметкой для обучения ^[1] модели ИИ? Сгенерировать синтетически! Меня зовут Никита, я работаю в ИСП РАН и веду блог по ИИ ^[2]. Выложил в открытый доступ проект, который поможет дата-саентистам самостоятельно создавать датасеты сканированных документов с разметкой, используя только CPU. Подойдет для предобучения мультимодальных трансформеров, OCR и проч.

Кратко:

1. Генератор создает синтетические PNG-картинки сканированных документов с разметкой. Ссылка на гитхаб ^[3]

2. Язык текста: можно выбрать.

3. Принцип работы: парсим страницу Википедии, переносим инфу в DOCX, рандомизируем стилизацию, из DOCX рендерим изображения, изображения аугментируем под сканы.

4. Разметка: bounding box (координаты ограничивающей рамки) каждого слова и само слово. Можно обучать модели VrDU, OCR и проч.

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

6. Требуемое железо: только CPU.

Примеры изображений

Принцип работы

Требования к датасету просты:

1. Нужны изображения сканов документов, достаточно реалистичные и разнообразные;

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

3. В разметке нужны слова и координаты их bounding box.

Текст на любом языке можно спарсить с интернета и получить форматирование через HTML, а затем разместить текст в DOCX. А там и рандомизировать стили и зарендерить изображение несложно. Однако тут встает вопрос – а как получить разметку? Можно воспользоваться OCR, но тогда неизбежно будут получаться ошибки ^[4] распознавания. Хотелось бы как-то доставать координаты из DOCX напрямую, однако из коробки для этого способа нет. Пришлось изобретать!)

Идею подсказали коллеги Андрей Перминов ^[5] и Настя Зыкина ^[6], которые создавали таким образом датасет документов для обучения детектора абзацев. Я улучшил алгоритм для детекции каждого слова и сделал из этого open-source систему.

Итак, в ворде есть инструмент “Заливка”. Он создает цветной фон вокруг слова.

А еще есть инструмент “Цвет шрифта”. Что если выбрать и там, и там одинаковый цвет? Получатся вот такие веселые прямоугольники вместо слов.

А вот веселые цветные прямоугольники на идеально белом фоне можно быстро находить классическим компьютерным зрением ^[7]. Ставим границы всех таблиц в белый цвет, чтобы не мешали. Переводим DOCX в PNG, с помощью OpenCV бинаризуем изображение, находим контуры и извлекаем из них координаты прямоугольников. 

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

Осталось сопоставить к каждому bbox’y текстовое слово. Сделать это несложно. Красить каждое слово будем в уникальный цвет. Затем сохраним соотношение HEX_код_цвета -> текст_слова в хэш-таблице. Тогда после детекции bbox’a сможем достать из хэш-таблицы слово, используя цвет прямоугольника. При глубине цвета 8 бит можно закодировать слов (вычитаем белый цвет, так как он сольется с фоном). PNG сжимает изображения без потерь, RGB-значения в каждом пикселе будут точно соответствовать цвету прямоугольника в DOCX документе. Поэтому способ сработает и для близких оттенков. 

Пока что генератор поддерживает в качестве источника текстов Википедию. Полный алгоритм генерации одного скана представлен на блок-схеме ниже:

Архитектура системы

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

Затем создаются несколько DocumentGenerator –  по одному на каждый процесс. Каждый экземпляр получает подмножество из списка URL и для каждой ссылки запускает приведенный выше алгоритм.

Вся работа с ворд-документом ведется через python-docx ^[8]. На данный момент бутылочное горлышко всего проекта – перевод DOCX в PNG. Я использую Unoserver ^[9] для перевода в PDF, а затем pdf2image ^[10] для рендера PNG, что достаточно затратно. Если кто знает, как можно это сделать проще – пишите!

Для ускорения поднимаются несколько экземляров Unoserver’a в разных процессах. Внутри каждого DocumentGenerator зашит мультитрединг, который не позволяет простаивать системе, пока Unoserver рендерит PDF. Так удалось добиться генерации одного скана в среднем за 0.5 секунды на Intel(R) Xeon(R) Gold 6338.

Алгоритм поиска контуров работает на OpenCV:

thr = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
thr = cv2.threshold(thr, 254, 255, cv2.THRESH_BINARY_INV)[1]
cnts = cv2.findContours(thr, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[0]

for c in cnts:
    peri = cv2.arcLength(c, True)
    approx = cv2.approxPolyDP(c, 0.015 * peri, True)

    if len(approx) == 4:
        x, y, w, h = cv2.boundingRect(approx)
        rgb_color = image_pil.getpixel((x+1, y+1))
        color = '#%02x%02x%02x' % (rgb_color)

Для аугментаций используется библиотека Augraphy ^[11], которая заточена под реалистичные “скановые” аугментации. Пришлось даже сделать мердж реквест, чтобы поправить там баг. Также для ускорения операций с изображениями используется Pillow-SIMD ^[12].

Так и получился DOcument GEnerator – DoGe.

Итого

Надеюсь, DoGe поможет вам в обучении ваших моделей. Код системы можно найти в открытом репозитории ^[13]. Я открыт для идей по улучшению генератора, а также помощи в их реализации. По плану:

• Добавление парсинга и разметки изображений

• Добавление разметки заголовкам, таблицам, абзацам

• Добавление новых форматов вывода (например, Parquet)

• Добавление дополнительной информации в разметку через языковые модели

• Повышение производительности: узким местом является преобразование Docx -> Pdf -> Png.

Подписывайтесь на мой телеграм канал про ИИ ^[2], где я делюсь своим опытом ^[14] в сфере глубокого обучения. Минимум репостов новостей, только свой опыт и мысли.