from transformers import TrOCRProcessor, VisionEncoderDecoderModel
from PIL import Image
import torch
import gradio as gr

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

# Define prediction function
def recognize_text(image):
    image = Image.fromarray(image).convert("RGB")
    pixel_values = processor(images=image, return_tensors="pt").pixel_values
    generated_ids = model.generate(pixel_values)
    text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
    return text

# Create Gradio interface
interface = gr.Interface(
    fn=recognize_text,
    inputs=gr.Image(type="numpy", label="Upload Image"),
    outputs=gr.Textbox(label="Recognized Text"),
    title="Handwritten Text Recognition",
    description="Upload an image of handwritten text to recognize it using TrOCR.",
)

# Launch interface
interface.launch(share=True)