""" Lightweight Image Preprocessing Pipeline for IDP Uses OpenCV and Pillow for CPU-friendly operations Optimizes document images for OCR quality """ import cv2 import numpy as np from PIL import Image, ImageEnhance from typing import Tuple, Optional import logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) class DocumentPreprocessor: """ Lightweight document preprocessing pipeline optimized for OCR All operations are CPU-friendly and designed for speed """ def __init__( self, max_width: int = 2048, max_height: int = 2048, enable_deskew: bool = True, enable_denoise: bool = True, enable_contrast: bool = True, ): """ Args: max_width: Maximum width for resizing max_height: Maximum height for resizing enable_deskew: Enable deskewing correction enable_denoise: Enable noise reduction enable_contrast: Enable contrast enhancement """ self.max_width = max_width self.max_height = max_height self.enable_deskew = enable_deskew self.enable_denoise = enable_denoise self.enable_contrast = enable_contrast def preprocess( self, image: np.ndarray, adaptive_threshold: bool = False ) -> np.ndarray: """ Complete preprocessing pipeline Args: image: Input image as numpy array (BGR or RGB) adaptive_threshold: Apply adaptive thresholding for poor quality scans Returns: Preprocessed image ready for OCR """ logger.info("Starting preprocessing pipeline") # Step 1: Auto-rotation (from EXIF metadata) image = self._auto_rotate(image) # Step 2: Resize to optimal dimensions image = self._resize_image(image) # Step 3: Deskew correction if self.enable_deskew: image = self._deskew_image(image) # Step 4: Denoise if self.enable_denoise: image = self._denoise_image(image) # Step 5: Contrast enhancement if self.enable_contrast: image = self._enhance_contrast(image) # Step 6: Adaptive thresholding (optional, for very poor scans) if adaptive_threshold: image = self._adaptive_threshold(image) logger.info("Preprocessing complete") return image def _auto_rotate(self, image: np.ndarray) -> np.ndarray: """ Auto-rotate image based on EXIF orientation For images without EXIF, uses simple heuristics """ # Convert to PIL to read EXIF if len(image.shape) == 2: pil_image = Image.fromarray(image) else: # OpenCV uses BGR, PIL uses RGB rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) pil_image = Image.fromarray(rgb_image) # Try to get EXIF orientation try: exif = pil_image._getexif() if exif: orientation = exif.get(274) # 274 is orientation tag if orientation == 3: pil_image = pil_image.rotate(180, expand=True) elif orientation == 6: pil_image = pil_image.rotate(270, expand=True) elif orientation == 8: pil_image = pil_image.rotate(90, expand=True) except (AttributeError, KeyError, TypeError): # No EXIF data, skip auto-rotation pass # Convert back to numpy image = np.array(pil_image) if len(image.shape) == 3: image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) return image def _resize_image(self, image: np.ndarray) -> np.ndarray: """ Resize image to optimal dimensions for OCR Maintains aspect ratio """ h, w = image.shape[:2] # Calculate scaling factor scale = min(self.max_width / w, self.max_height / h, 1.0) if scale < 1.0: new_w = int(w * scale) new_h = int(h * scale) image = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_AREA) logger.info(f"Resized from ({w}, {h}) to ({new_w}, {new_h})") return image def _deskew_image(self, image: np.ndarray) -> np.ndarray: """ Detect and correct skew using Hough line transform Fast and efficient for typical document skews """ # Convert to grayscale if needed if len(image.shape) == 3: gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) else: gray = image.copy() # Edge detection edges = cv2.Canny(gray, 50, 150, apertureSize=3) # Detect lines lines = cv2.HoughLines(edges, 1, np.pi / 180, 200) if lines is not None and len(lines) > 0: # Calculate angles angles = [] for rho, theta in lines[:, 0]: angle = np.degrees(theta) - 90 if -45 < angle < 45: # Only consider reasonable skew angles angles.append(angle) if angles: # Median angle is most robust skew_angle = np.median(angles) # Only correct if skew is significant (> 0.5 degrees) if abs(skew_angle) > 0.5: logger.info(f"Detected skew: {skew_angle:.2f} degrees") # Rotate image h, w = image.shape[:2] center = (w // 2, h // 2) M = cv2.getRotationMatrix2D(center, skew_angle, 1.0) image = cv2.warpAffine( image, M, (w, h), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_REPLICATE ) return image def _denoise_image(self, image: np.ndarray) -> np.ndarray: """ Apply bilateral filter for noise reduction Preserves edges while smoothing noise """ # Bilateral filter: smooths noise but preserves edges # d: diameter of pixel neighborhood # sigmaColor: filter sigma in color space # sigmaSpace: filter sigma in coordinate space denoised = cv2.bilateralFilter(image, d=5, sigmaColor=50, sigmaSpace=50) logger.info("Applied bilateral denoising") return denoised def _enhance_contrast(self, image: np.ndarray) -> np.ndarray: """ Enhance contrast using CLAHE (Contrast Limited Adaptive Histogram Equalization) More effective than global histogram equalization for documents """ # Convert to LAB color space for better contrast adjustment if len(image.shape) == 3: lab = cv2.cvtColor(image, cv2.COLOR_BGR2LAB) l, a, b = cv2.split(lab) else: l = image.copy() # Apply CLAHE to L channel clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8)) l = clahe.apply(l) # Merge back if len(image.shape) == 3: lab = cv2.merge([l, a, b]) image = cv2.cvtColor(lab, cv2.COLOR_LAB2BGR) else: image = l logger.info("Applied CLAHE contrast enhancement") return image def _adaptive_threshold(self, image: np.ndarray) -> np.ndarray: """ Apply adaptive thresholding for poor quality scans Converts to binary image """ # Convert to grayscale if needed if len(image.shape) == 3: gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) else: gray = image.copy() # Adaptive threshold binary = cv2.adaptiveThreshold( gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, blockSize=11, C=2 ) logger.info("Applied adaptive thresholding") return binary def preprocess_for_ocr( image_path: str, max_width: int = 2048, adaptive_threshold: bool = False ) -> np.ndarray: """ Convenience function to preprocess an image file for OCR Args: image_path: Path to input image max_width: Maximum width for resizing adaptive_threshold: Apply adaptive thresholding Returns: Preprocessed image as numpy array """ # Load image image = cv2.imread(image_path) if image is None: raise ValueError(f"Could not load image from {image_path}") # Preprocess preprocessor = DocumentPreprocessor(max_width=max_width) processed = preprocessor.preprocess(image, adaptive_threshold=adaptive_threshold) return processed if __name__ == "__main__": # Example usage import sys if len(sys.argv) < 2: print("Usage: python preprocessing.py ") sys.exit(1) input_path = sys.argv[1] output_path = "preprocessed_output.png" # Preprocess processed = preprocess_for_ocr(input_path) # Save result cv2.imwrite(output_path, processed) print(f"Preprocessed image saved to {output_path}")