import torch
import re
from PIL import Image
from transformers import DonutProcessor
#from transformers import VisionEncoderDecoderModel

import warnings
warnings.filterwarnings("ignore")

from sconf import Config
from donut import DonutConfig, DonutModel

config = Config(default="./config.yaml")

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
processor = DonutProcessor.from_pretrained("marzanconsulting/donut-dr-matriculas-ocr")

model = DonutModel.from_pretrained(
                "marzanconsulting/donut-dr-matriculas-ocr",
                input_size=config.input_size,
                max_length=config.max_length,
                align_long_axis=config.align_long_axis,
                ignore_mismatched_sizes=True,
            )

model.to(device)

def load_and_preprocess_image(image_path: str, processor):
    """
    Load an image and preprocess it for the model.
    """
    image = Image.open(image_path).convert("RGB")
    pixel_values = processor(image, return_tensors="pt").pixel_values
    return pixel_values

def generate_text_from_image(model, image_path: str, processor, device):
    """
    Generate text from an image using the trained model.
    """
    # Load and preprocess the image
    pixel_values = load_and_preprocess_image(image_path, processor)
    pixel_values = pixel_values.to(device)

    decoder_input_ids = processor.tokenizer(task_prompt="", 
                                            add_special_tokens=False,
                                            return_tensors="pt").input_ids    

    decoded_text = model.inference(image_tensors=pixel_values, 
                                   prompt_tensors=decoder_input_ids)["predictions"][0]

    return decoded_text

# Example usage
image_path = "path_to_your_image"  # Replace with your image path
extracted_text = generate_text_from_image(model, image_path, processor, device)
print("Extracted Text:", extracted_text)

@article{DBLP:journals/corr/abs-2111-15664,
  author    = {Geewook Kim and
               Teakgyu Hong and
               Moonbin Yim and
               Jinyoung Park and
               Jinyeong Yim and
               Wonseok Hwang and
               Sangdoo Yun and
               Dongyoon Han and
               Seunghyun Park},
  title     = {Donut: Document Understanding Transformer without {OCR}},
  journal   = {CoRR},
  volume    = {abs/2111.15664},
  year      = {2021},
  url       = {https://arxiv.org/abs/2111.15664},
  eprinttype = {arXiv},
  eprint    = {2111.15664},
  timestamp = {Thu, 02 Dec 2021 10:50:44 +0100},
  biburl    = {https://dblp.org/rec/journals/corr/abs-2111-15664.bib},
  bibsource = {dblp computer science bibliography, https://dblp.org}
}