Где взять пару миллионов документов с разметкой для обучения модели ИИ? Сгенерировать синтетически! Меня зовут Никита, я работаю в ИСП РАН и веду блог по ИИ. Выложил в открытый доступ проект, который поможет дата-саентистам самостоятельно создавать датасеты сканированных документов с разметкой, используя только CPU. Подойдет для предобучения мультимодальных трансформеров, OCR и проч.
Кратко:
Генератор создает синтетические PNG-картинки сканированных документов с разметкой. Ссылка на гитхаб
Язык текста: можно выбрать.
Принцип работы: парсим страницу Википедии, переносим инфу в DOCX, рандомизируем стилизацию, из DOCX рендерим изображения, изображения аугментируем под сканы.
Разметка: bounding box (координаты ограничивающей рамки) каждого слова и само слово. Можно обучать модели VrDU, OCR и проч.
Содержание документа: абзацы текста, заголовки, таблицы. Есть различные форматирования: жирный, подчеркивание, курсив. Одна или две колонки. Рандомизация шрифтов, настроек форматирования.
Требуемое железо: только CPU.
Примеры изображений
Принцип работы
Требования к датасету просты:
Нужны изображения сканов документов, достаточно реалистичные и разнообразные;
В документах должен содержаться осмысленный текст с разным форматированием, требуется поддержка русского языка;
В разметке нужны слова и координаты их bounding box.
Текст на любом языке можно спарсить с интернета и получить форматирование через HTML, а затем разместить текст в DOCX. А там и рандомизировать стили и зарендерить изображение несложно. Однако тут встает вопрос – а как получить разметку? Можно воспользоваться OCR, но тогда неизбежно будут получаться ошибки распознавания. Хотелось бы как-то доставать координаты из DOCX напрямую, однако из коробки для этого способа нет. Пришлось изобретать!)
Идею подсказали коллеги Андрей Перминов и Настя Зыкина, которые создавали таким образом датасет документов для обучения детектора абзацев. Я улучшил алгоритм для детекции каждого слова и сделал из этого open-source систему.
Итак, в ворде есть инструмент "Заливка". Он создает цветной фон вокруг слова.
А еще есть инструмент "Цвет шрифта". Что если выбрать и там, и там одинаковый цвет? Получатся вот такие веселые прямоугольники вместо слов.
А вот веселые цветные прямоугольники на идеально белом фоне можно быстро находить классическим компьютерным зрением. Ставим границы всех таблиц в белый цвет, чтобы не мешали. Переводим DOCX в PNG, с помощью OpenCV бинаризуем изображение, находим контуры и извлекаем из них координаты прямоугольников.
Проблем быть не должно, так как документ не содержит ничего, кроме закрашенных прямоугольников на белом фоне, а изображение не содержит никаких помех.
Осталось сопоставить к каждому bbox'y текстовое слово. Сделать это несложно. Красить каждое слово будем в уникальный цвет. Затем сохраним соотношение HEX_код_цвета -> текст_слова в хэш-таблице. Тогда после детекции bbox'a сможем достать из хэш-таблицы слово, используя цвет прямоугольника. При глубине цвета 8 бит можно закодировать 
слов (вычитаем белый цвет, так как он сольется с фоном). PNG сжимает изображения без потерь, RGB-значения в каждом пикселе будут точно соответствовать цвету прямоугольника в DOCX документе. Поэтому способ сработает и для близких оттенков.
Пока что генератор поддерживает в качестве источника текстов Википедию. Полный алгоритм генерации одного скана представлен на блок-схеме ниже:
Архитектура системы
Теперь о технической реализации. Сначала инициализируется класс Manager и создает список ссылок. Парсинг ссылок работает аналогично алгоритму обхода графа в ширину. Сначала сканируются все ссылки на стартовой странице, затем их дочерние ссылки и так далее, а для хранения непросмотренных URL используется очередь. Начальная страница задается пользователем. По умолчанию стоит начальная страница английской Википедии, но можно поставить страницу любого языка.
Затем создаются несколько DocumentGenerator – по одному на каждый процесс. Каждый экземпляр получает подмножество из списка URL и для каждой ссылки запускает приведенный выше алгоритм.
Вся работа с ворд-документом ведется через python-docx. На данный момент бутылочное горлышко всего проекта – перевод DOCX в PNG. Я использую Unoserver для перевода в PDF, а затем pdf2image для рендера 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, которая заточена под реалистичные "скановые" аугментации. Пришлось даже сделать мердж реквест, чтобы поправить там баг. Также для ускорения операций с изображениями используется Pillow-SIMD.
Так и получился DOcument GEnerator – DoGe.
Итого
Надеюсь, DoGe поможет вам в обучении ваших моделей. Код системы можно найти в открытом репозитории. Я открыт для идей по улучшению генератора, а также помощи в их реализации. По плану:
Добавление парсинга и разметки изображений
Добавление разметки заголовкам, таблицам, абзацам
Добавление новых форматов вывода (например, Parquet)
Добавление дополнительной информации в разметку через языковые модели
Повышение производительности: узким местом является преобразование Docx -> Pdf -> Png.
Подписывайтесь на мой телеграм канал про ИИ, где я делюсь своим опытом в сфере глубокого обучения. Минимум репостов новостей, только свой опыт и мысли.