Понимание документов без OCR с помощью Donut

Используйте недавно выпущенную модель Transformers для создания JSON-представлений данных вашего документа.

Визуальное понимание документов (VDU) — это активно изучаемая новая область глубокого обучения и науки о данных, особенно потому, что в PDF-файлах или отсканированных документах содержится множество неструктурированных данных. Последние модели, такие как LayoutLM, используют архитектуру модели глубокого обучения трансформеры для маркировки слов или ответов на заданные вопросы на основе изображения документа (например, вы можете выделить и пометить номер счета, аннотировав само изображение). , или спросите модель: Какой номер счета?). Такие библиотеки, как transformers от HuggingFace, упрощают работу с моделями трансформеров с открытым исходным кодом.

Большинство традиционных ответов на проблему VDU полагаются на синтаксический анализ выходных данных OCR этого изображения вместе с визуальным кодированием, но OCR требует больших вычислительных ресурсов (поскольку обычно требует установки механизма OCR, такого как Tesseract) и включения еще одной модели в полный конвейер приводит к другой модели, которую необходимо обучить и настроить, а неточная модель OCR приведет к распространению ошибок в модели VDU.

Так, исследователи из Naver CLOVA предложили сквозное решение VDU [1], использующее архитектуру модели преобразования кодер-декодер, и недавно сделали его доступным для использования с библиотекой HuggingFace transformers. Другими словами, он кодирует изображение (разделенное на фрагменты с помощью Swin Transformer) в векторы-токены, которые затем можно декодировать или преобразовать в выходную последовательность в виде структура данных (которая затем может быть преобразована в JSON) с использованием модели декодера BART, публично предварительно обученной на многоязычных наборах данных. Любые подсказки, подаваемые в модель во время вывода, также могут быть декодированы в той же архитектуре.

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

Я получил результат:

{
    nm: "Presentation"
}

что указывает на то, что он обнаружил, что заголовок «Презентация» является названием пункта меню или чека.

Авторы также предоставили сценарии обучения и тестирования, чтобы мы могли продемонстрировать, как на самом деле использовать модели на практике (я буду использовать набор данных SROIE [2], набор данных помеченных квитанций и счетов, чтобы продемонстрировать тонкую настройку). в пользовательском наборе данных). Я бы рекомендовал запускать код на графическом процессоре, так как вывод и обучение займут много времени на процессоре. Google Colab предлагает бесплатный доступ к графическому процессору, и его должно быть достаточно для точной настройки (перейдите в раздел Среда выполнения › Изменить тип среды выполнения, чтобы переключиться с ЦП на ГП).

Во-первых, давайте удостоверимся, что у нас есть доступ к графическому процессору.

import torch
print("CUDA available:", torch.cuda.is_available())
!nvcc --version

И теперь мы можем загрузить соответствующие файлы и библиотеки. Следующие строки кода должны установить все зависимости, включая библиотеку пончиков (хотя вы можете установить ее вручную с помощью pip install donut-python, кодовая база, клонированная из Github, включает важные сценарии обучения и тестирования).

!git clone https://github.com/clovaai/donut.git
!cd donut && pip install .

Вывод с использованием точной модели CORD

Во-первых, мы продемонстрируем базовое использование модели.

from donut import DonutModel
from PIL import Image
import torch
model = DonutModel.from_pretrained("naver-clova-ix/donut-base-finetuned-cord-v2")
if torch.cuda.is_available():
    model.half() 
    device = torch.device("cuda") 
    model.to(device) 
else: 
    model.encoder.to(torch.bfloat16)
model.eval() 
image = Image.open("./donut/misc/sample_image_cord_test_receipt_00004.png")
    .convert("RGB")
output = model.inference(image=image, prompt="<s_cord-v2>")
output

В вызове DonutModel.from_pretrained() я просто указал имя предварительно обученной модели из HuggingFace Hub (необходимые файлы загружаются в это время), хотя я также могу указать локальный путь к папке модели, как мы покажем позже. . Кодовая база Donut также включает образец изображения (показан ниже), который я передал в модель, но вы можете протестировать модель с любым изображением, которое вам нравится.

Вы должны получить вывод, например

{'predictions': [{'menu': [{'cnt': '2', 'nm': 'ICE BLAOKCOFFE', 'price': '82,000'}, 
    {'cnt': '1', 'nm': 'AVOCADO COFFEE', 'price': '61,000'}, 
    {'cnt': '1', 'nm': 'Oud CHINEN KATSU FF', 'price': '51,000'}],
    'sub_total': {'discount_price': '19,400', 'subtotal_price': '194,000'}, 
    'total': {'cashprice': '200,000', 
    'changeprice': '25,400', 
    'total_price': '174,600'}}]}

(Примечание. Если вам, как и мне, любопытно, и вы хотите знать, какие выходные данные даст вам предварительно обученная магистраль donut-base, я пошел дальше и протестировал ее. Потребовалось много времени, чтобы получить выходные данные, прежде чем произошел сбой, потому что он занимал слишком много оперативной памяти. )

Тонкая настройка Donut на пользовательском наборе данных

Чтобы продемонстрировать тонкую настройку, я буду использовать набор данных SROIE, набор данных сканирования квитанций и счетов вместе с их основной информацией в форме JSON, а также ограничивающими рамками и текстом на уровне слов. Он содержит 626 изображений, но я буду тренироваться только на 100, чтобы продемонстрировать эффективность Donut. Это меньший набор данных, чем CORD (который содержит около 1000 изображений), а также гораздо меньше меток (только компания, дата, адрес и общая сумма).

Скачивание и парсинг SROIE

Чтобы загрузить набор данных, вам нужно скачать только папку data из основного репозитория. Вы можете сделать это либо путем клонирования всего репозитория, либо с помощью чего-то вроде Каталога загрузки, чтобы загрузить только одну папку.

Но теперь нам нужно проанализировать набор данных в формате, требуемом библиотекой HuggingFace datasets, которую Donut использует под капотом для загрузки пользовательского набора данных в виде таблицы строк изображений. (Если вам нужна документация, Donut использует скрипт загрузки imagefolder.)

Вот желаемый формат набора данных:

dataset_name
├── test
│   ├── metadata.jsonl
│   ├── {image_path0}
│   ├── {image_path1}
│             .
│             .
├── train
│   ├── metadata.jsonl
│   ├── {image_path0}
│   ├── {image_path1}
│             .
│             .
└── validation
    ├── metadata.jsonl
    ├── {image_path0}
    ├── {image_path1}
              .
              .

Где metadata.jsonl – это документ строк JSON, который выглядит как

{"file_name": {image_path0}, "ground_truth": "{\"gt_parse\": {ground_truth_parse}, ... {other_metadata_not_used} ... }"}
{"file_name": {image_path1}, "ground_truth": "{\"gt_parse\": {ground_truth_parse}, ... {other_metadata_not_used} ... }"}

Другими словами, мы хотим преобразовать аннотации каждого документа (находящиеся в папке key) в строку достоверности JSON, которая выглядит как "{\"gt_parse\": {actual JSON content}"}". Вот пример аннотации:

{
    "company": "BOOK TA .K (TAMAN DAYA) SDN BHD",
    "date": "25/12/2018",
    "address": "NO.53 55,57 & 59, JALAN SAGU 18, TAMAN DAYA, 81100 JOHOR BAHRU, JOHOR.",
    "total": "9.00"
}

Вот скрипт, который я использовал для преобразования данных в файлы строк JSON, а также для копирования изображений в соответствующие папки:

import os
import json
import shutil
from tqdm.notebook import tqdm
lines = []
images = []
for ann in tqdm(os.listdir("./sroie/key")[:100]):
  if ann != ".ipynb_checkpoints":
    with open("./sroie/key/" + ann) as f:
      data = json.load(f)
images.append(ann[:-4] + "jpg")
    line = {"gt_parse": data}
    lines.append(line)
with open("./sroie-donut/train/metadata.jsonl", 'w') as f:
  for i, gt_parse in enumerate(lines):
    line = {"file_name": images[i], "ground_truth": json.dumps(gt_parse)}
    f.write(json.dumps(line) + "\n")
shutil.copyfile("./sroie/img/" + images[i], "./sroie-donut/train/" + images[i])

Я просто запускал этот скрипт три раза, каждый раз меняя имена папок и срез списка ([:100]), так что у меня было 100 примеров в train и по 20 примеров в validation и проверить.

dataset_name_or_paths: ["../sroie-donut"]
train_batch_sizes: [1]
check_val_every_n_epochs: 10
max_steps: -1 # infinite, since max_epochs is specified

cd donut && python train.py --config config/train_sroie.yaml

from donut import DonutModel
from PIL import Image
import torch
model = DonutModel.from_pretrained("./donut/result/train_sroie/20220804_214401")
if torch.cuda.is_available():
    model.half()
    device = torch.device("cuda")
    model.to(device)
else:
    model.encoder.to(torch.bfloat16)
model.eval()
image = Image.open("./sroie-donut/test/099.jpg").convert("RGB")
output = model.inference(image=image, prompt="<s_sroie-donut>")
output

{'predictions': [{'address': 'NO 290, JALAN AIR PANAS. SETAPAK. 53200, KUALA LUMPUR.',
   'company': 'SYARIKAT PERNIAGAAN GIN KEE',
   'date': '04/12/2017',
   'total': '47.70'}]}

*' il " i
- ' s ' -
W =
o o
ok S
?flfi (€
rgm"f"; o ;
L i 4

cd ./donut &&
python test.py --dataset_name_or_path ../sroie-donut --pretrained_model_name_or_path ./result/train_sroie/20220804_214401 --save_path ./result/train_sroie/output.json