"""Image preprocessing for optimal OCR results.""" from pathlib import Path from typing import List import numpy as np import cv2 try: import pdf2image PDF_AVAILABLE = True except ImportError: PDF_AVAILABLE = False class ImagePreprocessor: """Preprocess receipt images for OCR.""" def load_image(self, path: Path) -> np.ndarray: """Load image from file.""" image = cv2.imread(str(path)) if image is None: raise ValueError(f"Could not load image: {path}") return image def pdf_to_images(self, path: Path, dpi: int = 300) -> List[np.ndarray]: """ Convert PDF to images. Args: path: Path to PDF file dpi: Resolution (300 = fast & good quality, 400 = better but slower) """ if not PDF_AVAILABLE: raise RuntimeError("pdf2image not available. Install with: pip install pdf2image") images = pdf2image.convert_from_path(str(path), dpi=dpi) return [np.array(img) for img in images] def preprocess(self, image: np.ndarray, high_quality: bool = True) -> np.ndarray: """ Apply LIGHT preprocessing - better for clear PDFs. Heavy binarization can destroy text on clear images. """ return self.preprocess_light(image) def preprocess_light(self, image: np.ndarray) -> np.ndarray: """ Light preprocessing for CLEAR images (PDFs, good scans). Preserves original quality, only enhances contrast. """ # 1. Grayscale if len(image.shape) == 3: gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) else: gray = image.copy() # 2. Resize if too small height, width = gray.shape if width < 1500: scale = 1500 / width gray = cv2.resize(gray, None, fx=scale, fy=scale, interpolation=cv2.INTER_CUBIC) # 3. Deskew gray = self._deskew(gray) # 4. Light contrast enhancement only clahe = cv2.createCLAHE(clipLimit=1.5, tileGridSize=(8, 8)) enhanced = clahe.apply(gray) # NO binarization, NO morphological ops - preserve original quality return enhanced def preprocess_heavy(self, image: np.ndarray) -> np.ndarray: """ Heavy preprocessing for FADED thermal receipts. Aggressive binarization to recover faded text. """ # 1. Grayscale if len(image.shape) == 3: gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) else: gray = image.copy() # 2. Resize if too small (larger = better OCR) height, width = gray.shape if width < 1500: scale = 1500 / width gray = cv2.resize(gray, None, fx=scale, fy=scale, interpolation=cv2.INTER_CUBIC) # 3. Deskew gray = self._deskew(gray) # 4. Contrast enhancement with CLAHE clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8)) enhanced = clahe.apply(gray) # 5. Denoise denoised = cv2.fastNlMeansDenoising(enhanced, h=8, templateWindowSize=7, searchWindowSize=21) # 6. Sharpening gaussian = cv2.GaussianBlur(denoised, (0, 0), 2.0) sharpened = cv2.addWeighted(denoised, 1.5, gaussian, -0.5, 0) # 7. Adaptive thresholding (binarization) binary = cv2.adaptiveThreshold( sharpened, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, blockSize=11, C=5 ) # 8. Morphological operations kernel_close = cv2.getStructuringElement(cv2.MORPH_RECT, (2, 2)) result = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel_close) return result def get_all_variants(self, image: np.ndarray) -> List[np.ndarray]: """ Generate 2 preprocessing variants for OCR (fast mode). Returns: [light_processed, heavy_processed] """ return [ self.preprocess_light(image), self.preprocess_heavy(image), ] def _deskew(self, image: np.ndarray) -> np.ndarray: """Correct image rotation/skew using Hough lines.""" edges = cv2.Canny(image, 50, 150, apertureSize=3) lines = cv2.HoughLinesP( edges, 1, np.pi / 180, threshold=100, minLineLength=100, maxLineGap=10 ) if lines is None: return image angles = [] for line in lines: x1, y1, x2, y2 = line[0] angle = np.arctan2(y2 - y1, x2 - x1) * 180 / np.pi if abs(angle) < 45: angles.append(angle) if not angles: return image median_angle = np.median(angles) if abs(median_angle) < 0.5: return image h, w = image.shape[:2] center = (w // 2, h // 2) M = cv2.getRotationMatrix2D(center, median_angle, 1.0) return cv2.warpAffine( image, M, (w, h), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_REPLICATE )