import gradio as gr
from transformers import TrOCRProcessor, VisionEncoderDecoderModel
import requests
from PIL import Image
from io import BytesIO

# Load TrOCR model
processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten")
model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten")

def process_image(image):
    """Processes an image with TrOCR and returns the extracted text."""
    # Prepare image
    pixel_values = processor(image, return_tensors="pt").pixel_values

    # Generate text
    generated_ids = model.generate(pixel_values, max_new_tokens=50)

    # Decode text
    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]

    return generated_text

# Gradio interface
title = "Handwritten Text OCR with TrOCR"
description = "Upload a handwritten text image and get the extracted text using Microsoft's TrOCR model."
article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2109.10282'>TrOCR: Transformer-based Optical Character Recognition</a> | <a href='https://github.com/microsoft/unilm/tree/master/trocr'>GitHub Repo</a></p>"

examples = [
    ["https://fki.tic.heia-fr.ch/static/img/a01-122-02.jpg"],
    ["https://upload.wikimedia.org/wikipedia/commons/8/8d/Handwriting-sample.jpg"]
]

iface = gr.Interface(fn=process_image,
                     inputs=gr.Image(type="pil"),
                     outputs=gr.Textbox(),
                     title=title,
                     description=description,
                     article=article,
                     examples=examples)

iface.launch()