| import re |
| import transformers |
| from PIL import Image |
| from transformers import DonutProcessor, VisionEncoderDecoderModel |
| import torch |
| import random |
| import numpy as np |
| import gradio as gr |
|
|
| access_token = "" |
|
|
| transformers.logging.disable_default_handler() |
| processor = DonutProcessor.from_pretrained("daquarti/donut-base-sroie", use_auth_token=access_token) |
| model = VisionEncoderDecoderModel.from_pretrained("daquarti/donut-base-sroie", use_auth_token=access_token) |
| device = "cuda" if torch.cuda.is_available() else "cpu" |
| model.to(device) |
|
|
| def load_image (f): |
| with Image.open(f) as img: |
| a = img.load() |
| return img.convert('RGB') |
|
|
| def pred (a): |
| |
| |
| pixel_values = processor(a, return_tensors="pt").pixel_values |
| task_prompt = "<s>" |
| decoder_input_ids = processor.tokenizer(task_prompt, add_special_tokens=False, return_tensors="pt").input_ids |
|
|
| outputs = model.generate( |
| pixel_values.to(device), |
| decoder_input_ids=decoder_input_ids.to(device), |
| max_length=model.decoder.config.max_position_embeddings, |
| early_stopping=True, |
| pad_token_id=processor.tokenizer.pad_token_id, |
| eos_token_id=processor.tokenizer.eos_token_id, |
| use_cache=True, |
| num_beams=1, |
| bad_words_ids=[[processor.tokenizer.unk_token_id]], |
| return_dict_in_generate=True, |
| ) |
| prediction = processor.batch_decode(outputs.sequences)[0] |
| prediction = processor.token2json(prediction) |
| return str (prediction) |
|
|
| examples = ['1.jpg', '2.jpg'] |
| demo = gr.Interface(fn=pred, inputs="image", outputs= "text", examples= examples) |
|
|
|
|
| demo.launch(share= False) |
