Views
No views yet
1from transformers import DonutProcessor, VisionEncoderDecoderModel
2from PIL import Image
3import torch
4import re
5
6model_name = 'hsarfraz/irs-tax-form-1040-2023-doc-parser'
7
8processor = DonutProcessor.from_pretrained(model_name)
9model = VisionEncoderDecoderModel.from_pretrained(model_name)
10
11device = "cuda" if torch.cuda.is_available() else "cpu"
12model.to(device)
13model.eval()
14
15image_name = 'replace with name of the form 1040 (2023) image file '
16
17img = Image.open(image_name)
18new_width = 1536
19new_height = 1536
20
21# resize input image to finetuned images size
22img = img.resize((new_width, new_height), Image.LANCZOS)
23
24pixel_values = processor(img.convert("RGB"), return_tensors="pt").pixel_values
25pixel_values = pixel_values.to(device)
26
27# prompt
28task_prompt = "<s_cord-v2>"
29decoder_input_ids = processor.tokenizer(task_prompt, add_special_tokens=False, return_tensors="pt")["input_ids"]
30decoder_input_ids = decoder_input_ids.to(device)
31
32outputs = model.generate(pixel_values,decoder_input_ids=decoder_input_ids,
33 max_length=model.decoder.config.max_position_embeddings,
34 early_stopping=True,
35 pad_token_id=processor.tokenizer.pad_token_id,
36 eos_token_id=processor.tokenizer.eos_token_id,
37 use_cache=True,
38 num_beams=1,
39 bad_words_ids=[[processor.tokenizer.unk_token_id]],
40 return_dict_in_generate=True,
41 # output_scores=True,
42 )
43
44
45sequence = processor.batch_decode(outputs.sequences)[0]
46sequence = sequence.replace(processor.tokenizer.eos_token, "").replace(processor.tokenizer.pad_token, "")
47sequence = re.sub(r"<.*?>", "", sequence, count=1).strip() # remove first task start token
48output_json = processor.token2json(sequence)
49
50print('----------------------------------')
51print('--- Parsed data in json format ---')
52print('----------------------------------')
53print(output_json)
54
1{
2 "lbl_0_03": "Michael Evans",
3 "lbl_0_04": "Caldwell",
4 "lbl_0_05": "741-52-5353",
5 "lbl_0_06": "None",
6 "lbl_0_07": "None",
7 "lbl_0_08": "None",
8 "lbl_0_09": "289 Blackwell Land Suite 380 New Tiffany, NH 07548",
9 "lbl_0_11": "East Amandaport",
10 "lbl_0_12": "VI",
11 "lbl_0_13": "47832",
12 "lbl_0_14": "None",
13 "lbl_0_15": "None",
14 "lbl_0_16": "25677",
15 "lbl_0_55": "385321.36",
16 "lbl_0_56": "None",
17 "lbl_0_57": "None",
18 "lbl_0_58": "None",
19 "lbl_0_59": "None",
20 "lbl_0_60": "None",
21 "lbl_0_61": "None",
22 "lbl_0_62": "None",
23 "lbl_0_63": "None",
24 "lbl_0_67": "None",
25 "lbl_0_68": "481161.23",
26 "lbl_0_69": "None",
27 "lbl_0_70": "None",
28 "lbl_0_71": "None",
29 "lbl_0_72": "749100.68",
30 "lbl_0_73": "418381-6",
31 "lbl_0_74": "None",
32 "lbl_0_77": "755042.64",
33 "lbl_0_78": "None",
34 "lbl_0_79": "560928.32",
35 "lbl_0_80": "493913.73",
36 "lbl_0_81": "None",
37 "lbl_0_82": "738597.72",
38 "lbl_0_83": "34990.46"
39}