Spaces:
Sleeping
Sleeping
File size: 3,838 Bytes
de76d9d | 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 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 | """
OCR Processor Module
Applies image preprocessing to improve Tesseract OCR accuracy.
"""
import logging
import re
import cv2
import numpy as np
import pytesseract
from PIL import Image
logger = logging.getLogger(__name__)
class OCRProcessor:
"""
Wraps Tesseract OCR with image preprocessing pipeline:
- Grayscale conversion
- Adaptive thresholding (handles uneven lighting)
- Noise removal via median blur
- Deskewing for rotated scans
"""
def __init__(self, tesseract_cmd: str = None):
if tesseract_cmd:
pytesseract.pytesseract.tesseract_cmd = tesseract_cmd
# Tesseract config: PSM 3 = fully automatic page segmentation
self.tsv_config = r"--oem 3 --psm 3"
def preprocess(self, image: Image.Image) -> np.ndarray:
"""
Apply a preprocessing pipeline to improve OCR quality.
Returns a numpy array (grayscale, thresholded image).
"""
# Convert PIL to OpenCV format
img_array = np.array(image.convert("RGB"))
img_bgr = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)
# Step 1: Convert to grayscale
gray = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY)
# Step 2: Denoise
denoised = cv2.medianBlur(gray, 3)
# Step 3: Adaptive thresholding (works well for uneven lighting)
thresh = cv2.adaptiveThreshold(
denoised,
255,
cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
cv2.THRESH_BINARY,
31, # block size
10, # constant subtracted from mean
)
# Step 4: Morphological dilation to connect broken characters
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1, 1))
processed = cv2.dilate(thresh, kernel, iterations=1)
return processed
def extract(self, image: Image.Image) -> str:
"""
Run OCR on a PIL Image.
Returns cleaned, normalized text.
"""
try:
processed = self.preprocess(image)
# Convert back to PIL for pytesseract
processed_pil = Image.fromarray(processed)
raw_text = pytesseract.image_to_string(processed_pil, config=self.tsv_config)
return self._clean_ocr_output(raw_text)
except Exception as e:
logger.error(f"OCR extraction failed: {e}")
# Fallback: try without preprocessing
try:
raw_text = pytesseract.image_to_string(image, config=self.tsv_config)
return self._clean_ocr_output(raw_text)
except Exception as e2:
logger.error(f"OCR fallback also failed: {e2}")
return ""
def _clean_ocr_output(self, text: str) -> str:
"""
Clean raw OCR output:
- Remove control characters
- Normalize whitespace
- Remove page numbers and artifacts
- Merge broken hyphenated lines
"""
if not text:
return ""
# Normalize unicode
text = text.encode("utf-8", errors="ignore").decode("utf-8")
# Merge hyphenated line breaks (e.g., "docu-\nment" → "document")
text = re.sub(r"-\n(\w)", r"\1", text)
# Replace multiple spaces/tabs with single space
text = re.sub(r"[ \t]+", " ", text)
# Keep at most 2 consecutive newlines
text = re.sub(r"\n{3,}", "\n\n", text)
# Remove lines that are pure noise (single chars, page numbers)
lines = text.splitlines()
cleaned_lines = []
for line in lines:
stripped = line.strip()
# Skip very short lines that look like artifacts
if len(stripped) <= 1 and not stripped.isalpha():
continue
cleaned_lines.append(stripped)
return "\n".join(cleaned_lines).strip()
|