Views
No views yet
<s_receipt> for the V2 (previously <s_cord-v2> for V1). Two new keys <s_svc> and <s_discount> have been added, <s_telephone> has been renamed to <s_phone>.
1import torch
2import re
3from PIL import Image
4from transformers import DonutProcessor, VisionEncoderDecoderModel
5
6device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
7processor = DonutProcessor.from_pretrained("AdamCodd/donut-receipts-extract")
8model = VisionEncoderDecoderModel.from_pretrained("AdamCodd/donut-receipts-extract")
9model.to(device)
10
11def load_and_preprocess_image(image_path: str, processor):
12 """
13 Load an image and preprocess it for the model.
14 """
15 image = Image.open(image_path).convert("RGB")
16 pixel_values = processor(image, return_tensors="pt").pixel_values
17 return pixel_values
18
19def generate_text_from_image(model, image_path: str, processor, device):
20 """
21 Generate text from an image using the trained model.
22 """
23 # Load and preprocess the image
24 pixel_values = load_and_preprocess_image(image_path, processor)
25 pixel_values = pixel_values.to(device)
26
27 # Generate output using model
28 model.eval()
29 with torch.no_grad():
30 task_prompt = "<s_receipt>" # <s_cord-v2> for v1
31 decoder_input_ids = processor.tokenizer(task_prompt, add_special_tokens=False, return_tensors="pt").input_ids
32 decoder_input_ids = decoder_input_ids.to(device)
33 generated_outputs = model.generate(
34 pixel_values,
35 decoder_input_ids=decoder_input_ids,
36 max_length=model.decoder.config.max_position_embeddings,
37 pad_token_id=processor.tokenizer.pad_token_id,
38 eos_token_id=processor.tokenizer.eos_token_id,
39 early_stopping=True,
40 bad_words_ids=[[processor.tokenizer.unk_token_id]],
41 return_dict_in_generate=True
42 )
43
44 # Decode generated output
45 decoded_text = processor.batch_decode(generated_outputs.sequences)[0]
46 decoded_text = decoded_text.replace(processor.tokenizer.eos_token, "").replace(processor.tokenizer.pad_token, "")
47 decoded_text = re.sub(r"<.*?>", "", decoded_text, count=1).strip() # remove first task start token
48 decoded_text = processor.token2json(decoded_text)
49 return decoded_text
50
51# Example usage
52image_path = "path_to_your_image" # Replace with your image path
53extracted_text = generate_text_from_image(model, image_path, processor, device)
54print("Extracted Text:", extracted_text)1@article{DBLP:journals/corr/abs-2111-15664,
2 author = {Geewook Kim and
3 Teakgyu Hong and
4 Moonbin Yim and
5 Jinyoung Park and
6 Jinyeong Yim and
7 Wonseok Hwang and
8 Sangdoo Yun and
9 Dongyoon Han and
10 Seunghyun Park},
11 title = {Donut: Document Understanding Transformer without {OCR}},
12 journal = {CoRR},
13 volume = {abs/2111.15664},
14 year = {2021},
15 url = {https://arxiv.org/abs/2111.15664},
16 eprinttype = {arXiv},
17 eprint = {2111.15664},
18 timestamp = {Thu, 02 Dec 2021 10:50:44 +0100},
19 biburl = {https://dblp.org/rec/journals/corr/abs-2111-15664.bib},
20 bibsource = {dblp computer science bibliography, https://dblp.org}
21}