# Install Gradio if you haven't already


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

# Load the processor and model
processor = TrOCRProcessor.from_pretrained('openthaigpt/thai-trocr')
model = VisionEncoderDecoderModel.from_pretrained('openthaigpt/thai-trocr')

# Define the prediction function
def extract_text_from_image(image):
    # Process the input image and run the OCR model
    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]
    return generated_text

# Set up the Gradio interface
interface = gr.Interface(
    fn=extract_text_from_image,  # Function to process the image
    inputs=gr.Image(type="pil"),  # Input is an image (PIL format)
    outputs="text",               # Output is text
    title="Thai OCR with TrOCR",
    description="Upload an image containing Thai text, and this model will extract the text using a Thai-adapted TrOCR model.",
    examples=["path/to/example_image.jpg"]  # Optional: add a sample image path here for users to try
)

# Launch the interface
interface.launch()