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

# Setup logging
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)

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

# Function to recognize handwriting
def recognize_handwriting(image):
    try:
        logger.info("Received an image for handwriting recognition.")
        if isinstance(image, dict):
            image = image.get("image")
        
        if image is None:
            logger.error("No image found in the input.")
            return "No image found"

        image = Image.fromarray(image).convert("RGB")
        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]
        logger.info("Handwriting recognized successfully.")
        return generated_text
    except Exception as e:
        logger.error(f"Error during handwriting recognition: {e}")
        return f"Error: {str(e)}"

# Create Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("## Handwritten Text Recognition")
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(tool="editor", type="numpy", label="Draw or Upload an Image")
            submit_button = gr.Button("Submit")
        with gr.Column():
            output_text = gr.Textbox(label="Recognized Text")
    
    submit_button.click(fn=recognize_handwriting, inputs=image_input, outputs=output_text)

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