import gradio as gr
from transformers import TrOCRProcessor, VisionEncoderDecoderModel


def inference(img):
    pretrained_model_path = "trocr-small-handwritten"
    image = img.convert("RGB")

    processor = TrOCRProcessor.from_pretrained(pretrained_model_path)
    model = VisionEncoderDecoderModel.from_pretrained(pretrained_model_path)
    pixel_values = processor(images=image, return_tensors="pt").pixel_values

    generated_ids = model.generate(pixel_values)
    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
    print(generated_text)

title = "OCR:trocr-small-handwritten"
description = "Gradio Demo for trocr-small-handwritten. To use it, simply upload your image, or click one of the examples to load them."
article = "<p style='text-align: center'><a href='https://github.com/bryandlee/animegan2-pytorch' target='_blank'>Github Repo Pytorch</a></p> <center><img src='https://visitor-badge.glitch.me/badge?page_id=akhaliq_animegan' alt='visitor badge'></center></p>"
examples=[['ocr_example.jpg']]


demo = gr.Interface(
    fn=inference,
    inputs=[gr.inputs.Image(type="pil")],
    outputs=gr.outputs.Textbox(),
    title=title,
    description=description,
    article=article,
    examples=examples)

demo.launch()