Views
No views yet

1from transformers import DonutProcessor, VisionEncoderDecoderModel
2
3import re
4import cv2
5import json
6import torch
7from tqdm.auto import tqdm
8import numpy as np
9
10from donut import JSONParseEvaluator
11
12processor = DonutProcessor.from_pretrained("sourinkarmakar/kyc_v1-donut-demo")
13model = VisionEncoderDecoderModel.from_pretrained("sourinkarmakar/kyc_v1-donut-demo")
14
15# Need to install python-donut
16# !pip install -q donut-python
17
18# Images stored inside a folder 'unseen_samples'
19dataset = glob.glob(os.path.join(basepath, "unseen_samples/*"))
20
21output_list = []
22
23for idx, sample in tqdm(enumerate(dataset), total=len(dataset)):
24# prepare encoder inputs
25img = cv2.imread(sample)
26img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
27pixel_values = processor(img, return_tensors="pt").pixel_values
28pixel_values = pixel_values.to(device)
29
30# prepare decoder inputs
31task_prompt = "<s_cord-v2>"
32decoder_input_ids = processor.tokenizer(task_prompt, add_special_tokens=False, return_tensors="pt").input_ids
33decoder_input_ids = decoder_input_ids.to(device)
34
35# autoregressively generate sequence
36outputs = model.generate(
37pixel_values,
38decoder_input_ids=decoder_input_ids,
39max_length=model.decoder.config.max_position_embeddings,
40early_stopping=True,
41pad_token_id=processor.tokenizer.pad_token_id,
42eos_token_id=processor.tokenizer.eos_token_id,
43use_cache=True,
44num_beams=1,
45bad_words_ids=[[processor.tokenizer.unk_token_id]],
46return_dict_in_generate=True,
47)
48
49# turn into JSON
50seq = processor.batch_decode(outputs.sequences)[0]
51seq = seq.replace(processor.tokenizer.eos_token, "").replace(processor.tokenizer.pad_token, "")
52seq = re.sub(r"<.*?>", "", seq, count=1).strip() # remove first task start token
53seq = processor.token2json(seq)
54
55output_list.append(seq)
56
57print(output_list)