answer-evaluation-app / utils /image_processor.py
yeswanthvarma's picture
update utils/image_processor.py
706dc31 verified
raw
history blame
1.46 kB
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."