import torch
from transformers import TrOCRProcessor, VisionEncoderDecoderModel
from PIL import Image
import os

# Load processor and model only once
try:
    processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten")
    model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten")
    model.to("cpu")  # You can set to "cuda" if running with GPU
    print("✅ TrOCR model loaded successfully.")
    ocr_available = True
except Exception as e:
    print(f"❌ TrOCR initialization failed: {str(e)}")
    ocr_available = False

def extract_text_from_image(image_path):
    """
    Extract handwritten text from an image using TrOCR (OCR based on transformers).
    """
    print(f"📂 Reading image from: {image_path}")
    try:
        if not ocr_available:
            return "TrOCR is not available."

        # Open image
        image = Image.open(image_path).convert("RGB")

        # Preprocess image
        pixel_values = processor(images=image, return_tensors="pt").pixel_values

        # Run inference
        generated_ids = model.generate(pixel_values)

        # Decode text
        generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]

        return generated_text.strip() or "Text extraction failed. Please enter text manually."

    except Exception as e:
        print(f"OCR failed: {str(e)}")
        return "Text extraction failed. Please enter text manually."