preprocessing.py

"""
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 <image_path>")
        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}")