執筆者：Handbook編集部
画像内の構造を理解してQAを行うpix2structの紹介
pix2stract の概要
pix2stract とは、入力された画像の構造や意味を把握し、質問文に対して適切な出力をするモデルです。
今まで、画像認識というと画像中に含まれているオブジェクトのみの認識で文字に対するアプローチはできませんでした。そのため、画像中の文字を認識するためには別の OCR モデルを使用する必要があり、モデルの複雑化や計算コストがかかるといった問題がありました。
pix2stract では、入力画像を Web ページと捉え、html の構造を出力するように学習することで入力画像の構造を把握することを可能としています。
この記事では、この pix2stract のモデル構造、学習工程について紹介します。
pix2stract のモデル
モデルには、Visual Transformer (ViT)が使用されています。一般的な ViT と異なる点としては入力画像を固定の大きさにする必要がないことです。
pix2stract は入力画像全体の構造を把握する必要があるため、画像サイズを固定すると歪みや構造が崩れてしまう可能性があります。そのため、入力画像を拡大・縮小をして入力画像の性質を保ったまま入力することで構造を把握しやすくしています。また、この時に可変の解像度に対応するために二次元の位置情報を埋め込んでいます。
pix2stract の学習方法
ウォームアップ
入力画像から html を出力するタスクをいきなり始めると学習が不安定で遅くなるそうです。そのため、まずは画像のようなランダムな色とフォントの文字画像から元の文字を出力する単純なタスクをウォームアップとして導入しています。
このウォームアップによって学習が安定し速く収束するほか、fine tuning の際に精度が向上する効果がありました。
このウォームアップ手法は、Dessurt モデルでも使用されており、SOTA を達成している Donut の事前学習の簡略化バージョンと言える学習方法になります。
事前学習
pix2stract は fine tuning を行うことで様々なタスクに流用できますが、その前に事前学習と呼ばれる html 構造を理解するための学習を行います。
ここで使用された学習データは、80M の Web スクリーンショットと html のペアで自己教師あり学習を行います。自己教師あり学習なのですが、html をそのまま使用するのではなく、可視要素を持つノードに限定されています。論文中では可視要素のみでしたが、スタイルやタイトル、URL なども情報として保持することも可能なようです。
学習では、BERT と同じように画像にランダムな 50%のマスク処理を行います。
fine tuning
pix2stract では画像のような複数のタスクに対する fine tuning が行われます。この時に、入力画像と入力テキストは別々で処理するのではなく入力画像のヘッダーとして入力テキストを埋め込んだ画像として入力します。選択肢のある入力テキストに対してもヘッダーに埋め込み、あくまでも画像としてモデルに入力をします。
pix2stract の性能
pix2stract は、画像から文字、図についての QA タスクで以前まで最先端のモデルであった OCR を使用しない Donut モデルの精度を上回りました。また、ドキュメントから文字の抽出には既存の OCR 手法には及ばないものの匹敵するくらいの精度をしていました。
pix2stract の Demo
pix2struct の DocVQA タスクは簡単に試すことができます。
(Demo: Pix2Struct for DocVQA)[https://huggingface.co/spaces/RamAnanth1/Pix2Struct]ここで読み取りたいドキュメント画像をアップロードし質問文を入力することで簡単にレスポンスを得ることができます。
また、pdf から QA をするタスクではコードによる実装もすることができます。このコードでは、画像化したドキュメントも使用することができます。hugging face の transformers で用意されているので簡単に実装することができます。
1. パッケージのインストール
pip install git+https://github.com/huggingface/transformers pdf2image
apt install poppler-utils12diff
2. 必要パッケージのインポート
from PIL import Image
from pdf2image import convert_from_path, convert_from_bytes
import torch
from functools import partial
from transformers import Pix2StructForConditionalGeneration as psg
from transformers import Pix2StructProcessor as psp
3. モデルの読み込み
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
model = psg.from_pretrained("google/pix2struct-docvqa-large").to(DEVICE)
processor = psp.from_pretrained("google/pix2struct-docvqa-large")
4. QA を行うための関数
def generate(model, processor, img, questions):
inputs = processor(images=[img for _ in range(len(questions))],
text=questions, return_tensors="pt").to(DEVICE)
predictions = model.generate(**inputs, max_new_tokens=256)
return zip(questions, processor.batch_decode(predictions, skip_special_tokens=True))
def convert_pdf_to_image(filename, page_no):
return convert_from_path(filename)[page_no-1]
4. ドキュメントから QA を行う
doc_file = "ドキュメントのパス"
page_no = "QAを行うページ番号"
questions = ["質問文(複数可)"]
# 画像を直接使用する場合はパスできます
image = convert_pdf_to_image(doc_file, page_no)
generator = partial(generate, model, processor)
completions = generator(image, questions)
# 結果を表示します
for completion in completions:
print(f"{completion}")
ここでは、DocVQA の簡単な実装について紹介しましたが、より様々なタスクや設定をしたい場合は(公式 github)[https://github.com/google-research/pix2struct]を確認してみてください。
まとめ
今回は、画像から構造を把握する pix2stract について紹介しました。論文中では、各タスクについて SOTA を達成していたことから今後さらに主流になりそうなモデルになります。今後、CLIP から unCLIP のようにテキストから Web ページを作成するモデルなどが出ることを個人的に期待しています。
参考
Hakky ではエンジニアを募集中です！まずは話してみたいなどでも構いませんので、ぜひお気軽に採用ページからお問い合わせくださいませ。
