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

# Load the model and processor from Hugging Face
model = VisionEncoderDecoderModel.from_pretrained("paudelanil/denvagari-TrOCR")
processor = TrOCRProcessor.from_pretrained("paudelanil/denvagari-TrOCR")

def predict(image):
    # Preprocess the image
    image = Image.open(image).convert("RGB")
    pixel_values = processor(image, return_tensors="pt").pixel_values
    
    # Generate text from the image
    generated_ids = model.generate(pixel_values)
    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
    
    return generated_text

# Create the Gradio interface
interface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="filepath"),
    outputs="text",
    title="Devanagari OCR with TrOCR",
    description="Upload an image with Devanagari script and get the text prediction using a pre-trained Vision-Text model."
)

# Launch the interface
interface.launch()