OCR-VQA with Donut

3 min readApr 9, 2024

Introduction

Imagine being able to ask a question about a photograph and having the answer revealed not just by the image itself, but by the text within it! This is the fascinating world of Optical Character Recognition (OCR) and Visual Question Answering (VQA).

OCR technology bridges the gap between the physical and the digital, allowing computers to ‘read’ text in images. VQA goes one step further. By combining OCR’s ability to extract text with advanced image analysis, VQA systems can answer questions that rely on understanding both the visual content and the embedded text.

Enter the Document Understanding Transformer (Donut) model, a powerful tool that harnesses the potential of OCR and VQA. In the following sections, we’ll take a closer look at how Donut works and how you can use it to unlock the hidden insights in your images and text.

Get Mert’s stories in your inbox

Join Medium for free to get updates from this writer.

Remember me for faster sign in

While the inner workings of Donut are fascinating, this blog focuses on its practical application. We’ll bypass the deep dives into theory and instead equip you with the knowledge to use Donut in your own projects. We’ll walk you through the steps of getting started, from setting up the model to extracting valuable information from your images and text. For those curious about the inner workings of the donut model, there are some links at the end. To mention it, you can also use Donut for document image classification and form understanding.

How to use Donut for OCR-VQA

Step 1: Installing the necessary libraries

pip install transformers torch opencv-python

Step 2: Importing libraries

from transformers import DonutProcessor, VisionEncoderDecoderModel
import cv2
import torch
import re

Step 3: Loading Model and Processor

processor = DonutProcessor.from_pretrained("naver-clova-ix/donut-base-finetuned-docvqa")
model = VisionEncoderDecoderModel.from_pretrained("naver-clova-ix/donut-base-finetuned-docvqa")
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)

If you have a GPU, you should use it, as you will get much faster results.

Step 4: Loading your Image

image_path = r"YourImagePath"
image = cv2.imread(image_path)

Step 5: Give your prompt

# prepare decoder inputs
task_prompt = "<s_docvqa><s_question>{user_input}</s_question><s_answer>"
question = "What's the population of Dade?"
prompt = task_prompt.replace("{user_input}", question)
decoder_input_ids = processor.tokenizer(prompt, add_special_tokens=False, return_tensors="pt").input_ids

Step 6: Prepare Input of Decoder

decoder_input_ids = processor.tokenizer(prompt, add_special_tokens=False, return_tensors="pt").input_ids

pixel_values = processor(image, return_tensors="pt").pixel_values

outputs = model.generate(
    pixel_values.to(device),
    decoder_input_ids=decoder_input_ids.to(device),
    max_length=model.decoder.config.max_position_embeddings,
    pad_token_id=processor.tokenizer.pad_token_id,
    eos_token_id=processor.tokenizer.eos_token_id,
    use_cache=True,
    bad_words_ids=[[processor.tokenizer.unk_token_id]],
    return_dict_in_generate=True,
)

Step 7: Get the result

sequence = processor.batch_decode(outputs.sequences)[0]
sequence = sequence.replace(processor.tokenizer.eos_token, "").replace(processor.tokenizer.pad_token, "")
sequence = re.sub(r"<.*?>", "", sequence, count=1).strip()  # remove first task start token
print(processor.token2json(sequence))

Conclusion

In this tutorial we have learned how to use Donut for OCR-VQA. If you found this code helpful, please clap your hands and comment on this post! I would also love for you to follow me to learn more about Data Science and other related topics. Thanks for reading!