Spaces:
Sleeping
Sleeping
File size: 1,461 Bytes
706dc31 3497db5 706dc31 436a6c3 706dc31 3497db5 706dc31 3497db5 706dc31 3497db5 436a6c3 5888c6a 706dc31 5888c6a 6aedaf3 706dc31 3497db5 706dc31 3497db5 706dc31 3497db5 706dc31 3497db5 706dc31 3497db5 706dc31 986c935 5888c6a 6aedaf3 3497db5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 |
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."
|