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

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

# Ensure model is on appropriate device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# OCR function
def extract_text(image):
    if image is None:
        return "Please upload an image."
    
    image = image.convert("RGB")
    pixel_values = processor(images=image, return_tensors="pt").pixel_values.to(device)

    generated_ids = model.generate(pixel_values)
    text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
    
    return text.strip()

# Build Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("![Logo](logo.png)")
    gr.Markdown("""
    # Evaluator! Upgrade Your Writing Skill
    """)
    
    with gr.Row():
        input_image = gr.Image(type="pil", label="Upload Handwritten Image")
        output_text = gr.Textbox(label="Extracted Text", lines=8)

    submit_btn = gr.Button("Extract Text")
    submit_btn.click(fn=extract_text, inputs=input_image, outputs=output_text)

# Run the app
if __name__ == "__main__":
    demo.launch()