Threshold U-Net: как мы отказались от высокого разрешения и выиграли в скорости бинаризации

Хоть современный ИИ уже почти в состоянии написать симфонию и превратить холст в шедевр, некоторые простые задачи все еще не так просты, если наложить ограничения на решение. Так как главной задачей у нас, в Smart Engines, является распознавание и оцифровка документов, то для нас ограничениями являются скорость работы и возможность запуска системы на вычислительно слабых конечных устройствах. В этой статье мы вспомним об одной классической задаче обработки изображений – бинаризации документов. А также расскажем, как нам удалось существенно оптимизировать нейросетевую модель бинаризации, сохранив исходное качество, и причем тут современные модели детекции текста.

Бинаризация документов… это еще актуально?

Начнем с того, что ответим на важный вопрос – а зачем вообще нужна бинаризация документов в 2025 году? Ведь все современные системы распознавания документов основаны на нейросетях, которые способны обрабатывать сложные входные данные и не требуют какого-либо препроцессинга… И в целом это будет правдой: непосредственно для распознавания текста она не нужна. Нужна же она для применения в не-нейросетевых алгоритмах, которых также предостаточно: например, при выделении линий таблиц, или сжатии исходного изображения. Последнее крайне актуально: во многих сферах, будь то бизнес или архивы с историческими документами, оцифровка документов не заканчивается на извлечении ключевой информации. Необходимо также сохранить само изображение документа, а для уменьшения размера файла используется бинаризация.

Задача бинаризации сама по себе довольно старая и хорошо изученная: всем известные метод Оцу (Otsu’s method) и метод Ниблэка (Niblack’s method) широко используются с 80-х годов прошлого века. Значительный скачок в качестве бинаризации произошел в 2015 году с появлением нейросетевой модели U-Net (Рис. 1). С тех пор, нейросетевые модели стали де-факто стандартом в задачах семантической сегментации. Как бы не называлась очередная представленная модель, топология ее архитектуры будет U-Net-подобной. Это значит, что она будет обязательно состоять из энкодера, выполняющего отображение входного изображения в латентное пространство, и декодера, который восстанавливает изображение из его латентного представления в высоком разрешении с сохранением мелких деталей.

В 2017 году с помощью U-Net нами даже был выигран [1, 2] конкурс бинаризации исторических документов DIBCO-2017, организованный на базе крупнейшей конференции по анализу документов ICDAR (International Conference on Document Analysis and Recognition). Казалось бы, инструмент для качественной бинаризации документов есть, проблему можно считать решенной. Однако на практике все уперлось в сложность применения модели на конечных устройствах: U-Net подобные модели вычислительно сложны в исполнении и потребляют много оперативной памяти. Главным образом из-за того, что они обрабатывают изображение в большом разрешении: в случае деловых документов, это более тысячи пикселей по ширине. Можно, конечно, поделить исходное изображение на перекрывающиеся небольшие патчи, обработать их по отдельности, и затем агрегировать результаты в финальное изображение. Однако это займет значительно больше времени, чем применение нейросети в “полносверточном” режиме. В целом, вопрос оптимизации модели стал крайне актуальным.

Оптимизация U-Net’а за счет переформулирования задачи

Первым шагом оптимизации любой нейросетевой модели является сокращение числа операций за счет уменьшения параметров слоев. Например, за счет уменьшения количества сверточных фильтров в слоях. Также, можно попробовать сквантовать сеть в 8-битный, 4-битный или даже 4.6-битный [3] формат инференса. Хоть это все и эффективно, но не решает ключевую проблему U-Net подобных архитектур для бинаризации, а именно – обилие сверточных слоев, работающих на картах признаков в высоком разрешении. И тут возник вопрос – а можно ли как-то упростить декодер, убрав из него слои, работающие на высоком входном разрешении? (Рис. 2)

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

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

Например, если вход был размера , а выход сети (предсказанная карта порогов) имеет размер , то каждый предсказанный порог служит для бинаризации области размера k*m на исходном изображении.

В случае документов, подбор такого порога задача несложная: классические локальные методы бинаризации в виде метода Ниблэка или Саволы показывают хорошие результаты при правильной настройке параметров (Рис. 3). В данном случае, большое рецептивное поле сети и автоматическая настройка параметров на обучающей выборке позволяют добиться хорошего результата при применении одного порога к небольшим регионам изображения: например, размера или .

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

Differentiable Binarization, или причем тут детекция текста

В 2019 году был опубликован доклад “Real-time Scene Text Detection with Differentiable Binarization” (в 2022 году вышла ее доработанная журнальная версия [5]), в котором авторы представили новый детектор текста DBNet. Как и многие другие аналогичные детекторы текста, выходом их модели является тепловая карта, на которой подсвечены слова и строки текста. Главное нововведение заключалось в том, что помимо тепловой карты текста авторы предложили предсказывать еще и карту порогов, с которыми необходимо ее бинаризовать (вместо фиксированного одного порога в качестве параметра постпроцессинга, как в других методах). Модель показала хорошее качество на основных бенчмарках, и была включена в такие open-source системы распознавания документов, как DocTR и PaddleOCR.

Для обучения такой модели авторы использовали дифференцируемую бинаризацию. Если обычная бинаризация выражается следующим образом:

то дифференцируемая бинаризация имеет следующий вид:

где – аппроксимированный бинарный выход, – предсказание сети, – предсказанный порог для бинаризации в данной точке. Коэффициент k подбирается эмпирически, у авторов он был равен 50. Как можно заметить, это простая логистическая функция, которая имеет непрерывную производную, в отличии от обычной функции бинаризации (Рис. 5). А это значит, что можно учить сети с таким модулем с обычным градиентным спуском.

При этом всем стоит отметить, что сама по себе дифференцируемая бинаризация не является чем то новым и доселе неизвестным: например, ее использовали для выражения метода Ниблэка в виде нейросетевой архитектуры в работе [6]. Однако успех модели DBNet напомнил сообществу о данном полезном модуле.

Модель Threshold U-Net

Таким образом, с применением дифференцируемой бинаризации для обучения модели предсказывать пороги не нужен ground truth в виде этих самых порогов. Достаточно лишь применить предсказанные пороги к входному изображению, а по полученному псевдобинарному изображению посчитать функцию ошибки, которая в нашем случае – бинарная кросс-энтропия. Т.е., для обучения необходим ground truth в том же формате, что и для обучения обычного U-Net. Архитектура модели, которую мы назвали Threshold U-Net, представлена на Рис. 6.

Для тестирования модели, мы обучили ее бинаризации исторических документов из наборов данных DIBCO. Сеть обучалась на патчах размера 128×128; были обучены вариации модели с размером выхода 16×16, 32×32 и 64×64 (т.е. один предсказанный порог отвечал за бинаризацию области размера 8×8, 4×4 и 2×2 соответственно). Результаты приведены в таблице 1: для замеров качества использовались метрики DRD (Distance-Reciprocal Distortion), F-score, Pseudo F-score (модификация F-score с учетом весов пикселей), а также PSNR (Peak signal-to-noise ratio).

А вот так выглядит выход сети (карта порогов) в сыром виде:

Ну и наконец результаты бинаризации на изображении документа из DIBCO-2017:

Заключение

“Пороговый” U-Net оказался заметно быстрее классической модели U-Net, при этом без особой разницы по качеству – как по метрикам, так и визуально. В тоже время, предложенная модификация проста в реализации как обучения, так и инференса.

Статья написана на основе нашего доклада на конференции ICMV-2023: K. E. Lihota, A. V. Gayer and V. V. Arlazarov, “Threshold U-Net: speed up document binarization with adaptive thresholds,” ICMV 2023, 13072 ed., Wolfgang Osten, Dmitry Nikolaev, Johan Debayle, Ed., Bellingham, Washington 98227-0010 USA, Society of Photo-Optical Instrumentation Engineers (SPIE), Apr. 2024, vol. 13072, ISSN 0277-786X, ISBN 9781510674622, vol. 13072, pp. 130720T1-130720T8, 2024, DOI: 10.1117/12.3023176.

А чтобы внимательно изучить все наши публикации – приглашаем вас к нам на сайт.