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."