# app.py
import gradio as gr
from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline
from PIL import Image
import numpy as np

# Load model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("microsoft/trocr-base-handwritten")
model = AutoModelForTokenClassification.from_pretrained("microsoft/trocr-base-handwritten")

# Create OCR pipeline
ocr_pipeline = pipeline(
    "image-to-text",
    model=model,
    tokenizer=tokenizer,
    feature_extractor=tokenizer.init_feature_extractor()
)

def predict_handwriting(image):
    """
    Function to process handwritten text image and return transcription
    """
    try:
        # Preprocess image
        image = image.convert("RGB")
        image = np.array(image)
        
        # Run OCR
        result = ocr_pipeline(image)
        
        # Extract text from results
        transcription = " ".join([word["value"] for word in result])
        return transcription
    
    except Exception as e:
        return f"Error processing image: {str(e)}"

# Create Gradio interface
demo = gr.Interface(
    fn=predict_handwriting,
    inputs=gr.Image(type="pil", label="Upload Handwritten Text Image"),
    outputs=gr.Textbox(label="Transcription"),
    title="Handwritten Text to Text Converter",
    description="Upload a handwritten text image and get the transcribed text. Best results with clear, high-contrast images."
)

if __name__ == "__main__":
    demo.launch()