Spaces:
Sleeping
Sleeping
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." | |