本記事では、LINE CLOVA OCRを開発するCLOVA AIが発表した文書理解AI Docutを用いて文書分類、情報抽出、VQAを行う方法をご紹介します。

Donut

概要

Donutは、Document Understanding Transformerは、Transformerベースの文書理解技術です。

Donutの特徴は、OCRエンジンを必要せずOCR freeで、文書分類、情報抽出、Visual Question Answering(VQA)を実現している点です。これにより、OCRの高い計算コストやOCRエラーによる後継プロセスのエラーなどの問題を解消しています。Donutでは、画像をTransformerによるエンコーダー・デコーダーを用いてそれぞれのタスクに応じたデコードを実現しています。

詳細はこちらの論文をご参照ください。

本記事では上記手法を用いて、文書画像から情報抽出

デモ(Colaboratory)

それでは、実際に動かしながらDonutによる文書分類などを行っていきます。
ソースコードは本記事にも記載していますが、下記のGitHubでも取得可能です。
GitHub - Colaboratory demo

また、下記から直接Google Colaboratoryで開くこともできます。

なお、このデモはPythonで実装しています。
Pythonの実装に不安がある方、Pythonを使った機械学習について詳しく勉強したい方は、以下の書籍やオンライン講座などがおすすめです。

環境セットアップ

それではセットアップしていきます。 Colaboratoryを開いたら下記を設定しGPUを使用するようにしてください。

初めに、ライブラリをインストールします。

!pip install donut-python

最後にライブラリをインポートします。

import torch
from PIL import Image

from donut import DonutModel

import IPython
import pprint

以上で環境セットアップは完了です。

テスト画像のセットアップ

続いて、モデルに入力するテスト画像をセットアップします。

!rm -rf test_01.jpg

# 英語のスライド
!wget -c https://media.slidesgo.com/storage/162635/conversions/1-market-share-infographics-thumb.jpg \
      -O ./test_01.jpg

input_img_path = './test_01.jpg'

本記事では、以下の画像を使用します。

Document Classification

それでは、初めに文書分類を行います。
まず学習済みモデルのロードをロードします。

pretrained_model = DonutModel.from_pretrained("naver-clova-ix/donut-base-finetuned-rvlcdip")

if torch.cuda.is_available():
    pretrained_model.half()
    device = torch.device("cuda")
    pretrained_model.to(device)
else:
    pretrained_model.encoder.to(torch.bfloat16)

pretrained_model.eval()

ロードしたモデルに画像を入力します。

task_prompt = "<s_rvlcdip>"

input_img = Image.open(input_img_path)
output = pretrained_model.inference(image=input_img, prompt=task_prompt)["predictions"][0]

出力結果は以下の通りです。RVL-CLIPで学習している本モデルですが、問題なくpresentationに分類されています。

pprint.pprint(output)
# {'class': 'presentation'}

Document Information Extraction

次に、文書から情報抽出を行います。
実行方法を先ほどと同じで、モデルとpromptが異なるのみです。

pretrained_model = DonutModel.from_pretrained("naver-clova-ix/donut-base-finetuned-cord-v2")

if torch.cuda.is_available():
    pretrained_model.half()
    device = torch.device("cuda")
    pretrained_model.to(device)
else:
    pretrained_model.encoder.to(torch.bfloat16)

pretrained_model.eval()

task_prompt = "<s_cord-v2>"

input_img = Image.open(input_img_path)
output = pretrained_model.inference(image=input_img, prompt=task_prompt)["predictions"][0]

出力結果は以下の通りです。
こちらは、CORDと呼ばれるレシートのデータセットでトレーニングしているためスライドの情報抽出としては精度がいまいちです。

pprint.pprint(output)
{'menu': {'cnt': '1',
          'nm': 'Market Share Infographics',
          'price': {'cnt': '1', 'unitprice': 'a gas giant and'},
          'unitprice': 'Rupiter'},
 'sub_total': {'cnt': '1', 'price': 'Mercury', 'unitprice': '30%'}}

Document Visual Question Answering

次に、文書画像と質問を与えてモデルに回答してもらうVQAを行います。
実行方法を先ほどと同じで、モデルとpromptが異なるのみです。

pretrained_model = DonutModel.from_pretrained("naver-clova-ix/donut-base-finetuned-docvqa")

if torch.cuda.is_available():
    pretrained_model.half()
    device = torch.device("cuda")
    pretrained_model.to(device)
else:
    pretrained_model.encoder.to(torch.bfloat16)

pretrained_model.eval()

task_name = "docvqa"
task_prompt = "<s_docvqa><s_question>{user_input}</s_question><s_answer>"

question = "What has the highest market share?" # @param {type:"string"}
user_prompt = task_prompt.replace("{user_input}", question)

input_img = Image.open(input_img_path)
output = pretrained_model.inference(input_img, prompt=user_prompt)["predictions"][0]

出力結果は以下の通りです。
35%を見落としていますね。

pprint.pprint(output)
{'answer': '30%', 'question': 'What has the highest market share?'}

まとめ

本記事では、本記事では、LINE CLOVA OCRを開発するCLOVA AIが発表した文書理解AI Docutを用いて文書分類、情報抽出、VQAを行う方法をご紹介しました。
公開済みモデルの精度はこれからといった印象ですが、アーキテクチャ設計自体が新しいのでこれからの発展に注目です。

また本記事では、機械学習を動かすことにフォーカスしてご紹介しました。
もう少し学術的に体系立てて学びたいという方には以下の書籍などがお勧めです。ぜひご一読下さい。

また動かせるだけから理解して応用できるエンジニアの足掛かりに下記のUdemyなどもお勧めです。

参考文献

1.  論文 - OCR-free Document Understanding Transformer

2. GitHub - clovaai/donut