Передаточное обучение YOLOv4 для распознавания структуры отсканированного документа

Когда вы впервые слышите о «YOLO», вы, вероятно, интуитивно думаете о «You Only Live Once». Но это также название системы обнаружения объектов под названием «Вы только посмотрите один раз».

Как показано выше, YOLO может обнаруживать различные объекты в изображении и прогнозировать метку для каждого из них за очень короткое время (1-2 секунды для изображения с высоким разрешением на GPU P2). Это, несомненно, популярно в области обнаружения объектов в реальном времени. Однако он редко применяется в области определения структуры документа или таблицы. Пока я нашел только одну статью по этой теме. Учитывая многообещающий результат статьи и возможности обобщения YOLO, я решил реализовать трансферное обучение YOLOv4 для этой задачи.

В этой статье дается краткое введение в YOLO и основное внимание уделяется реализации обучения передачи YOLOv4 для обнаружения структуры отсканированного документа с использованием набора данных Marmot. Имейте в виду, что любое коммерческое использование набора данных должно быть разрешено основателем Apabi. Текущая статья предназначена только для исследовательских целей.

Краткое введение в YOLO

Последняя версия YOLOv4, разработанная Алексеем Бочковским, Чиен-Яо Ван и Хун-Юань Марком Ляо, была выпущена в апреле 2020 года. Как показано на рисунке ниже, она достигла высочайшего уровня производительности по сравнению с другими моделями при время, на 12% быстрее и на 10% точнее, чем YOLOv3.

На высоком уровне он делит входное изображение на сетки S x S. Если центр объекта попадает в ячейку сетки, эта ячейка будет отвечать за его классификацию. Каждая ячейка сетки предсказывает фиксированное количество ограничивающих рамок. В конце концов, для обучения и прогнозирования используется только поле с наивысшим показателем достоверности.

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

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

Внедрение индивидуального обучения YOLOv4

Подготовка набора данных

Набор данных Marmot состоит из страниц на китайском и английском языках с 15 различными метками. Я использовал только набор данных на английском языке, который содержит 1006 изображений и пять меток: Table, TableBody, TableFootnote, TableCaption, Paragraph. Данные аннотации представлены в формате XML.

Я разработал код для просмотра каждого XML-файла и нахождения информации об ограничивающей рамке для целевых меток. Сложная часть состоит в том, чтобы преобразовать исходную информацию о ограничивающем прямоугольнике, которая находится в шестнадцатеричном формате, в нормализованные значения от 0 до 1, необходимые для модели YOLO. Чтобы оценить свое преобразование, я использовал Yolo_mark, чтобы аннотировать одно и то же изображение и сравнить значения координат. Если у вас возникнут проблемы с OpenCV при компиляции Yolo_mark, вот пост, который мне помог.

Для каждого позитивного изображения вам необходимо создать файл .txt в том же каталоге с тем же именем. Для негативных изображений вы можете включить их в тот же каталог без файла .txt. Каждая строка в файле содержит идентификатор класса объекта и координаты этого объекта.

<object-class> <x_center> <y_center> <width> <height>

Класс объекта варьируется от 0 до общего количества классов минус 1. В моем случае это от 0 до 4 для пяти ярлыков.

Четыре координаты представляют собой значения относительные (от 0 до 1) к общей ширине и высоте изображения. Обратите особое внимание на ‹x_center› ‹y_center›, поскольку они являются центром прямоугольника, а не верхней левой точкой. Ниже приведен пример моего файла train.txt.

2 0.274538413 0.22350818 0.346147266 0.0254018459
1 0.275281816 0.45046899 0.365293325 0.3980142464
...

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

Как обучить вашу индивидуальную модель YOLOv4

Перед тем как начать, я настоятельно рекомендую вам ознакомиться с учебником и требованиями на официальной странице YOLO GitHub. Следующие шаги выполняются в терминале в Linux.

Git клонирует репо.

git clone https://github.com/AlexeyAB/darknet.git

2. Перейдите в каталог darknet и измените параметры в Makefile. Установите для GPU, CUDNN и OPENCV значение 1, чтобы ускорить процесс обучения, если у вас есть GPU. Убедитесь, что ваша система имеет правильные версии, необходимые для YOLOv4, а пути к файлам CUDA и CUDNN верны. Затем введите следующую команду для компиляции модели.

cd darknet
make

3. Создайте obj.data, obj.name, train.txt, test.txt и . yolo-obj.cfg следуя инструкциям здесь. Отрегулируйте параметры в соответствии с вашими потребностями. В моем случае я выбрал для своего обучения 10 000 эпох. Yolo-obj.cfg - это абстрактный каркас модели YOLO.

Obj.name содержит метки классов, и последовательность имеет значение для результата прогнозирования.

Table
TableBody
TableCaption
TableFootnote
Paragraph

Obj.data включает путь к разным файлам и папкам.

classes= 5
train = ./build/darknet/x64/data/train.txt
valid = ./build/darknet/x64/data/test.txt
names = ./build/darknet/x64/data/obj.names
backup = ./build/darknet/x64/backup

Train.txt и test.txt указывают расположение изображений.

./build/darknet/x64/data/obj/10.1.1.6.2197_3.jpg
./build/darknet/x64/data/obj/10.1.1.36.2833_22.jpg
...

4. Затем загрузите yolov4.weights, а затем вы можете начать обучение модели с помощью следующей командной строки. Измените пути к файлам в зависимости от вашего текущего каталога и каталога файлов.

./darknet detector train ./build/darknet/x64/data/obj.data ./build/darknet/x64/yolo-obj.cfg ./build/darknet/x64/yolov4.conv.137

Для каждых 100 итераций модель будет сохранять веса как yolo_obj_last.weight в папку ./build/darknet/x64/backup. Для каждых 1000 итераций модель будет сохранять веса как yolo-obj-xxxx.weights в той же папке.

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

./darknet detector train ./build/darknet/x64/data/obj.data ./build/darknet/x64/yolo-obj.cfg ./build/darknet/x64/backup/yolo-obj_last.weights

5. Остановите обучение, когда увидите, что средний убыток 0.xxx avg больше не уменьшается на многих итерациях.

Оценка модели

Вот мои результаты за 9000 эпох. Как видите, модель начинает переобучаться после 6000 эпох.

MAP - это аббревиатура от Mean Average Precision (Средняя средняя точность).

Давайте разберем производительность по меткам классов в наборе данных проверки. Table, TableBody, Paragraph обычно имеют лучшую точность, чем TableCaption и TableFootnote. Возможные причины могут заключаться в том, что последние два являются меньшими объектами и имеют меньше обучающих данных, чем остальные. Например, в обучающих данных всего 135 TableFootnote.

Обнаружение модели

Выбирайте веса в зависимости от вашего варианта использования. Здесь я выбрал веса 6000 эпох, так как он дает наивысший балл в наборе данных проверки. Чтобы сделать прогноз для одного изображения, выполните следующую команду.

./darknet detector test ./build/darknet/x64/data/obj.data ./build/darknet/x64/yolo-obj.cfg ./build/darknet/x64/backup/yolo-obj_6000.weights

Чтобы обработать несколько изображений и сохранить результаты в файл .txt, вы можете использовать следующее. Здесь ./custom_pred/custom_train.txt - это файл, содержащий пути к тестовым изображениям, а ./custom_pred/custom_result.txt - выходной файл результата прогноза.

./darknet detector test ./build/darknet/x64/data/obj.data ./build/darknet/x64/yolo-obj.cfg ./build/darknet/x64/backup/yolo-obj_6000.weights -dont_show -ext_output  < ./custom_pred/custom_train.txt > ./custom_pred/custom_result.txt

Вот визуализация результата прогноза на тестовом изображении.

Следующие шаги:

Вот несколько идей по улучшению характеристик модели:

Установить флаг random=1 в .cfg-файле
Для обучения и прогнозирования увеличьте разрешение сети в файле.cfg (height=608, width=608 или любое значение, кратное 32)
Самое главное, используйте Yolo_mark, чтобы добавить больше обучающих данных для классов с несколькими примерами, такими как TableFootnote.
Используйте естественные изображения в качестве негативных данных для обучения.