Views
No views yet
<s_matricula> for the V1.
1import torch
2import re
3from PIL import Image
4from transformers import DonutProcessor
5#from transformers import VisionEncoderDecoderModel
6
7import warnings
8warnings.filterwarnings("ignore")
9
10from sconf import Config
11from donut import DonutConfig, DonutModel
12
13config = Config(default="./config.yaml")
14
15device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
16processor = DonutProcessor.from_pretrained("marzanconsulting/donut-dr-matriculas-ocr")
17
18model = DonutModel.from_pretrained(
19 "marzanconsulting/donut-dr-matriculas-ocr",
20 input_size=config.input_size,
21 max_length=config.max_length,
22 align_long_axis=config.align_long_axis,
23 ignore_mismatched_sizes=True,
24 )
25
26model.to(device)
27
28def load_and_preprocess_image(image_path: str, processor):
29 """
30 Load an image and preprocess it for the model.
31 """
32 image = Image.open(image_path).convert("RGB")
33 pixel_values = processor(image, return_tensors="pt").pixel_values
34 return pixel_values
35
36def generate_text_from_image(model, image_path: str, processor, device):
37 """
38 Generate text from an image using the trained model.
39 """
40 # Load and preprocess the image
41 pixel_values = load_and_preprocess_image(image_path, processor)
42 pixel_values = pixel_values.to(device)
43
44 decoder_input_ids = processor.tokenizer(task_prompt="<s_matricula>",
45 add_special_tokens=False,
46 return_tensors="pt").input_ids
47
48 decoded_text = model.inference(image_tensors=pixel_values,
49 prompt_tensors=decoder_input_ids)["predictions"][0]
50
51 return decoded_text
52
53# Example usage
54image_path = "path_to_your_image" # Replace with your image path
55extracted_text = generate_text_from_image(model, image_path, processor, device)
56print("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}