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