Привет, Хабр! Меня зовут Владислав, я CV Engineer в компании YADRO. В этой статье я расскажу, как мы разрабатывали и обучали алгоритм детекции документов для нашего планшета Kvadra_T. Я подробно описал нюансы задачи и весь наш путь — от классического подхода до генерации недостающих датасетов и обучения на них нашей собственной нейросети. Постарался сделать историю интересной как для новичков в теме, так и для более опытных читателей. Режим детекции, кстати, уже доступен в kvadraOS.

Сканирование документов — стандартная функция телефонов и планшетов. Наводишь камеру на документ, видишь его границы в режиме реального времени, делаешь снимок и получаешь скан без фона. Обычно в этом режиме фотографируют паспорта, водительские права, листы А4, книги и так далее.

Задача детекции документов кажется в целом простой, так как поиск объектов на фото — это уже хорошо изученный класс задач. Основные сложности связаны с тем, что модель должна работать в реальном времени и что документы фотографируют в разных условиях, — это может заметно отразиться на результате.

После анализа требований к функционалу мы свели задачу детекции к трем пунктам:

1. На вход модели подается фотография одного документа (и не более, это обязательное условие).

2. Модель предсказывает координаты четырех углов документа.

3. Первоначальный вид документа восстанавливается через перспективное преобразование по четырем заданным опорным точкам.

Решаем классическим методом

Вначале мы решили попробовать классический подход без нейронных сетей, поскольку задача кажется простой и, следовательно, решаться она должна с использование простых алгоритмов CV. Взяли готовый алгоритм компьютерного зрения с открытым кодом и переписали на C++. Порядок работы алгоритма:

1. На вход идет изображение документа.

2. Границы документа определяются с помощью фильтра Canny.

3. Определяются и удаляются с изображения пиксели, соответствующие тексту.

4. На картинке при помощи HoughLines определяются все линии и их пересечения.

5. Перебираются все возможные четырехугольники.

6. Выбирается самый большой четырехугольник, который принимается за документ.

На мобильном устройстве алгоритм заработал со скоростью 20–100 мс (10–50 fps). Это решение мы использовали как временное, пока параллельно разрабатывали нейросеть.

Предпоследний этап занимает бо́льшую часть времени работы алгоритма. Я обозначил средний диапазон в 20–100 мс, но в зависимости от числа линий время может вырасти до целой секунды, то есть уже совсем не в реальном времени. Иногда алгоритм лагает, иногда его путают соседние объекты, в результате чего за углы документов принимаются пересечения совсем других линий. С учетом скорости работы и числа ошибок останавливаться на этом решении было нельзя.

Нейросетевой подход

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

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

• SmartDoc 2015 Challenge 1. Документы на листах А4, простой фон, горизонтальное положение камеры (~24 000 изображений).

• MIDV 500, MIDV 2019. Размеченные водительские права и паспорта разных стран, изображения на сложном фоне, в том числе частично обрезанные (~36 000 изображений). 

• Invoice3D. Синтетический датасет из отрисованных 3D-моделей мятых и изогнутых документов (~35 000 изображений).

• COCO 2017. Размеченные фотографии объектов, похожих на документы — экраны, ноутбуки, коробки, телефоны (~9 000 изображений). Необходим для детекции подобных предметов. В нем не все объекты имеют четыре угла — позже расскажу, зачем нужен был такой датасет и как мы решили проблему с углами.

В сумме во всех датасетах у нас было около 105 000 картинок. Их лицензии не разрешают использование в коммерческих продуктах, поэтому для релизной модели, которая работает в планшете, мы взяли свой датасет, состоящий из сгенерированных (об этом ниже) и собранных разметчиками изображений. Его мы сбалансировали так, чтобы он был похож на тот, что мы использовали для экспериментов.

Переходим к самой нейросети. Мы взяли за основу модель LDRNet, которая предсказывает углы и границы документа (около 100 точек). Недостаток LDRNet в том, что ее тренировали только на горизонтальных изображениях с документом, поэтому она ошибается на любых других. Также я считаю скорее недостатком, что LDRNet написана на TensorFlow — это затягивает создание пайплайнов.

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

Минусов у нее много, но в статье были хорошие метрики, и нейросеть работала в реальном времени. Так что мы решили не отказываться от модели, а переписать ее на Pytorch Lightning под новым названием QuadraTNet и дотренировать на большем количестве юзкейсов. Вот как модель стала выглядеть после наших изменений:

Разберем эту схему пошагово. Модель принимает на вход изображение размером 224 на 224 пикселя — к этому формату приводится любое входящее фото без сохранения соотношения сторон. Далее мы раскладываем обычную картинку на несколько входов, которые конкатенируются. Первый вход — это фильтр Собеля. Он находит границы изображения и не требует много ресурсов. Эта небольшая модификация позволяет добавить инвариантность к входу: для похожих картинок границы будут одинаковые.

Фильтр не требует изменения препроцессинга, его можно зашить в модель при помощи свертки. Сеть с таким изменением ведет себя стабильнее при обучении.

Фильтр Собеля давал результат только в одном цвете, а нейросеть принимала три, так что мы добавили дополнительный вход — CoordConv. По сути, это картинка из двух каналов, где вместо изображения записаны координаты каждого пикселя. Она нужна, чтобы модель лучше понимала, где какие пиксели находятся, и ей было легче предсказать координаты:

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

Теперь расскажу об архитектуре нейросети. Изначально это была MobileNetV2, но мы взяли MobileNetV3 — улучшенную версию, которая занимает меньше места. К ней добавили пару слоев и пулингов на выходе, которые в ходе экспериментов улучшали результат.

Переходим к выходу — по сути, это и есть предсказание сети, которое можно превратить в координаты на изображении. Исходная модель содержала 100 точек на выходе, и предсказываются они не всегда корректно. Чтобы исправить это, мы сократили количество предсказываемых точек до четырех — остальные нам не нужны.

Также я исправил классификацию: просто добавил в датасет картинки без документов. И сделал два линейных слоя вместо одного: на скорость и качество это не влияет. Так исторически сложилось из-за классификации: Cross-entropy Loss влиял на регрессию координат и поэтому MLP обучался отдельно. Для него мы делали detach после пуллинга, но тогда регрессия начинала влиять на классификацию, если точки чуть съезжали в сторону. После замены на FocalLoss сеть сразу начала лучше классифицировать и предсказывать точки на документе. Чтобы не проводить дополнительные эксперименты, мы оставили два линейных слоя и обучали их совместно. 

Обучение нейросети

Для обучения сети ей нужно как-то показать, как сильно она ошиблась. То, насколько предсказание далеко от истинно верного значения (ground truth), измеряет функция ошибки, или лосс (loss function). На основе лоссов работает метод оптимизации для настройки параметров, или весов модели. В нашей задаче корректировать надо так, чтобы сеть выдавала правильные координаты и вероятность присутствия документа на изображении.

Charbonnier Loss

Для координат обычно используют среднее квадратичное расстояние для каждой точки (MSE Loss). В ней по каждой оси ошибка считается отдельно. На практике лучше взять корень ошибки для каждой точки (RMSE Loss), ведь если мы попадем по одной оси точно, а по другой сильно ошибемся, это будет сильно хуже, чем меньшая ошибка сразу по двум осям. Под знак корня добавляют небольшое положительное значение epsilon (Charbonnier Loss), чтобы обеспечить стабильность значений производной при малых ошибках. Во время тренировки они часто приводят к значению NaN.

TClip Loss

Следующая функция — TСlip Loss. Это дополнительная функция ошибки, предписывающая сети выдавать координаты в ограниченном диапазоне, чтобы точки не выходили за границы изображения. Этот лосс нам необходим, поскольку в некоторых наборах данных документы видны лишь частично.

Focal Loss

Кроме того, у нас есть классификация, чтобы определять, есть ли вообще документ на изображении. Ниже приведено два графика для обоих лоссов: синий (Cross-entropy Loss) — исходный, оранжевый (Focal Loss) — тот, на который мы его заменили.

Синий график легко можно продолжить, он достигнет оси под углом, а оранжевый будет к ней только приближаться. При большой уверенности нейронной сети Cross-entropy Loss будет тянуть ее дальше, хотя в принципе уже некуда, поскольку вероятность ограничена 1. Оранжевый (Focal Loss) подходит нам лучше, так как несильно выкручивает результат, если сеть уже угадала наличие документа. Также классификация может влиять на предсказание координат, поскольку нейросеть во время обучения будет делать больший акцент на наличии документа, чем на координатах.

Chamfer Loss

Последний лосс, Chamfer Loss, затрагивает координаты, когда в датасете у нас больше четырех точек.

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

Для обучения установили следующие параметры:

• AdamW с learning rate = 0.00064 и weight_decay = 0.001 и batch size = 256.

• Scheduler — CosineAnnealingWarmRestarts с T_0 = 1000 шагов и T_mult = 2.

• Вес каждого лосса:

  • Charbonnier Loss — 1.0,

  • TClip Loss — 1.0,

  • Focal Loss — 9.0,

  • Chamfer Loss — 0.03.

• Для итоговой модели — Exponential Moving Average с beta = 0.9995

• Продолжительность — 1000 эпох (но лучший чекпоинт получался примерно на 300 эпохе). На одной видеокарте RTX 4090 тренировка занимает десять дней. 

Натренировав сеть, мы можем построить пайплайн детекции. На вход даем картинку, предсказываем координаты и делаем варпинг по четырем точкам.

Улучшение пайплайна

На иллюстрации выше видно, что детекция неидеальна. Иногда отклонение составляет 5–10% при допустимых 2%. Это значение мы установили для себя эмпирическим путем: посмотрели много картинок и приняли такой порог.

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

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

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

В итоге для уточняющей сети мы переиспользовали QuadraTNet, но здесь на выходе она выдает не классификацию и четыре точки, а просто одну точку, то есть две координаты. Выше приведены картинки с разными шумами для тренировки уточняющей сети. Красный кружок — это предсказание сети, зеленый — точка из разметки. Мы брали те же датасеты, что и в прошлый раз, но делали кроп размером от 14% оригинала с центром в предсказанном уголке. Почему именно 14%? Исходили из 5% на ошибки грубой детекции с двух сторон и 4% запаса.

В итоге весь пайплайн сканера теперь состоит из пяти шагов:

1. На вход поступает картинка.

2. На превью срабатывает первая модель для грубой детекции.

3. Пользователь делает фото.

4. Вторая модель уточняет координаты точек.

5. Перспективное преобразование по четырем точкам и сохранение в галерею.

Вот как это работает на планшете. Если обе сети вдруг ошибаются, можно самому скорректировать углы документа. Это стандартная возможность в режиме документов.

Синтетический датасет

Я рассказал, почему детекция не такая точная, как хотелось бы. Но есть еще одна проблема: детекция может быть неверна в принципе. Причины могут быть разные:

• При повороте камеры точки начинают съезжать в сторону.

• Сеть может путаться при появлении теней от документа.

В датасете просто отсутствуют эти примеры, и их нельзя получить простыми преобразованиями — перспективное преобразование добавит паддинги, на которых сеть переобучится, а тени из аугментаций не похожи на реальные. Мы решили дополнить датасет недостающими кейсами, но не размечать их руками, а сгенерировать. Можно было взять 3D-модели, как в одном из датасетов, и отрендерить их. Тогда у нас будет и разметка, и фото. Другой, более доступный вариант — создать датасет с помощью модели для генерации изображений.

Мы выбрали Stable Diffusion. Она генерирует картинки, но для них также нужно было получить разметку. У Stable Diffusion есть множество модификаций, которые позволяют контролировать генерацию. Одна из них, ControlNet, позволяет вместе с текстом использовать для генерации картинки в качестве подсказки еще и изображения — сегментацию, карты границ, контуры.

С помощью ControlNet мы сгенерировали датасет с разметкой. Генерация происходила в три шага:

1. Отрисовали аннотацию в виде линий.

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

3. Отфильтровали изображения вручную: предварительно сгенерировали на 30% больше и просто удалили то, что не подходит. С возможностями нейросетей гораздо проще удалять, чем редактировать.

Мы сделали отдельный инструмент для создания аннотаций. Он берет четыре точки какого-нибудь изображения из ранее описанных наборов данных и отрисовывает линии. Сначала случайным образом задние, фоновые для документа. Затем линии границ документа по четырем точкам. И наконец, случайные линии на самом документе. Мы создали линии, промпты для достаточно консистентной генерации и, совместив это, сгенерировали изображения.

Так мы получили около 80 000 разнообразных изображений с новыми вариантами документов. Это почти половина всех изображений, на которых мы тренировали модели.

Что получилось

По итогам работы мы сравнили свою натренированную модель с базовой LDRNet по совпадению предсказанных углов документов с реальными. Наша модель может обработать больше разнообразных кейсов и умещается в 10 МБ.

Если суммировать всю проделанную работу, то мы:

• написали пайплайн тренировки, легко адаптируемый под другую архитектуру, и пайплайн конвертации и инференса на устройствах,

• добились работы превью-режима со скоростью 20 мс на кадр (50 fps),

• добавили детекцию документов в обновлении kvadraOS 1.4,

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

Вот как это выглядит на планшете. Для запуска режима документов есть отдельная кнопка. Если нажать ее, то в превью-режиме у нас появляются предварительные границы документа:

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

Когда вышло обновление kvadraOS 1.4, нам попался обзор, в котором пользователь пробовал детекцию документов на коробке. Благодаря тому, что мы предусмотрели разные альтернативные кейсы, модель успешно смогла ее определить.

Режим сканирования документов со столь же подробной визуализацией есть у у мобильных устройств Samsung. Мы сравнили их и новые Kvadra. Основное различие в том, что у нас рамки документа при движении пытаются адаптироваться на ходу и в процессе чуть съезжают, а Samsung просто прячет их и перерисовывает заново. По производительности во всех кейсах мы оказались примерно на том же уровне. Отставание было заметно только в обработке помятых документов — нам есть куда расти, и мы продолжим развивать алгоритм детекции в новых версиях kvadraOS.

Напоследок хочу поблагодарить коллег из YADRO Ивана Вихрева и Антона Клыкова за совместную работу над проектом и помощь в подготовке статьи. Если вас заинтересовала реализация AI-функций в Kvadra_T, оцените другой пост в нашем блоге — об исследовании энергоэффективности инференса нейросетей на планшете.