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

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

def ocr_recognition(image):
    # Open the image
    image = Image.open(image).convert("RGB")

    # Process the image and generate text
    pixel_values = processor(image, return_tensors="pt").pixel_values
    generated_ids = model.generate(pixel_values)
    
    # Decode the output text
    text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
    
    return text

# Create Gradio interface
iface = gr.Interface(
    fn=ocr_recognition,
    inputs=gr.Image(type="pil"),  # Ensures PIL image input
    outputs="text",
    title="Handwritten OCR Extraction",
    description="Upload a handwritten image to extract text using TrOCR."
)

# Launch the Gradio app
iface.launch()